(ignite-3) 02/02: Fix Collation pass-through for sorted map aggregate

[am...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

amashenkov pushed a commit to branch ignite-21897
in repository https://gitbox.apache.org/repos/asf/ignite-3.git

commit 5746d704f3d5adef638ee9d7543c13489745616f
Author: amashenkov <an...@gmail.com>
AuthorDate: Wed Mar 27 21:09:35 2024 +0300

    Fix Collation pass-through for sorted map aggregate
---
 .../engine/rel/agg/IgniteSortAggregateBase.java    | 49 ++++++++++++++++------
 1 file changed, 36 insertions(+), 13 deletions(-)

diff --git a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/rel/agg/IgniteSortAggregateBase.java b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/rel/agg/IgniteSortAggregateBase.java
index d763c49b95..e5e88d9e18 100644
--- a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/rel/agg/IgniteSortAggregateBase.java
+++ b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/rel/agg/IgniteSortAggregateBase.java
@@ -17,6 +17,7 @@
 
 package org.apache.ignite.internal.sql.engine.rel.agg;
 
+import static java.util.function.Predicate.not;
 import static org.apache.ignite.internal.sql.engine.util.Commons.maxPrefix;
 
 import it.unimi.dsi.fastutil.ints.IntList;
@@ -29,9 +30,11 @@ import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.util.ImmutableBitSet;
 import org.apache.calcite.util.Pair;
+import org.apache.calcite.util.mapping.Mapping;
 import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
 import org.apache.ignite.internal.sql.engine.trait.TraitUtils;
 import org.apache.ignite.internal.sql.engine.trait.TraitsAwareIgniteRel;
+import org.apache.ignite.internal.sql.engine.util.Commons;
 
 /**
  * Defines common methods for {@link IgniteRel relational nodes} that implement sort-base aggregates.
@@ -77,24 +80,35 @@ interface IgniteSortAggregateBase extends TraitsAwareIgniteRel {
             RelTraitSet nodeTraits, List<RelTraitSet> inputTraits
     ) {
         RelCollation required = TraitUtils.collation(nodeTraits);
-        ImmutableBitSet requiredKeys = ImmutableBitSet.of(required.getKeys());
-        RelCollation collation;
 
-        if (getGroupSet().contains(requiredKeys)) {
-            List<RelFieldCollation> newCollationFields = new ArrayList<>(getGroupSet().cardinality());
-            newCollationFields.addAll(required.getFieldCollations());
-
-            ImmutableBitSet keysLeft = getGroupSet().except(requiredKeys);
+        if (getGroupSet().isEmpty()) {
+            return passThroughCollation(nodeTraits, inputTraits, RelCollations.EMPTY); // Erase collation for a single group.
+        } else if (required.getFieldCollations().isEmpty()) {
+            return passThroughCollation(nodeTraits, inputTraits, TraitUtils.createCollation(getGroupSet().asList())); // No match.
+        }
 
-            keysLeft.forEach(fieldIdx -> newCollationFields.add(TraitUtils.createFieldCollation(fieldIdx)));
+        ImmutableBitSet groupingColumns = ImmutableBitSet.range(getGroupSet().cardinality());
+        IntList prefix = maxPrefix(required.getKeys(), groupingColumns.asSet());
 
-            collation = RelCollations.of(newCollationFields);
-        } else {
-            collation = TraitUtils.createCollation(getGroupSet().toList());
+        if (prefix.isEmpty()) {
+            return passThroughCollation(nodeTraits, inputTraits, TraitUtils.createCollation(getGroupSet().asList())); // No match.
         }
 
-        return Pair.of(nodeTraits.replace(collation),
-                List.of(inputTraits.get(0).replace(collation)));
+        // Rearrange grouping columns to satisfy required collation (as much as possible).
+        List<RelFieldCollation> newCollationColumns = new ArrayList<>(getGroupSet().cardinality());
+        Mapping mapping = Commons.trimmingMapping(groupingColumns.cardinality(), groupingColumns);
+
+        // Add required columns first.
+        prefix.intStream().map(mapping::getTarget)
+                .forEach(fieldIdx -> newCollationColumns.add(TraitUtils.createFieldCollation(fieldIdx)));
+
+        // Then add missed grouping columns.
+        groupingColumns.asList().stream()
+                .filter(not(prefix::contains))
+                .map(mapping::getTarget)
+                .forEach(fieldIdx -> newCollationColumns.add(TraitUtils.createFieldCollation(fieldIdx)));
+
+        return passThroughCollation(nodeTraits, inputTraits, RelCollations.of(newCollationColumns));
     }
 
     /** {@inheritDoc} */
@@ -118,4 +132,13 @@ interface IgniteSortAggregateBase extends TraitsAwareIgniteRel {
                 inputTraits
         ));
     }
+
+    private static Pair<RelTraitSet, List<RelTraitSet>> passThroughCollation(
+            RelTraitSet nodeTraits,
+            List<RelTraitSet> inputTraits,
+            RelCollation collation
+    ) {
+        return Pair.of(nodeTraits.replace(collation),
+                List.of(inputTraits.get(0).replace(collation)));
+    }
 }