Skip to content

Commit 25e3ae0

Browse files
tjbangharttjbanghart
committed
[CALCITE-7254] Add rule for sharing trivially equivalent RelNodes within Combine
1 parent df3de96 commit 25e3ae0

File tree

9 files changed

+1582
-1
lines changed

9 files changed

+1582
-1
lines changed

‎core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableCombine.java‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to you under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.calcite.adapter.enumerable;
18+
19+
import org.apache.calcite.linq4j.Ord;
20+
import org.apache.calcite.linq4j.tree.BlockBuilder;
21+
import org.apache.calcite.linq4j.tree.Expression;
22+
import org.apache.calcite.linq4j.tree.Expressions;
23+
import org.apache.calcite.linq4j.tree.ParameterExpression;
24+
import org.apache.calcite.linq4j.tree.Types;
25+
import org.apache.calcite.plan.RelOptCluster;
26+
import org.apache.calcite.plan.RelTraitSet;
27+
import org.apache.calcite.rel.RelNode;
28+
import org.apache.calcite.rel.core.Combine;
29+
import org.apache.calcite.rel.type.RelDataType;
30+
import org.apache.calcite.util.BuiltInMethod;
31+
32+
import java.util.ArrayList;
33+
import java.util.List;
34+
35+
/** Implementation of {@link org.apache.calcite.rel.core.Combine} in
36+
* {@link org.apache.calcite.adapter.enumerable.EnumerableConvention enumerable calling convention}. */
37+
public class EnumerableCombine extends Combine implements EnumerableRel {
38+
public EnumerableCombine(RelOptCluster cluster, RelTraitSet traitSet,
39+
List<RelNode> inputs) {
40+
super(cluster, traitSet, inputs);
41+
}
42+
43+
@Override public EnumerableCombine copy(RelTraitSet traitSet, List<RelNode> inputs) {
44+
return new EnumerableCombine(getCluster(), traitSet, inputs);
45+
}
46+
47+
@Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
48+
final BlockBuilder builder = new BlockBuilder();
49+
final RelDataType rowType = getRowType();
50+
final List<Expression> fieldExpressions = new ArrayList<>();
51+
52+
// Implement each input and collect their results
53+
// Convert each Enumerable to a List since the row type is STRUCT<QUERY_0: ARRAY<...>, ...>
54+
for (Ord<RelNode> ord : Ord.zip(inputs)) {
55+
EnumerableRel input = (EnumerableRel) ord.e;
56+
final Result result = implementor.visitChild(this, ord.i, input, pref);
57+
Expression childExp =
58+
builder.append(
59+
"child" + ord.i,
60+
result.block);
61+
62+
// Check if the input has multiple columns (row is Object[])
63+
// If so, convert each row to a List for proper toString() formatting
64+
final int fieldCount = input.getRowType().getFieldCount();
65+
Expression enumerableToConvert;
66+
if (fieldCount > 1) {
67+
// Multi-column rows: transform each Object[] to List
68+
// child.select(row -> Arrays.asList((Object[]) row))
69+
ParameterExpression row = Expressions.parameter(Object.class, "row" + ord.i);
70+
Expression selectLambda =
71+
Expressions.lambda(
72+
Expressions.call(
73+
BuiltInMethod.ARRAYS_AS_LIST.method,
74+
Expressions.convert_(row, Object[].class)),
75+
row);
76+
enumerableToConvert =
77+
builder.append("converted" + ord.i,
78+
Expressions.call(
79+
childExp,
80+
BuiltInMethod.SELECT.method,
81+
selectLambda));
82+
} else {
83+
// Single column: use directly
84+
enumerableToConvert = childExp;
85+
}
86+
87+
// Convert Enumerable to List for the array field
88+
// Each field in the struct is an ARRAY type, which maps to List in Java
89+
Expression listExp =
90+
builder.append(
91+
"list" + ord.i,
92+
Expressions.call(
93+
enumerableToConvert,
94+
Types.lookupMethod(
95+
org.apache.calcite.linq4j.Enumerable.class,
96+
"toList")));
97+
fieldExpressions.add(listExp);
98+
}
99+
100+
// The physical type represents the struct of all query results
101+
final PhysType physType =
102+
PhysTypeImpl.of(
103+
implementor.getTypeFactory(),
104+
rowType,
105+
pref.prefer(JavaRowFormat.CUSTOM));
106+
107+
// Create the struct record with the list fields
108+
final Expression recordExpression =
109+
physType.record(fieldExpressions);
110+
111+
// Return a singleton enumerable containing the struct as the single row
112+
// The Combine operator produces exactly one row containing all the child results
113+
final Expression singletonExpression =
114+
Expressions.call(
115+
Types.lookupMethod(
116+
org.apache.calcite.linq4j.Linq4j.class,
117+
"singletonEnumerable",
118+
Object.class),
119+
recordExpression);
120+
121+
builder.add(singletonExpression);
122+
123+
return implementor.result(physType, builder.toBlock());
124+
}
125+
}

