diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java index bc1809a41aa2..1e0b2be72c5e 100644 --- a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java @@ -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); diff --git a/testkit/src/main/java/org/apache/calcite/test/catalog/MockCatalogReader.java b/testkit/src/main/java/org/apache/calcite/test/catalog/MockCatalogReader.java index 01aeb5b8b2c9..164c2825ef71 100644 --- a/testkit/src/main/java/org/apache/calcite/test/catalog/MockCatalogReader.java +++ b/testkit/src/main/java/org/apache/calcite/test/catalog/MockCatalogReader.java @@ -316,7 +316,7 @@ public static class MockTable extends Prepare.AbstractPreparingTable protected final double rowCount; protected final List> columnList = new ArrayList<>(); - protected final List keyList = new ArrayList<>(); + protected final List keyList = new ArrayList<>(); protected final List referentialConstraints = new ArrayList<>(); protected RelDataType rowType; @@ -378,7 +378,7 @@ private MockTable(MockCatalogReader catalogReader, List names, */ protected MockTable(MockCatalogReader catalogReader, boolean stream, boolean temporal, double rowCount, - List> columnList, List keyList, + List> columnList, List keyList, RelDataType rowType, List collationList, List names, Set monotonicColumnSet, StructKind kind, @Nullable ColumnResolver resolver, @@ -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 getKeys() { - if (keyList.isEmpty()) { - return ImmutableList.of(); - } - return ImmutableList.of(ImmutableBitSet.of(keyList)); + return keyList; } @Override public List getReferentialConstraints() { @@ -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 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); @@ -712,7 +740,7 @@ private MockModifiableViewRelOptTable(MockModifiableViewTable modifiableViewTabl */ private MockModifiableViewRelOptTable(MockModifiableViewTable modifiableViewTable, MockCatalogReader catalogReader, boolean stream, double rowCount, - List> columnList, List keyList, + List> columnList, List keyList, RelDataType rowType, List collationList, List names, Set monotonicColumnSet, StructKind kind, @Nullable ColumnResolver resolver, diff --git a/testkit/src/test/java/org/apache/calcite/test/MockTableTest.java b/testkit/src/test/java/org/apache/calcite/test/MockTableTest.java new file mode 100644 index 000000000000..8becd8a9ca03 --- /dev/null +++ b/testkit/src/test/java/org/apache/calcite/test/MockTableTest.java @@ -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); + } +}