diff --git a/core/src/main/java/org/apache/calcite/rel/core/Collect.java b/core/src/main/java/org/apache/calcite/rel/core/Collect.java index 475cc966d3a..d2f89eeec3b 100644 --- a/core/src/main/java/org/apache/calcite/rel/core/Collect.java +++ b/core/src/main/java/org/apache/calcite/rel/core/Collect.java @@ -29,9 +29,12 @@ import org.apache.calcite.sql.SqlKind; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.sql.type.SqlTypeUtil; +import org.apache.calcite.util.Litmus; import com.google.common.collect.Iterables; +import org.checkerframework.checker.nullness.qual.Nullable; + import java.util.List; import static java.util.Objects.requireNonNull; @@ -184,6 +187,31 @@ public RelNode copy(RelTraitSet traitSet, RelNode input) { return new Collect(getCluster(), traitSet, input, rowType()); } + @Override public boolean isValid(Litmus litmus, @Nullable Context context) { + final RelDataTypeFactory typeFactory = getCluster().getTypeFactory(); + final RelDataType inputRow = getInput().getRowType(); + if (getCollectionType() == SqlTypeName.MAP && inputRow.getFieldCount() != 2) { + return litmus.fail("MAP requires an input with exactly two fields;" + + " input row type is {}", inputRow); + } + final RelDataType derived = + deriveRowType(typeFactory, getCollectionType(), getFieldName(), inputRow); + if (rowType().equals(derived)) { + return super.isValid(litmus, context); + } + // A Collect created for a collection query constructor derives its element + // type from the input row type; see #create(RelNode, SqlKind, String). + final RelDataType derivedForQuery = + deriveRowType(typeFactory, getCollectionType(), getFieldName(), + SqlTypeUtil.deriveCollectionQueryComponentType(typeFactory, + getCollectionType(), inputRow)); + if (rowType().equals(derivedForQuery)) { + return super.isValid(litmus, context); + } + return litmus.fail("row type {} does not match the type {}" + + " derived from the input", rowType(), derived); + } + @Override public RelNode accept(RelShuttle shuttle) { return shuttle.visit(this); } diff --git a/core/src/test/java/org/apache/calcite/rel/core/CollectTest.java b/core/src/test/java/org/apache/calcite/rel/core/CollectTest.java new file mode 100644 index 00000000000..694eebf1f93 --- /dev/null +++ b/core/src/test/java/org/apache/calcite/rel/core/CollectTest.java @@ -0,0 +1,80 @@ +/* + * 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.rel.core; + +import org.apache.calcite.plan.Convention; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.test.RelBuilderTest; +import org.apache.calcite.tools.RelBuilder; +import org.apache.calcite.util.Litmus; + +import org.junit.jupiter.api.Test; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; + +/** + * Test cases for [CALCITE-7717] + * Add a Collect.isValid method to check type invariants. + */ +class CollectTest { + @Test void testIsValid() { + final RelBuilder b = RelBuilder.create(RelBuilderTest.config().build()); + final RelNode oneColumn = b.values(new String[] {"i"}, 1, 2).build(); + final RelDataTypeFactory typeFactory = oneColumn.getCluster().getTypeFactory(); + + // Element type is the input row type. 'x' is the name of the result field + final Collect collect0 = Collect.create(oneColumn, SqlKind.ARRAY_QUERY_CONSTRUCTOR, "x"); + assertThat(collect0.isValid(Litmus.IGNORE, null), is(true)); + + // Element type is the type of the sole input column. + final Collect collect1 = + new Collect(oneColumn.getCluster(), + oneColumn.getCluster().traitSetOf(Convention.NONE), oneColumn, + Collect.deriveRowType(typeFactory, SqlTypeName.ARRAY, "x", + oneColumn.getRowType().getFieldList().get(0).getType())); + assertThat(collect1.isValid(Litmus.IGNORE, null), is(true)); + + // Array over two columns is invalid + final RelNode twoColumns = b.values(new String[] {"k", "v"}, 1, "a").build(); + final Collect mismatched = + new Collect(oneColumn.getCluster(), + oneColumn.getCluster().traitSetOf(Convention.NONE), oneColumn, + Collect.deriveRowType(typeFactory, SqlTypeName.ARRAY, "x", + twoColumns.getRowType())); + assertThat(mismatched.isValid(Litmus.IGNORE, null), is(false)); + + final RelDataType mapRowType = + Collect.deriveRowType(typeFactory, SqlTypeName.MAP, "x", twoColumns.getRowType()); + + // A MAP(subquery) over an input that does not have exactly two columns is invalid + final Collect mapOverOneColumn = + new Collect(oneColumn.getCluster(), + oneColumn.getCluster().traitSetOf(Convention.NONE), oneColumn, mapRowType); + assertThat(mapOverOneColumn.isValid(Litmus.IGNORE, null), is(false)); + + // The same MAP row type over the two-column input is valid. + final Collect mapOverTwoColumns = + new Collect(twoColumns.getCluster(), + twoColumns.getCluster().traitSetOf(Convention.NONE), twoColumns, mapRowType); + assertThat(mapOverTwoColumns.isValid(Litmus.IGNORE, null), is(true)); + } +}