‎core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableCombineRule.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to you under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.calcite.adapter.enumerable;
18+
19+
import org.apache.calcite.plan.Convention;
20+
import org.apache.calcite.plan.RelTraitSet;
21+
import org.apache.calcite.rel.RelNode;
22+
import org.apache.calcite.rel.convert.ConverterRule;
23+
import org.apache.calcite.rel.core.Combine;
24+
import org.apache.calcite.util.Util;
25+
26+
import java.util.List;
27+
28+
/**
29+
* Rule to convert a {@link Combine} to an {@link EnumerableCombine}.
30+
*
31+
* @see EnumerableRules#ENUMERABLE_COMBINE_RULE
32+
*/
33+
class EnumerableCombineRule extends ConverterRule {
34+
/** Default configuration. */
35+
static final Config DEFAULT_CONFIG = Config.INSTANCE
36+
.withConversion(Combine.class, Convention.NONE,
37+
EnumerableConvention.INSTANCE, "EnumerableCombineRule")
38+
.withRuleFactory(EnumerableCombineRule::new);
39+
40+
/** Called from the Config. */
41+
protected EnumerableCombineRule(Config config) {
42+
super(config);
43+
}
44+
45+
@Override public RelNode convert(RelNode rel) {
46+
final Combine combine = (Combine) rel;
47+
final EnumerableConvention out = EnumerableConvention.INSTANCE;
48+
final RelTraitSet traitSet = rel.getCluster().traitSet().replace(out);
49+
final List<RelNode> newInputs =
50+
Util.transform(combine.getInputs(), n -> convert(n, traitSet));
51+
return new EnumerableCombine(rel.getCluster(), traitSet, newInputs);
52+
}
53+
}

