This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 2c0a465f884 Compare all semantic fields when spooling equivalent
stages (#19219)
2c0a465f884 is described below
commit 2c0a465f884ae9fab44ab9fc0a50ec7d90c20ae7
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Aug 26 16:32:01 2026 -0400
Compare all semantic fields when spooling equivalent stages (#19219)
---
.../planner/logical/EquivalentStagesFinder.java | 9 ++
.../pinot/query/planner/logical/RexExpression.java | 7 +-
.../logical/EquivalentStagesFinderTest.java | 134 +++++++++++++++++++++
.../query/planner/logical/StagesTestBase.java | 63 ++++++++++
.../src/test/resources/queries/Spool.json | 16 +++
5 files changed, 226 insertions(+), 3 deletions(-)
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java
index 00e0daf6cb6..2ae70264e5c 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinder.java
@@ -184,6 +184,9 @@ public class EquivalentStagesFinder {
return areBaseNodesEquivalent(node1, node2) &&
Objects.equals(node1.getAggCalls(), that.getAggCalls())
&& Objects.equals(node1.getFilterArgs(), that.getFilterArgs())
&& Objects.equals(node1.getGroupKeys(), that.getGroupKeys())
+ // Group keys only hold the union of the grouping columns, so two
different sets of grouping sets over the
+ // same columns are indistinguishable without this check (and
produce the same data schema).
+ && Objects.equals(node1.getGroupingSets(), that.getGroupingSets())
&& node1.getAggType() == that.getAggType()
&& node1.isLeafReturnFinalResult() ==
that.isLeafReturnFinalResult()
&& Objects.equals(node1.getCollations(), that.getCollations())
@@ -245,6 +248,9 @@ public class EquivalentStagesFinder {
&& Objects.equals(node1.getLeftKeys(), that.getLeftKeys())
&& Objects.equals(node1.getRightKeys(), that.getRightKeys())
&& Objects.equals(node1.getNonEquiConditions(),
that.getNonEquiConditions())
+ // ASOF joins keep their whole comparison here rather than in the
non-equi conditions, so without this
+ // check two ASOF joins differing only on the match condition look
equivalent.
+ && Objects.equals(node1.getMatchCondition(),
that.getMatchCondition())
&& node1.getJoinStrategy() == that.getJoinStrategy();
}
@@ -318,6 +324,9 @@ public class EquivalentStagesFinder {
&& Objects.equals(node1.getKeys(), that.getKeys())
&& Objects.equals(node1.getCollations(), that.getCollations())
&& node1.getWindowFrameType() == that.getWindowFrameType()
+ // The frame exclusion changes which rows feed the window
function, so two windows that only differ on it
+ // compute different values and must not be spooled together.
+ && node1.getExclude() == that.getExclude()
&& Objects.equals(node1.getConstants(), that.getConstants());
}
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java
index 0cf11775f5b..2d4c03a503a 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpression.java
@@ -151,13 +151,14 @@ public interface RexExpression {
return false;
}
FunctionCall that = (FunctionCall) o;
- return _isDistinct == that._isDistinct && _dataType == that._dataType &&
Objects.equals(_functionName,
- that._functionName) && Objects.equals(_functionOperands,
that._functionOperands);
+ return _isDistinct == that._isDistinct && _ignoreNulls ==
that._ignoreNulls && _dataType == that._dataType
+ && Objects.equals(_functionName, that._functionName)
+ && Objects.equals(_functionOperands, that._functionOperands);
}
@Override
public int hashCode() {
- return Objects.hash(_dataType, _functionName, _functionOperands,
_isDistinct);
+ return Objects.hash(_dataType, _functionName, _functionOperands,
_isDistinct, _ignoreNulls);
}
}
}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java
index 1e7f71f3960..45707bb9a4f 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/EquivalentStagesFinderTest.java
@@ -18,10 +18,13 @@
*/
package org.apache.pinot.query.planner.logical;
+import java.util.List;
import java.util.Map;
import org.apache.calcite.rel.RelDistribution;
+import org.apache.calcite.sql.SqlKind;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.planner.plannode.MailboxSendNode;
+import org.apache.pinot.query.planner.plannode.WindowNode.WindowExclusion;
import org.testng.annotations.Test;
import static org.testng.Assert.*;
@@ -153,6 +156,137 @@ public class EquivalentStagesFinderTest extends
StagesTestBase {
assertEquals(result.toString(), "[[0], [1], [2]]");
}
+ @Test
+ public void sameWindowKeepEquivalence() {
+ when(
+ join(
+ exchange(1, window(tableScan("T1"), lastValueRespectingNulls())),
+ exchange(2, window(tableScan("T1"), lastValueRespectingNulls()))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1, 2]]");
+ }
+
+ /// Two windows that both ignore nulls are still equivalent: the check must
compare the flag, not reject it.
+ @Test
+ public void sameIgnoreNullsKeepEquivalence() {
+ when(
+ join(
+ exchange(1, window(tableScan("T1"), lastValueIgnoringNulls())),
+ exchange(2, window(tableScan("T1"), lastValueIgnoringNulls()))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1, 2]]");
+ }
+
+ /// Same for a non-default frame exclusion.
+ @Test
+ public void sameWindowExclusionKeepEquivalence() {
+ when(
+ join(
+ exchange(1, window(tableScan("T1"), lastValueRespectingNulls(),
WindowExclusion.CURRENT_ROW)),
+ exchange(2, window(tableScan("T1"), lastValueRespectingNulls(),
WindowExclusion.CURRENT_ROW))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1, 2]]");
+ }
+
+ /// A window that ignores nulls computes different values than one that
respects them, so the two stages must not be
+ /// treated as equivalent.
+ @Test
+ public void differentIgnoreNullsBreakEquivalence() {
+ when(
+ join(
+ exchange(1, window(tableScan("T1"), lastValueIgnoringNulls())),
+ exchange(2, window(tableScan("T1"), lastValueRespectingNulls()))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1], [2]]");
+ }
+
+ /// The frame exclusion changes which rows feed the window function, so it
must break equivalence too.
+ @Test
+ public void differentWindowExclusionBreakEquivalence() {
+ when(
+ join(
+ exchange(1, window(tableScan("T1"), lastValueRespectingNulls(),
WindowExclusion.CURRENT_ROW)),
+ exchange(2, window(tableScan("T1"), lastValueRespectingNulls(),
WindowExclusion.NO_OTHERS))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1], [2]]");
+ }
+
+ /// The group keys only hold the union of the grouping columns, so
`ROLLUP(col1, col2)` and `CUBE(col1, col2)` agree
+ /// on every other field (including the data schema) and are only told apart
by the grouping sets themselves.
+ @Test
+ public void differentGroupingSetsBreakEquivalence() {
+ when(
+ join(
+ exchange(1, aggregate(tableScan("T1"), List.of(List.of(0, 1),
List.of(0), List.of()))),
+ exchange(2, aggregate(tableScan("T1"), List.of(List.of(0, 1),
List.of(0), List.of(1), List.of())))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1], [2]]");
+ }
+
+ @Test
+ public void sameGroupingSetsKeepEquivalence() {
+ List<List<Integer>> rollup = List.of(List.of(0, 1), List.of(0), List.of());
+ when(
+ join(
+ exchange(1, aggregate(tableScan("T1"), rollup)),
+ exchange(2, aggregate(tableScan("T1"), rollup))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1, 2]]");
+ }
+
+ /// ASOF joins carry their comparison in the match condition rather than in
the non-equi conditions, so two joins
+ /// that only differ there must not be spooled together.
+ @Test
+ public void differentAsofMatchConditionBreakEquivalence() {
+ when(
+ join(
+ exchange(1, asofJoin(tableScan("T1"), tableScan("T2"),
greaterThan())),
+ exchange(2, asofJoin(tableScan("T1"), tableScan("T2"),
greaterThanOrEqual()))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1], [2]]");
+ }
+
+ @Test
+ public void sameAsofMatchConditionKeepEquivalence() {
+ when(
+ join(
+ exchange(1, asofJoin(tableScan("T1"), tableScan("T2"),
greaterThan())),
+ exchange(2, asofJoin(tableScan("T1"), tableScan("T2"),
greaterThan()))
+ )
+ );
+ GroupedStages result =
EquivalentStagesFinder.findEquivalentStages(stage(0));
+ assertEquals(result.toString(), "[[0], [1, 2]]");
+ }
+
+ private static RexExpression greaterThan() {
+ return comparison(SqlKind.GREATER_THAN);
+ }
+
+ private static RexExpression greaterThanOrEqual() {
+ return comparison(SqlKind.GREATER_THAN_OR_EQUAL);
+ }
+
+ private static RexExpression comparison(SqlKind kind) {
+ return new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN,
kind.name(),
+ List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(1)));
+ }
+
@Test
public void differentDataSchemaBreakEquivalence() {
when(
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java
index 3475f42c70a..3c0e6617616 100644
---
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/StagesTestBase.java
@@ -34,12 +34,14 @@ import org.apache.calcite.rel.core.JoinRelType;
import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.query.planner.partitioning.KeySelector;
+import org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.JoinNode;
import org.apache.pinot.query.planner.plannode.MailboxReceiveNode;
import org.apache.pinot.query.planner.plannode.MailboxSendNode;
import org.apache.pinot.query.planner.plannode.PlanNode;
import org.apache.pinot.query.planner.plannode.PlanNodeVisitor;
import org.apache.pinot.query.planner.plannode.TableScanNode;
+import org.apache.pinot.query.planner.plannode.WindowNode;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
@@ -154,6 +156,67 @@ public class StagesTestBase {
List.of());
}
+ /// Creates an ASOF join node whose whole comparison lives in the match
condition, as ASOF joins require.
+ public SimpleChildBuilder<JoinNode> asofJoin(
+ SimpleChildBuilder<? extends PlanNode> leftBuilder,
+ SimpleChildBuilder<? extends PlanNode> rightBuilder,
+ RexExpression matchCondition) {
+ return (stageId, mySchema, myHints) -> {
+ PlanNode left = leftBuilder.build(stageId);
+ PlanNode right = rightBuilder.build(stageId);
+ return new JoinNode(stageId, mySchema, myHints, List.of(left, right),
JoinRelType.LEFT, List.of(0), List.of(0),
+ List.of(), JoinNode.JoinStrategy.ASOF, matchCondition);
+ };
+ }
+
+ /// Creates a `COUNT` aggregate node grouping on the first two columns with
the given grouping sets.
+ public SimpleChildBuilder<AggregateNode> aggregate(SimpleChildBuilder<?
extends PlanNode> childBuilder,
+ List<List<Integer>> groupingSets) {
+ return (stageId, mySchema, myHints) -> {
+ PlanNode input = childBuilder.build(stageId);
+ DataSchema schema = mySchema != null ? mySchema : input.getDataSchema();
+ RexExpression.FunctionCall count =
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.LONG,
"COUNT", List.of());
+ return new AggregateNode(stageId, schema, myHints, List.of(input),
List.of(count), List.of(), List.of(0, 1),
+ AggregateNode.AggType.DIRECT, false, List.of(), -1, groupingSets);
+ };
+ }
+
+ /// Creates a window node over the given child with a single aggregate call
and no frame exclusion.
+ public SimpleChildBuilder<WindowNode> window(SimpleChildBuilder<? extends
PlanNode> childBuilder,
+ RexExpression.FunctionCall aggCall) {
+ return window(childBuilder, aggCall, WindowNode.WindowExclusion.NO_OTHERS);
+ }
+
+ /// Creates a window node over the given child with a single aggregate call.
+ ///
+ /// The frame is `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`
partitioned on the first column with no
+ /// collations. Only the fields that tests need to vary are exposed as
parameters; add more if a test needs them.
+ public SimpleChildBuilder<WindowNode> window(SimpleChildBuilder<? extends
PlanNode> childBuilder,
+ RexExpression.FunctionCall aggCall, WindowNode.WindowExclusion exclude) {
+ return (stageId, mySchema, myHints) -> {
+ PlanNode input = childBuilder.build(stageId);
+ DataSchema schema = mySchema != null ? mySchema : input.getDataSchema();
+ return new WindowNode(stageId, schema, myHints, List.of(input),
List.of(0), List.of(), List.of(aggCall),
+ WindowNode.WindowFrameType.ROWS, Integer.MIN_VALUE, 0, exclude,
List.of());
+ };
+ }
+
+ /// Creates a `LAST_VALUE` window aggregate call over the first input column
that respects nulls.
+ public static RexExpression.FunctionCall lastValueRespectingNulls() {
+ return lastValue(false);
+ }
+
+ /// Creates a `LAST_VALUE` window aggregate call over the first input column
that ignores nulls.
+ public static RexExpression.FunctionCall lastValueIgnoringNulls() {
+ return lastValue(true);
+ }
+
+ private static RexExpression.FunctionCall lastValue(boolean ignoreNulls) {
+ return new RexExpression.FunctionCall(DataSchema.ColumnDataType.INT,
"LAST_VALUE",
+ List.of(new RexExpression.InputRef(0)), false, ignoreNulls);
+ }
+
/// Looks for the mailbox that corresponds to the given stageId.
/// @throws IllegalStateException if the mailbox is not found.
public MailboxSendNode stage(int stageId) {
diff --git a/pinot-query-runtime/src/test/resources/queries/Spool.json
b/pinot-query-runtime/src/test/resources/queries/Spool.json
index 4e335d87e5a..dc67255ea19 100644
--- a/pinot-query-runtime/src/test/resources/queries/Spool.json
+++ b/pinot-query-runtime/src/test/resources/queries/Spool.json
@@ -24,6 +24,17 @@
["foo", "bob", 3, 3.1416, true],
["alice", "alice", 4, 2.7183, false]
]
+ },
+ "tbl3" : {
+ "schema": [
+ {"name": "grpCol", "type": "STRING"},
+ {"name": "ordCol", "type": "INT"}
+ ],
+ "inputs": [
+ ["foo", 1],
+ ["foo", 2],
+ ["foo", 3]
+ ]
}
},
"queries": [
@@ -46,6 +57,11 @@
"description": "Spool of an aggregated stage reused by both arms of a
UNION ALL",
"sql": "SET timeoutMs=10000; SET useSpools=true; WITH g AS (SELECT
strCol1, SUM(intCol1) AS s FROM {tbl1} GROUP BY strCol1) SELECT strCol1, s FROM
g UNION ALL SELECT strCol1, s FROM g",
"h2Sql": "WITH g AS (SELECT strCol1, SUM(intCol1) AS s FROM {tbl1}
GROUP BY strCol1) SELECT strCol1, s FROM g UNION ALL SELECT strCol1, s FROM g"
+ },
+ {
+ "description": "Two window stages that only differ on the frame
EXCLUDE must not be spooled together. Both windows must partition on the same
key so that the exclusion is the only thing telling the two stages apart.",
+ "sql": "SET timeoutMs=10000; SET useSpools=true; SELECT w1.ordCol,
w1.s, w2.ordCol, w2.s FROM (SELECT grpCol, ordCol, SUM(ordCol) OVER (PARTITION
BY grpCol ORDER BY ordCol ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
EXCLUDE CURRENT ROW) AS s FROM {tbl3}) w1 JOIN {tbl1} b ON w1.grpCol =
b.strCol1 JOIN (SELECT grpCol, ordCol, SUM(ordCol) OVER (PARTITION BY grpCol
ORDER BY ordCol ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW EXCLUDE NO
OTHERS) AS s FROM {tbl3}) w2 ON w2.grpC [...]
+ "h2Sql": "SELECT w1.ordCol, w1.s, w2.ordCol, w2.s FROM (SELECT grpCol,
ordCol, SUM(ordCol) OVER (PARTITION BY grpCol ORDER BY ordCol ROWS BETWEEN
UNBOUNDED PRECEDING AND CURRENT ROW EXCLUDE CURRENT ROW) AS s FROM {tbl3}) w1
JOIN {tbl1} b ON w1.grpCol = b.strCol1 JOIN (SELECT grpCol, ordCol, SUM(ordCol)
OVER (PARTITION BY grpCol ORDER BY ordCol ROWS BETWEEN UNBOUNDED PRECEDING AND
CURRENT ROW EXCLUDE NO OTHERS) AS s FROM {tbl3}) w2 ON w2.grpCol = b.strCol1"
}
]
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
