This is an automated email from the ASF dual-hosted git repository. zstan pushed a commit to branch ignite-2.18 in repository https://gitbox.apache.org/repos/asf/ignite.git
commit 7e15f04ab6649f8931bd295b2ce806c2baf5c8cb Author: Aleksey Plekhanov <[email protected]> AuthorDate: Fri Feb 6 09:23:05 2026 +0300 IGNITE-27738 SQL Calcite: Fix planner hang on multi-row values - Fixes #12697. Signed-off-by: Aleksey Plekhanov <[email protected]> (cherry picked from commit f0964f0c05bef8926dda345c8777a3adbb23f218) --- .../rel/agg/IgniteColocatedSortAggregate.java | 8 +- .../calcite/rel/agg/IgniteMapSortAggregate.java | 10 +- .../calcite/rel/agg/IgniteReduceSortAggregate.java | 8 +- .../processors/query/calcite/trait/TraitUtils.java | 68 ++++++++- .../integration/SortAggregateIntegrationTest.java | 15 ++ .../integration/TableDmlIntegrationTest.java | 15 ++ .../calcite/planner/JoinColocationPlannerTest.java | 28 ++++ .../calcite/planner/SortAggregatePlannerTest.java | 168 ++++++++++++++------- 8 files changed, 253 insertions(+), 67 deletions(-) diff --git a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteColocatedSortAggregate.java b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteColocatedSortAggregate.java index 50488bf4d4f..3c85ddc89bc 100644 --- a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteColocatedSortAggregate.java +++ b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteColocatedSortAggregate.java @@ -19,7 +19,6 @@ package org.apache.ignite.internal.processors.query.calcite.rel.agg; import java.util.List; import java.util.Objects; - import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptCost; import org.apache.calcite.plan.RelOptPlanner; @@ -74,8 +73,13 @@ public class IgniteColocatedSortAggregate extends IgniteColocatedAggregateBase i /** {@inheritDoc} */ @Override public Aggregate copy(RelTraitSet traitSet, RelNode input, ImmutableBitSet groupSet, List<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) { + RelCollation collation = TraitUtils.collation(input.getTraitSet()); + + assert collation.satisfies(TraitUtils.collation(traitSet)) + : "Unexpected collations: input=" + collation + ", traitSet=" + TraitUtils.collation(traitSet); + return new IgniteColocatedSortAggregate( - getCluster(), traitSet, input, groupSet, groupSets, aggCalls, TraitUtils.collation(traitSet)); + getCluster(), traitSet, input, groupSet, groupSets, aggCalls, collation); } /** {@inheritDoc} */ diff --git a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteMapSortAggregate.java b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteMapSortAggregate.java index 8efba40250c..15b37d76242 100644 --- a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteMapSortAggregate.java +++ b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteMapSortAggregate.java @@ -82,9 +82,15 @@ public class IgniteMapSortAggregate extends IgniteMapAggregateBase implements Ig RelNode input, ImmutableBitSet groupSet, List<ImmutableBitSet> groupSets, - List<AggregateCall> aggCalls) { + List<AggregateCall> aggCalls + ) { + RelCollation collation = TraitUtils.collation(input.getTraitSet()); + + assert collation.satisfies(TraitUtils.collation(traitSet)) + : "Unexpected collations: input=" + collation + ", traitSet=" + TraitUtils.collation(traitSet); + return new IgniteMapSortAggregate( - getCluster(), traitSet, input, groupSet, groupSets, aggCalls, TraitUtils.collation(traitSet)); + getCluster(), traitSet, input, groupSet, groupSets, aggCalls, collation); } /** {@inheritDoc} */ diff --git a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteReduceSortAggregate.java b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteReduceSortAggregate.java index 632df8863a4..7d1700bcd39 100644 --- a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteReduceSortAggregate.java +++ b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/agg/IgniteReduceSortAggregate.java @@ -19,7 +19,6 @@ package org.apache.ignite.internal.processors.query.calcite.rel.agg; import java.util.List; import java.util.Objects; - import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptCost; import org.apache.calcite.plan.RelOptPlanner; @@ -75,6 +74,11 @@ public class IgniteReduceSortAggregate extends IgniteReduceAggregateBase impleme /** {@inheritDoc} */ @Override public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) { + RelCollation collation = TraitUtils.collation(sole(inputs).getTraitSet()); + + assert collation.satisfies(TraitUtils.collation(traitSet)) + : "Unexpected collations: input=" + collation + ", traitSet=" + TraitUtils.collation(traitSet); + return new IgniteReduceSortAggregate( getCluster(), traitSet, @@ -83,7 +87,7 @@ public class IgniteReduceSortAggregate extends IgniteReduceAggregateBase impleme groupSets, aggCalls, rowType, - TraitUtils.collation(traitSet) + collation ); } diff --git a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/trait/TraitUtils.java b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/trait/TraitUtils.java index b007a973315..ff27c6e41c6 100644 --- a/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/trait/TraitUtils.java +++ b/modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/trait/TraitUtils.java @@ -18,6 +18,8 @@ package org.apache.ignite.internal.processors.query.calcite.trait; import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.BitSet; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -29,6 +31,7 @@ import java.util.stream.Collectors; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import org.apache.calcite.linq4j.Ord; +import org.apache.calcite.plan.AbstractRelOptPlanner; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.RelOptPlanner; import org.apache.calcite.plan.RelOptRule; @@ -61,6 +64,7 @@ import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRel; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSort; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTableSpool; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTrimExchange; +import org.apache.ignite.internal.processors.query.calcite.util.Commons; import org.apache.ignite.internal.util.typedef.F; import org.jetbrains.annotations.Nullable; @@ -167,6 +171,7 @@ public class TraitUtils { RelOptRule.convert( rel, rel.getTraitSet() + .replace(RewindabilityTrait.ONE_WAY) .replace(CorrelationTrait.UNCORRELATED) ), toTrait); @@ -419,15 +424,58 @@ public class TraitUtils { assert traits.size() <= 1; + if (!traits.isEmpty() && traits.get(0).left.satisfies(requiredTraits)) { + // Return most relaxed parent traits. + return Pair.of(requiredTraits, traits.get(0).right); + } + return F.first(traits); } + /** */ + public static List<RelTraitSet> removeDuplicates(List<RelTraitSet> traits) { + BitSet duplicates = null; + + for (int i = 0; i < traits.size() - 1; i++) { + if (duplicates != null && duplicates.get(i)) + continue; + + for (int j = i + 1; j < traits.size(); j++) { + if (duplicates != null && duplicates.get(j)) + continue; + + // Return most strict child traits. + if (traits.get(i).satisfies(traits.get(j))) + (duplicates == null ? duplicates = new BitSet() : duplicates).set(j); + else if (traits.get(j).satisfies(traits.get(i))) { + (duplicates == null ? duplicates = new BitSet() : duplicates).set(i); + break; + } + } + } + + if (duplicates == null) + return traits; + + List<RelTraitSet> newTraits = new ArrayList<>(traits.size() - duplicates.cardinality()); + + for (int i = 0; i < traits.size(); i++) { + if (!duplicates.get(i)) + newTraits.add(traits.get(i)); + } + + return newTraits; + } + /** */ public static List<RelNode> derive(TraitsAwareIgniteRel rel, List<List<RelTraitSet>> inTraits) { assert !F.isEmpty(inTraits); RelTraitSet outTraits = rel.getCluster().traitSetOf(IgniteConvention.INSTANCE); - Set<Pair<RelTraitSet, List<RelTraitSet>>> combinations = combinations(outTraits, inTraits); + + inTraits = Commons.transform(inTraits, TraitUtils::removeDuplicates); + + Set<Pair<RelTraitSet, List<RelTraitSet>>> combinations = combinations(rel, outTraits, inTraits); if (combinations.isEmpty()) return ImmutableList.of(); @@ -448,14 +496,19 @@ public class TraitUtils { } /** */ - private static Set<Pair<RelTraitSet, List<RelTraitSet>>> combinations(RelTraitSet outTraits, List<List<RelTraitSet>> inTraits) { + private static Set<Pair<RelTraitSet, List<RelTraitSet>>> combinations( + TraitsAwareIgniteRel rel, + RelTraitSet outTraits, + List<List<RelTraitSet>> inTraits + ) { Set<Pair<RelTraitSet, List<RelTraitSet>>> out = new HashSet<>(); - fillRecursive(outTraits, inTraits, out, new RelTraitSet[inTraits.size()], 0); + fillRecursive(rel, outTraits, inTraits, out, new RelTraitSet[inTraits.size()], 0); return out; } /** */ private static boolean fillRecursive( + TraitsAwareIgniteRel rel, RelTraitSet outTraits, List<List<RelTraitSet>> inTraits, Set<Pair<RelTraitSet, List<RelTraitSet>>> result, @@ -463,6 +516,13 @@ public class TraitUtils { int idx ) throws ControlFlowException { boolean processed = false, last = idx == inTraits.size() - 1; + + if (last) { + assert rel.getCluster().getPlanner() instanceof AbstractRelOptPlanner; + + ((AbstractRelOptPlanner)rel.getCluster().getPlanner()).checkCancel(); + } + for (RelTraitSet t : inTraits.get(idx)) { assert t.getConvention() == IgniteConvention.INSTANCE; @@ -471,7 +531,7 @@ public class TraitUtils { if (last) result.add(Pair.of(outTraits, ImmutableList.copyOf(combination))); - else if (!fillRecursive(outTraits, inTraits, result, combination, idx + 1)) + else if (!fillRecursive(rel, outTraits, inTraits, result, combination, idx + 1)) return false; } return processed; diff --git a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/SortAggregateIntegrationTest.java b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/SortAggregateIntegrationTest.java index 28d72a94720..8d9f9cc8535 100644 --- a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/SortAggregateIntegrationTest.java +++ b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/SortAggregateIntegrationTest.java @@ -140,6 +140,21 @@ public class SortAggregateIntegrationTest extends AbstractBasicIntegrationTransa assertEquals(ROWS, cursors.size()); } + /** */ + @Test + public void testNullsReordering() { + sql("CREATE TABLE t(a INTEGER, b INTEGER) WITH " + atomicity()); + sql("INSERT INTO t VALUES (1, 1), (2, 2), (1, 3), (3, 4), (NULL, 1), (1, NULL)"); + + assertQuery("SELECT a, SUM(b), COUNT(b), COUNT(*) FROM t GROUP BY a ORDER BY a NULLS LAST") + .ordered() + .returns(1, 4L, 2L, 3L) + .returns(2, 2L, 1L, 1L) + .returns(3, 4L, 1L, 1L) + .returns(null, 1L, 1L, 1L) + .check(); + } + /** * @param c Cache. * @param rows Rows count. diff --git a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/TableDmlIntegrationTest.java b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/TableDmlIntegrationTest.java index 4ef41473fa6..f94c45162fb 100644 --- a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/TableDmlIntegrationTest.java +++ b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/TableDmlIntegrationTest.java @@ -23,6 +23,7 @@ import java.sql.Timestamp; import java.time.Duration; import java.time.Period; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.UUID; @@ -657,6 +658,20 @@ public class TableDmlIntegrationTest extends AbstractBasicIntegrationTransaction assertThrows("INSERT INTO timestamp_t VALUES ('1900-1-1 00-00-00')", errType, errDate); } + /** */ + @Test + public void testInsertMultiRowValues() { + sql("CREATE TABLE test (id int, val int) WITH " + atomicity()); + + int rowsCnt = 50; + + String sql = "INSERT INTO test VALUES " + String.join(", ", Collections.nCopies(rowsCnt, "(?, ?)")); + + sql(sql, new Object[rowsCnt * 2]); + + assertQuery("SELECT * FROM test").resultSize(rowsCnt).check(); + } + /** */ private void checkDefaultValue(String sqlType, String sqlVal, Object expectedVal) { try { diff --git a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/JoinColocationPlannerTest.java b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/JoinColocationPlannerTest.java index 5376a79e8bb..0b06c0ddd75 100644 --- a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/JoinColocationPlannerTest.java +++ b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/JoinColocationPlannerTest.java @@ -33,6 +33,7 @@ import org.apache.ignite.internal.processors.query.calcite.trait.IgniteDistribut import org.apache.ignite.internal.util.typedef.internal.CU; import org.junit.Test; +import static org.apache.calcite.sql.type.SqlTypeName.INTEGER; import static org.hamcrest.CoreMatchers.nullValue; import static org.junit.Assert.assertThat; @@ -209,4 +210,31 @@ public class JoinColocationPlannerTest extends AbstractPlannerTest { ) ); } + + /** + * Re-hashing right hand for merge join. + */ + @Test + public void joinMergeJoinAffinityRehash() throws Exception { + IgniteSchema schema = createSchema( + createTable("ORDERS", IgniteDistributions.affinity(0, "orders", "hash"), + "ID", INTEGER, "REGION", INTEGER) + .addIndex("ORDER_ID_IDX", 0), + createTable("ORDER_ITEMS", IgniteDistributions.affinity(0, "order_items", "hash"), + "ID", INTEGER, "ORDER_ID", INTEGER, "AMOUNT", INTEGER) + .addIndex("ORDER_ITEMS_ORDER_ID_IDX", 1) + ); + + String sql = "SELECT sum(amount)" + + " FROM order_items i JOIN orders o ON o.id=i.order_id" + + " WHERE o.region = ?"; + + assertPlan(sql, schema, + nodeOrAnyChild(isInstanceOf(IgniteMergeJoin.class)) + .and(hasChildThat(isIndexScan("ORDERS", "ORDER_ID_IDX"))) + .and(hasChildThat(isInstanceOf(IgniteExchange.class) + .and(hasDistribution(IgniteDistributions.affinity(0, "orders", "hash"))) + .and(hasChildThat(isIndexScan("ORDER_ITEMS", "ORDER_ITEMS_ORDER_ID_IDX"))))) + ); + } } diff --git a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/SortAggregatePlannerTest.java b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/SortAggregatePlannerTest.java index 9116092c0d7..2f0e13cc24c 100644 --- a/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/SortAggregatePlannerTest.java +++ b/modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/SortAggregatePlannerTest.java @@ -18,20 +18,24 @@ package org.apache.ignite.internal.processors.query.calcite.planner; import java.util.Arrays; +import java.util.function.Predicate; import org.apache.calcite.plan.RelOptPlanner; import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelCollation; import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.ignite.internal.processors.query.calcite.metadata.ColocationGroup; import org.apache.ignite.internal.processors.query.calcite.prepare.MappingQueryContext; -import org.apache.ignite.internal.processors.query.calcite.rel.IgniteAggregate; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteCorrelatedNestedLoopJoin; +import org.apache.ignite.internal.processors.query.calcite.rel.IgniteExchange; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteLimit; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteRel; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteSort; import org.apache.ignite.internal.processors.query.calcite.rel.IgniteTableScan; import org.apache.ignite.internal.processors.query.calcite.rel.agg.IgniteColocatedSortAggregate; +import org.apache.ignite.internal.processors.query.calcite.rel.agg.IgniteMapSortAggregate; import org.apache.ignite.internal.processors.query.calcite.rel.agg.IgniteReduceSortAggregate; import org.apache.ignite.internal.processors.query.calcite.schema.IgniteSchema; import org.apache.ignite.internal.processors.query.calcite.trait.IgniteDistribution; @@ -43,6 +47,10 @@ import org.apache.ignite.internal.util.typedef.F; import org.apache.ignite.testframework.GridTestUtils; import org.junit.Test; +import static org.apache.calcite.sql.type.SqlTypeName.INTEGER; +import static org.apache.ignite.internal.processors.query.calcite.trait.IgniteDistributions.random; +import static org.apache.ignite.internal.processors.query.calcite.trait.IgniteDistributions.single; + /** * */ @@ -208,7 +216,7 @@ public class SortAggregatePlannerTest extends AbstractAggregatePlannerTest { @Test public void testEmptyCollationPassThroughLimit() throws Exception { IgniteSchema publicSchema = createSchema( - createTable("TEST", IgniteDistributions.single(), "A", Integer.class)); + createTable("TEST", single(), "A", Integer.class)); assertPlan("SELECT (SELECT test.a FROM test t ORDER BY 1 LIMIT 1) FROM test", publicSchema, hasChildThat(isInstanceOf(IgniteCorrelatedNestedLoopJoin.class) @@ -220,65 +228,111 @@ public class SortAggregatePlannerTest extends AbstractAggregatePlannerTest { /** */ @Test public void testCollationPassThrough() throws Exception { - IgniteSchema publicSchema = createSchema( - createTable("TEST", IgniteDistributions.single(), "A", Integer.class, "B", Integer.class)); - - // Sort order equals to grouping set. - assertPlan("SELECT a, b, COUNT(*) FROM test GROUP BY a, b ORDER BY a, b", publicSchema, - isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) - .and(input(isTableScan("TEST"))))), - HASH_AGG_RULES - ); - - // Sort order equals to grouping set (permuted collation). - assertPlan("SELECT a, b, COUNT(*) FROM test GROUP BY a, b ORDER BY b, a", publicSchema, - isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(1, 0)))) - .and(input(isTableScan("TEST"))))), - HASH_AGG_RULES - ); + for (boolean colocated : F.asList(true, false)) { + log.info("Test colocated=" + colocated); - // Sort order is a subset of grouping set. - assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY a", publicSchema, - isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) - .and(input(isTableScan("TEST"))))), - HASH_AGG_RULES - ); + String[] disabledRules = colocated ? HASH_AGG_RULES : + F.concat(HASH_AGG_RULES, "ColocatedSortAggregateConverterRule"); - // Sort order is a subset of grouping set (permuted collation). - assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY b", publicSchema, - isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(1, 0)))) - .and(input(isTableScan("TEST"))))), - HASH_AGG_RULES - ); + IgniteSchema publicSchema = createSchema( + createTable("TEST", colocated ? single() : random(), "A", INTEGER, "B", INTEGER)); - // Sort order is a superset of grouping set (additional sorting required). - assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY a, b, cnt", publicSchema, - isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1, 2)))) - .and(input(isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) + // Sort order equals to grouping set. + assertPlan("SELECT a, b, COUNT(*) FROM test GROUP BY a, b ORDER BY a, b", publicSchema, + isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(0, 1)), + input(isInstanceOf(IgniteSort.class) .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) - .and(input(isTableScan("TEST"))))))), - HASH_AGG_RULES - ); - - // Sort order is not equals to grouping set (additional sorting required). - assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY cnt, b", publicSchema, - isInstanceOf(IgniteSort.class) - .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(2, 1)))) - .and(input(isInstanceOf(IgniteAggregate.class) - .and(input(isInstanceOf(IgniteSort.class) + .and(input(isTableScan("TEST"))))), + disabledRules + ); + + // Sort order equals to grouping set, but order is not default. + RelCollation expectedCollation = RelCollations.of( + TraitUtils.createFieldCollation(0, false), + TraitUtils.createFieldCollation(1, true) + ); + + assertPlan("SELECT a, b, COUNT(*) FROM test GROUP BY a, b ORDER BY a desc, b", publicSchema, + isAggregateWithCollation(colocated, expectedCollation, + input(isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(expectedCollation)) + .and(input(isTableScan("TEST"))))), + disabledRules + ); + + // Sort order equals to grouping set (permuted collation). + assertPlan("SELECT a, b, COUNT(*) FROM test GROUP BY a, b ORDER BY b, a", publicSchema, + isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(1, 0)), + input(isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(1, 0)))) + .and(input(isTableScan("TEST"))))), + disabledRules + ); + + // Sort order is a subset of grouping set. + assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY a", publicSchema, + isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(0, 1)), + input(isInstanceOf(IgniteSort.class) .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) - .and(input(isTableScan("TEST"))))))), - HASH_AGG_RULES - ); + .and(input(isTableScan("TEST"))))), + disabledRules + ); + + // Sort order is a subset of grouping set (permuted collation). + assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY b", publicSchema, + isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(1, 0)), + input(isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(1, 0)))) + .and(input(isTableScan("TEST"))))), + disabledRules + ); + + // Sort order is a superset of grouping set (additional sorting required). + assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY a, b, cnt", publicSchema, + isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1, 2)))) + .and(input(isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(0, 1)), + input(isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) + .and(input(isTableScan("TEST"))))))), + disabledRules + ); + + // Sort order is not equals to grouping set (additional sorting required). + assertPlan("SELECT a, b, COUNT(*) cnt FROM test GROUP BY a, b ORDER BY cnt, b", publicSchema, + isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(2, 1)))) + .and(input(isAggregateWithCollation(colocated, TraitUtils.createCollation(F.asList(0, 1)), + input(isInstanceOf(IgniteSort.class) + .and(s -> s.collation().equals(TraitUtils.createCollation(F.asList(0, 1)))) + .and(input(isTableScan("TEST"))))))), + disabledRules + ); + } + } + + /** + * Predicate builder to check aggregate with collation. + */ + protected Predicate<? extends IgniteRel> isAggregateWithCollation( + boolean colocated, + RelCollation collation, + Predicate<RelNode> predicate + ) { + if (colocated) { + return isInstanceOf(IgniteColocatedSortAggregate.class) + .and(a -> a.collation().satisfies(collation)) + .and(predicate); + } + else { + return isInstanceOf(IgniteReduceSortAggregate.class) + .and(a -> a.collation().satisfies(collation)) + .and(input(isInstanceOf(IgniteExchange.class) + .and(input(isInstanceOf(IgniteMapSortAggregate.class) + .and(a -> a.collation().satisfies(collation)) + .and(predicate) + )) + )); + } } }
