Skip to content

Commit

Permalink
[CALCITE-6724] MockTable support for multiple (individual & composite…
Browse files Browse the repository at this point in the history
…) keys

MockTable currently allows only one key to be specified; the key can be either simple or composite. In real scenarios though a table can have multiple keys.
  • Loading branch information
zabetak committed Dec 10, 2024
1 parent c3478b9 commit fcf40bd
Show file tree
Hide file tree
Showing 3 changed files with 145 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -4198,9 +4198,10 @@ private static class CompositeKeysCatalogReader
// Register "T1" table.
final MockTable t1 =
MockTable.create(this, tSchema, "composite_keys_table", false, 7.0, null);
t1.addColumn("key1", typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
t1.addColumn("key2", typeFactory.createSqlType(SqlTypeName.VARCHAR), true);
t1.addColumn("key1", typeFactory.createSqlType(SqlTypeName.VARCHAR));
t1.addColumn("key2", typeFactory.createSqlType(SqlTypeName.VARCHAR));
t1.addColumn("value1", typeFactory.createSqlType(SqlTypeName.INTEGER));
t1.addKey("key1", "key2");
addSizeHandler(t1);
addDistinctRowcountHandler(t1);
addUniqueKeyHandler(t1);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -316,7 +316,7 @@ public static class MockTable extends Prepare.AbstractPreparingTable
protected final double rowCount;
protected final List<Map.Entry<String, RelDataType>> columnList =
new ArrayList<>();
protected final List<Integer> keyList = new ArrayList<>();
protected final List<ImmutableBitSet> keyList = new ArrayList<>();
protected final List<RelReferentialConstraint> referentialConstraints =
new ArrayList<>();
protected RelDataType rowType;
Expand Down Expand Up @@ -378,7 +378,7 @@ private MockTable(MockCatalogReader catalogReader, List<String> names,
*/
protected MockTable(MockCatalogReader catalogReader, boolean stream,
boolean temporal, double rowCount,
List<Map.Entry<String, RelDataType>> columnList, List<Integer> keyList,
List<Map.Entry<String, RelDataType>> columnList, List<ImmutableBitSet> keyList,
RelDataType rowType, List<RelCollation> collationList, List<String> names,
Set<String> monotonicColumnSet, StructKind kind,
@Nullable ColumnResolver resolver,
Expand Down Expand Up @@ -592,15 +592,11 @@ public static MockTable create(MockCatalogReader catalogReader,
}

@Override public boolean isKey(ImmutableBitSet columns) {
return !keyList.isEmpty()
&& columns.contains(ImmutableBitSet.of(keyList));
return keyList.stream().anyMatch(columns::contains);
}

@Override public List<ImmutableBitSet> getKeys() {
if (keyList.isEmpty()) {
return ImmutableList.of();
}
return ImmutableList.of(ImmutableBitSet.of(keyList));
return keyList;
}

@Override public List<RelReferentialConstraint> getReferentialConstraints() {
Expand Down Expand Up @@ -651,11 +647,43 @@ public void addColumn(String name, RelDataType type) {

public void addColumn(String name, RelDataType type, boolean isKey) {
if (isKey) {
keyList.add(columnList.size());
keyList.add(ImmutableBitSet.of(columnList.size()));
}
columnList.add(Pair.of(name, type));
}

public void addKey(String... columns) {
ImmutableBitSet.Builder keyBuilder = ImmutableBitSet.builder();
for (String c : columns) {
int i = columnIndex(c);
if (i < 0) {
throw new IllegalArgumentException("Column " + c + " not found in the table");
}
keyBuilder.set(i);
}
keyList.add(keyBuilder.build());
}

public void addKey(ImmutableBitSet key) {
for (Integer columnIndex : key) {
if (columnIndex >= columnList.size()) {
throw new IllegalArgumentException(
"Column index " + columnIndex + " exceeds the number of columns");
}
}
keyList.add(key);
}

private int columnIndex(String colName) {
for (int i = 0; i < columnList.size(); i++) {
Map.Entry<String, RelDataType> col = columnList.get(i);
if (colName.equals(col.getKey())) {
return i;
}
}
return -1;
}

public void addMonotonic(String name) {
monotonicColumnSet.add(name);
assert Pair.left(columnList).contains(name);
Expand Down Expand Up @@ -712,7 +740,7 @@ private MockModifiableViewRelOptTable(MockModifiableViewTable modifiableViewTabl
*/
private MockModifiableViewRelOptTable(MockModifiableViewTable modifiableViewTable,
MockCatalogReader catalogReader, boolean stream, double rowCount,
List<Map.Entry<String, RelDataType>> columnList, List<Integer> keyList,
List<Map.Entry<String, RelDataType>> columnList, List<ImmutableBitSet> keyList,
RelDataType rowType, List<RelCollation> collationList, List<String> names,
Set<String> monotonicColumnSet, StructKind kind,
@Nullable ColumnResolver resolver,
Expand Down
104 changes: 104 additions & 0 deletions testkit/src/test/java/org/apache/calcite/test/MockTableTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to you under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.calcite.test;

import org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.sql2rel.NullInitializerExpressionFactory;
import org.apache.calcite.test.catalog.MockCatalogReader;
import org.apache.calcite.util.ImmutableBitSet;

import org.junit.jupiter.api.Test;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.hasToString;
import static org.junit.jupiter.api.Assertions.assertThrows;

/**
* Tests for {@link org.apache.calcite.test.catalog.MockCatalogReader.MockTable}.
*/
public class MockTableTest {
private static final SqlTypeFactoryImpl TYPE_FACTORY =
new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);

@Test void testAddColumnCreatesIndividualKeys() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), true);
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), true);
assertThat(t.getKeys(), hasToString("[{0}, {1}]"));
}

@Test void testAddKeyWithOneEntryCreatesSimpleKey() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addKey("k1");
assertThat(t.getKeys(), hasToString("[{0}]"));
}

@Test void testAddKeyWithMultipleEntriesCreatesCompositeKey() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addKey("k1", "k2");
assertThat(t.getKeys(), hasToString("[{0, 1}]"));
}

@Test void testAddKeyWithMissingColumnNameThrowsException() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
assertThrows(IllegalArgumentException.class, () -> t.addKey("k1", "k3"));
}

@Test void testAddKeyUsingColumnIndex() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addKey(ImmutableBitSet.of(0, 1));
assertThat(t.getKeys(), hasToString("[{0, 1}]"));
}

@Test void testAddKeyUsingWrongIndexThrowsException() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
assertThrows(IllegalArgumentException.class, () -> t.addKey(ImmutableBitSet.of(0, 2)));
}

@Test void testAddKeyMultipleTimes() {
MockCatalogReader.MockTable t = newTable();
t.addColumn("k1", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k2", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addColumn("k3", TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
t.addKey("k1");
t.addKey("k2", "k3");
assertThat(t.getKeys(), hasToString("[{0}, {1, 2}]"));
}

private static MockCatalogReader.MockTable newTable() {

MockCatalogReader catalogReader = new MockCatalogReader(TYPE_FACTORY, false) {
@Override public MockCatalogReader init() {
return this;
}
};
return new MockCatalogReader.MockTable(catalogReader, "catalog", "schema", "table", false,
false, 0.0, null, NullInitializerExpressionFactory.INSTANCE);
}
}

0 comments on commit fcf40bd

Please sign in to comment.