This is an automated email from the ASF dual-hosted git repository.
kgyrtkirk pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git
The following commit(s) were added to refs/heads/master by this push:
new 8849b62 [CALCITE-3880] Add SortExchange support to RelFieldTrimmer
(Krisztian Kasa)
8849b62 is described below
commit 8849b62e7baff84dcd4262f47809e87a38be7b10
Author: kkasa <[email protected]>
AuthorDate: Thu Mar 12 14:16:18 2020 +0100
[CALCITE-3880] Add SortExchange support to RelFieldTrimmer (Krisztian Kasa)
Close #1879
---
.../org/apache/calcite/rel/RelDistributions.java | 2 +-
.../apache/calcite/sql2rel/RelFieldTrimmer.java | 45 ++++++++
.../calcite/sql2rel/RelFieldTrimmerTest.java | 127 +++++++++++++++++++++
3 files changed, 173 insertions(+), 1 deletion(-)
diff --git a/core/src/main/java/org/apache/calcite/rel/RelDistributions.java
b/core/src/main/java/org/apache/calcite/rel/RelDistributions.java
index 1de39fd..f7a524f 100644
--- a/core/src/main/java/org/apache/calcite/rel/RelDistributions.java
+++ b/core/src/main/java/org/apache/calcite/rel/RelDistributions.java
@@ -141,7 +141,7 @@ public class RelDistributions {
return getTraitDef().canonize(
new RelDistributionImpl(type,
ImmutableIntList.copyOf(
- Mappings.apply((Mapping) mapping, keys))));
+ Mappings.apply2((Mapping) mapping, keys))));
}
public boolean satisfies(RelTrait trait) {
diff --git a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
index 456b4eb..f8f7da7 100644
--- a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
+++ b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
@@ -21,6 +21,7 @@ import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.rel.RelCollation;
import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelFieldCollation;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Aggregate;
@@ -33,6 +34,7 @@ import org.apache.calcite.rel.core.Project;
import org.apache.calcite.rel.core.RelFactories;
import org.apache.calcite.rel.core.SetOp;
import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.SortExchange;
import org.apache.calcite.rel.core.TableScan;
import org.apache.calcite.rel.logical.LogicalTableFunctionScan;
import org.apache.calcite.rel.logical.LogicalTableModify;
@@ -548,6 +550,49 @@ public class RelFieldTrimmer implements ReflectiveVisitor {
return result(relBuilder.build(), inputMapping);
}
+ public TrimResult trimFields(
+ SortExchange sortExchange,
+ ImmutableBitSet fieldsUsed,
+ Set<RelDataTypeField> extraFields) {
+ final RelDataType rowType = sortExchange.getRowType();
+ final int fieldCount = rowType.getFieldCount();
+ final RelCollation collation = sortExchange.getCollation();
+ final RelDistribution distribution = sortExchange.getDistribution();
+ final RelNode input = sortExchange.getInput();
+
+ // We use the fields used by the consumer, plus any fields used as
sortExchange
+ // keys.
+ final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
+ for (RelFieldCollation field : collation.getFieldCollations()) {
+ inputFieldsUsed.set(field.getFieldIndex());
+ }
+ for (int keyIndex : distribution.getKeys()) {
+ inputFieldsUsed.set(keyIndex);
+ }
+
+ // Create input with trimmed columns.
+ final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
+ TrimResult trimResult =
+ trimChild(sortExchange, input, inputFieldsUsed.build(),
inputExtraFields);
+ RelNode newInput = trimResult.left;
+ final Mapping inputMapping = trimResult.right;
+
+ // If the input is unchanged, and we need to project all columns,
+ // there's nothing we can do.
+ if (newInput == input
+ && inputMapping.isIdentity()
+ && fieldsUsed.cardinality() == fieldCount) {
+ return result(sortExchange, Mappings.createIdentity(fieldCount));
+ }
+
+ relBuilder.push(newInput);
+ RelCollation newCollation = RexUtil.apply(inputMapping, collation);
+ RelDistribution newDistribution = distribution.apply(inputMapping);
+ relBuilder.sortExchange(newDistribution, newCollation);
+
+ return result(relBuilder.build(), inputMapping);
+ }
+
/**
* Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
* {@link org.apache.calcite.rel.logical.LogicalJoin}.
diff --git
a/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
b/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
new file mode 100644
index 0000000..d602433
--- /dev/null
+++ b/core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java
@@ -0,0 +1,127 @@
+/*
+ * 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.sql2rel;
+
+import org.apache.calcite.plan.RelTraitDef;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelDistributions;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.schema.SchemaPlus;
+import org.apache.calcite.sql.parser.SqlParser;
+import org.apache.calcite.test.CalciteAssert;
+import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.tools.Programs;
+import org.apache.calcite.tools.RelBuilder;
+
+import com.google.common.collect.Lists;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.apache.calcite.test.Matchers.hasTree;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+
+class RelFieldTrimmerTest {
+ public static Frameworks.ConfigBuilder config() {
+ final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
+ return Frameworks.newConfigBuilder()
+ .parserConfig(SqlParser.Config.DEFAULT)
+ .defaultSchema(
+ CalciteAssert.addSchema(rootSchema,
CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL))
+ .traitDefs((List<RelTraitDef>) null)
+ .programs(Programs.heuristicJoinOrder(Programs.RULE_SET, true, 2));
+ }
+
+ @Test public void testSortExchangeFieldTrimmer() {
+ final RelBuilder builder = RelBuilder.create(config().build());
+ final RelNode root =
+ builder.scan("EMP")
+ .project(builder.field("EMPNO"), builder.field("ENAME"),
builder.field("DEPTNO"))
+ .sortExchange(RelDistributions.hash(Lists.newArrayList(1)),
RelCollations.of(0))
+ .project(builder.field("EMPNO"), builder.field("ENAME"))
+ .build();
+
+ RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+ RelNode trimmed = fieldTrimmer.trim(root);
+
+ final String expected = ""
+ + "LogicalSortExchange(distribution=[hash[1]], collation=[[0]])\n"
+ + " LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
+ + " LogicalTableScan(table=[[scott, EMP]])\n";
+ assertThat(trimmed, hasTree(expected));
+ }
+
+ @Test public void testSortExchangeFieldTrimmerWhenProjectCannotBeMerged() {
+ final RelBuilder builder = RelBuilder.create(config().build());
+ final RelNode root =
+ builder.scan("EMP")
+ .project(builder.field("EMPNO"), builder.field("ENAME"),
builder.field("DEPTNO"))
+ .sortExchange(RelDistributions.hash(Lists.newArrayList(1)),
RelCollations.of(0))
+ .project(builder.field("EMPNO"))
+ .build();
+
+ RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+ RelNode trimmed = fieldTrimmer.trim(root);
+
+ final String expected = ""
+ + "LogicalProject(EMPNO=[$0])\n"
+ + " LogicalSortExchange(distribution=[hash[1]], collation=[[0]])\n"
+ + " LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
+ + " LogicalTableScan(table=[[scott, EMP]])\n";
+ assertThat(trimmed, hasTree(expected));
+ }
+
+ @Test public void testSortExchangeFieldTrimmerWithEmptyCollation() {
+ final RelBuilder builder = RelBuilder.create(config().build());
+ final RelNode root =
+ builder.scan("EMP")
+ .project(builder.field("EMPNO"), builder.field("ENAME"),
builder.field("DEPTNO"))
+ .sortExchange(RelDistributions.hash(Lists.newArrayList(1)),
RelCollations.EMPTY)
+ .project(builder.field("EMPNO"), builder.field("ENAME"))
+ .build();
+
+ RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+ RelNode trimmed = fieldTrimmer.trim(root);
+
+ final String expected = ""
+ + "LogicalSortExchange(distribution=[hash[1]], collation=[[]])\n"
+ + " LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
+ + " LogicalTableScan(table=[[scott, EMP]])\n";
+ assertThat(trimmed, hasTree(expected));
+ }
+
+ @Test public void testSortExchangeFieldTrimmerWithSingletonDistribution() {
+ final RelBuilder builder = RelBuilder.create(config().build());
+ final RelNode root =
+ builder.scan("EMP")
+ .project(builder.field("EMPNO"), builder.field("ENAME"),
builder.field("DEPTNO"))
+ .sortExchange(RelDistributions.SINGLETON, RelCollations.of(0))
+ .project(builder.field("EMPNO"), builder.field("ENAME"))
+ .build();
+
+ RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+ RelNode trimmed = fieldTrimmer.trim(root);
+
+ final String expected = ""
+ + "LogicalSortExchange(distribution=[single], collation=[[0]])\n"
+ + " LogicalProject(EMPNO=[$0], ENAME=[$1])\n"
+ + " LogicalTableScan(table=[[scott, EMP]])\n";
+ assertThat(trimmed, hasTree(expected));
+ }
+}