‎core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableRules.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,11 @@ private EnumerableRules() {
103103
public static final EnumerableUnionRule ENUMERABLE_UNION_RULE =
104104
EnumerableUnionRule.DEFAULT_CONFIG.toRule(EnumerableUnionRule.class);
105105

106+
/** Rule that converts a {@link org.apache.calcite.rel.core.Combine}
107+
* to an {@link EnumerableCombine}. */
108+
public static final EnumerableCombineRule ENUMERABLE_COMBINE_RULE =
109+
EnumerableCombineRule.DEFAULT_CONFIG.toRule(EnumerableCombineRule.class);
110+
106111
/** Rule that converts a {@link LogicalRepeatUnion} into an
107112
* {@link EnumerableRepeatUnion}. */
108113
public static final EnumerableRepeatUnionRule ENUMERABLE_REPEAT_UNION_RULE =
@@ -224,6 +229,7 @@ private EnumerableRules() {
224229
EnumerableRules.ENUMERABLE_UNCOLLECT_RULE,
225230
EnumerableRules.ENUMERABLE_MERGE_UNION_RULE,
226231
EnumerableRules.ENUMERABLE_UNION_RULE,
232+
EnumerableRules.ENUMERABLE_COMBINE_RULE,
227233
EnumerableRules.ENUMERABLE_REPEAT_UNION_RULE,
228234
EnumerableRules.ENUMERABLE_TABLE_SPOOL_RULE,
229235
EnumerableRules.ENUMERABLE_INTERSECT_RULE,

‎core/src/main/java/org/apache/calcite/plan/SpoolRelOptTable.java‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to you under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.calcite.plan;
18+
19+
import org.apache.calcite.linq4j.tree.Expression;
20+
import org.apache.calcite.rel.RelCollation;
21+
import org.apache.calcite.rel.RelDistribution;
22+
import org.apache.calcite.rel.RelDistributions;
23+
import org.apache.calcite.rel.RelNode;
24+
import org.apache.calcite.rel.RelReferentialConstraint;
25+
import org.apache.calcite.rel.type.RelDataType;
26+
import org.apache.calcite.rel.type.RelDataTypeField;
27+
import org.apache.calcite.schema.ColumnStrategy;
28+
import org.apache.calcite.schema.Statistic;
29+
import org.apache.calcite.schema.Statistics;
30+
import org.apache.calcite.schema.Table;
31+
import org.apache.calcite.schema.impl.ListTransientTable;
32+
import org.apache.calcite.util.ImmutableBitSet;
33+
34+
import com.google.common.collect.ImmutableList;
35+
36+
import org.checkerframework.checker.nullness.qual.Nullable;
37+
38+
import java.util.Collections;
39+
import java.util.List;
40+
41+
/**
42+
* Implementation of {@link RelOptTable} for temporary spool tables.
43+
*
44+
* <p>This table represents temporary storage used by spool operators
45+
* during query execution. It's used for planning purposes only and
46+
* will be converted to appropriate physical operators later.
47+
*/
48+
public class SpoolRelOptTable implements RelOptTable {
49+
private final @Nullable RelOptSchema schema;
50+
private final RelDataType rowType;
51+
private final String name;
52+
private final double rowCount;
53+
private final Table table;
54+
55+
/**
56+
* Creates a SpoolRelOptTable with explicit row count.
57+
*
58+
* @param schema the schema this table belongs to (can be null for temporary tables)
59+
* @param rowType the row type of the data that will be stored in this spool
60+
* @param name optional name for the spool table
61+
* @param rowCount the estimated number of rows that will be materialized in this spool
62+
*/
63+
public SpoolRelOptTable(
64+
@Nullable RelOptSchema schema,
65+
RelDataType rowType,
66+
String name,
67+
double rowCount) {
68+
this.schema = schema;
69+
this.rowType = rowType;
70+
this.name = name;
71+
this.rowCount = rowCount;
72+
// Use standard ListTransientTable with custom statistics for accurate cost estimation
73+
this.table = new ListTransientTable(name, rowType) {
74+
@Override public Statistic getStatistic() {
75+
return Statistics.of(rowCount, ImmutableList.of());
76+
}
77+
};
78+
}
79+
80+
@Override public RelNode toRel(ToRelContext context) {
81+
// This shouldn't be called during planning - spools are created differently
82+
throw new UnsupportedOperationException("SpoolRelOptTable.toRel should not be called");
83+
}
84+
85+
@Override public List<String> getQualifiedName() {
86+
return ImmutableList.of("TEMP", name);
87+
}
88+
89+
@Override public double getRowCount() {
90+
// Return the actual row count of the materialized data in this spool
91+
return rowCount;
92+
}
93+
94+
@Override public RelDataType getRowType() {
95+
return rowType;
96+
}
97+
98+
@Override public @Nullable RelOptSchema getRelOptSchema() {
99+
return schema;
100+
}
101+
102+
@Override public @Nullable RelDistribution getDistribution() {
103+
return RelDistributions.ANY;
104+
}
105+
106+
@Override public @Nullable List<ImmutableBitSet> getKeys() {
107+
// Spools typically don't have keys
108+
return ImmutableList.of();
109+
}
110+
111+
@Override public @Nullable List<RelReferentialConstraint> getReferentialConstraints() {
112+
// Temporary tables don't have referential constraints
113+
return ImmutableList.of();
114+
}
115+
116+
@Override public @Nullable List<RelCollation> getCollationList() {
117+
// Could be extended to preserve collations from the input
118+
return ImmutableList.of();
119+
}
120+
121+
@Override public boolean isKey(ImmutableBitSet columns) {
122+
return false;
123+
}
124+
125+
@Override public @Nullable Expression getExpression(Class clazz) {
126+
// Return null so EnumerableTableScanRule won't try to convert spool table scans
127+
// Spool table scans are handled within the spool operator itself
128+
return null;
129+
}
130+
131+
@Override public RelOptTable extend(List<RelDataTypeField> extendedFields) {
132+
throw new UnsupportedOperationException("SpoolRelOptTable.extend should not be called");
133+
}
134+
135+
@Override public List<ColumnStrategy> getColumnStrategies() {
136+
return Collections.emptyList();
137+
}
138+
139+
@Override public <C> @Nullable C unwrap(Class<C> aClass) {
140+
if (aClass.isInstance(table)) {
141+
return aClass.cast(table);
142+
}
143+
return null;
144+
}
145+
}

‎core/src/main/java/org/apache/calcite/rel/core/Combine.java‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
* This is used for multi-root optimization in the VolcanoPlanner.
3838
*/
3939
public class Combine extends AbstractRelNode {
40-
protected final ImmutableList<RelNode> inputs;
40+
protected ImmutableList<RelNode> inputs;
4141

4242
/** Creates a Combine. */
4343
public static Combine create(RelOptCluster cluster, RelTraitSet traitSet, List<RelNode> inputs) {
@@ -54,6 +54,25 @@ public Combine(RelOptCluster cluster, RelTraitSet traitSet, List<RelNode> inputs
5454
return inputs;
5555
}
5656

57+
@Override public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
58+
return new Combine(getCluster(), traitSet, inputs);
59+
}
60+
61+
@Override public void replaceInput(int ordinalInParent, RelNode rel) {
62+
// Combine has multiple inputs stored in an immutable list.
63+
// To replace an input, we need to create a new list with the replacement.
64+
ImmutableList.Builder<RelNode> newInputs = ImmutableList.builder();
65+
for (int i = 0; i < inputs.size(); i++) {
66+
if (i == ordinalInParent) {
67+
newInputs.add(rel);
68+
} else {
69+
newInputs.add(inputs.get(i));
70+
}
71+
}
72+
inputs = newInputs.build();
73+
}
74+
75+
5776
@Override public RelWriter explainTerms(RelWriter pw) {
5877
super.explainTerms(pw);
5978
for (Ord<RelNode> ord : Ord.zip(inputs)) {

0 commit comments

Comments
 (0)