You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2025/08/14 16:51:54 UTC
(pinot) branch master updated: [MSE] Enriched join operator (#16123)
This is an automated email from the ASF dual-hosted git repository.
gortiz 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 1cc901f49b0 [MSE] Enriched join operator (#16123)
1cc901f49b0 is described below
commit 1cc901f49b018b3aa29c6c48bd6fb4c5c2f13d0e
Author: Song Fu <13...@users.noreply.github.com>
AuthorDate: Thu Aug 14 09:51:48 2025 -0700
[MSE] Enriched join operator (#16123)
---
pinot-common/src/main/proto/plan.proto | 30 +
.../rel/logical/PinotLogicalEnrichedJoin.java | 248 ++++++++
.../PinotAggregateExchangeNodeInsertRule.java | 3 +
.../calcite/rel/rules/PinotEnrichedJoinRule.java | 396 +++++++++++++
.../org/apache/pinot/query/QueryEnvironment.java | 12 +-
.../planner/explain/ExplainNodeSimplifier.java | 6 +
.../explain/PhysicalExplainPlanVisitor.java | 9 +
.../query/planner/explain/PlanNodeMerger.java | 33 ++
.../query/planner/explain/PlanNodeSorter.java | 6 +
.../planner/logical/EquivalentStagesFinder.java | 16 +
.../query/planner/logical/PlanFragmenter.java | 6 +
.../planner/logical/PlanNodeToRelConverter.java | 38 ++
.../planner/logical/RelToPlanNodeConverter.java | 89 +++
.../query/planner/logical/SubPlanFragmenter.java | 6 +
.../planner/physical/DispatchablePlanVisitor.java | 7 +
.../planner/physical/v2/RelToPRelConverter.java | 4 +
.../planner/physical/v2/nodes/PhysicalJoin.java | 26 +-
.../plannode/DefaultPostOrderTraversalVisitor.java | 5 +
.../query/planner/plannode/EnrichedJoinNode.java | 160 +++++
.../query/planner/plannode/PlanNodeVisitor.java | 9 +
.../query/planner/serde/PlanNodeDeserializer.java | 42 ++
.../query/planner/serde/PlanNodeSerializer.java | 39 ++
.../validation/ArrayToMvValidationVisitor.java | 7 +
.../rel/rules/PinotEnrichedJoinRuleTest.java | 638 ++++++++++++++++++++
.../logical/RelToPlanNodeConverterTest.java | 73 +++
.../src/test/resources/queries/JoinPlans.json | 97 +++
.../query/runtime/InStageStatsTreeBuilder.java | 6 +
.../query/runtime/operator/BaseJoinOperator.java | 31 +-
.../runtime/operator/EnrichedHashJoinOperator.java | 417 +++++++++++++
.../query/runtime/operator/HashJoinOperator.java | 22 +-
.../query/runtime/plan/PlanNodeToOpChain.java | 35 ++
.../plan/server/ServerPlanRequestVisitor.java | 66 +++
.../operator/EnrichedHashJoinOperatorTest.java | 652 +++++++++++++++++++++
.../src/test/resources/queries/EnrichedJoin.json | 93 +++
.../apache/pinot/spi/utils/CommonConstants.java | 4 +-
35 files changed, 3314 insertions(+), 17 deletions(-)
diff --git a/pinot-common/src/main/proto/plan.proto b/pinot-common/src/main/proto/plan.proto
index 703abb96075..78f283d4141 100644
--- a/pinot-common/src/main/proto/plan.proto
+++ b/pinot-common/src/main/proto/plan.proto
@@ -40,6 +40,7 @@ message PlanNode {
ValueNode valueNode = 14;
WindowNode windowNode = 15;
ExplainNode explainNode = 16;
+ EnrichedJoinNode enrichedJoinNode = 17;
}
}
@@ -103,6 +104,35 @@ message JoinNode {
Expression matchCondition = 6;
}
+enum FilterProjectRexType {
+ FILTER = 0;
+ PROJECT = 1;
+}
+
+message ProjectAndResultSchema {
+ repeated Expression project = 1;
+ DataSchema schema = 2;
+}
+
+message FilterProjectRex {
+ FilterProjectRexType type = 1;
+ optional Expression filter = 2;
+ optional ProjectAndResultSchema projectAndResultSchema = 3;
+}
+
+message EnrichedJoinNode {
+ JoinType joinType = 1;
+ repeated int32 leftKeys = 2;
+ repeated int32 rightKeys = 3;
+ repeated Expression nonEquiConditions = 4;
+ JoinStrategy joinStrategy = 5;
+ Expression matchCondition = 6;
+ repeated FilterProjectRex filterProjectRex = 7;
+ DataSchema joinResultDataSchema = 8;
+ int32 fetch = 9;
+ int32 offset = 10;
+}
+
enum ExchangeType {
STREAMING = 0;
SUB_PLAN = 1;
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalEnrichedJoin.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalEnrichedJoin.java
new file mode 100644
index 00000000000..03244709934
--- /dev/null
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/logical/PinotLogicalEnrichedJoin.java
@@ -0,0 +1,248 @@
+/**
+ * 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.pinot.calcite.rel.logical;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import javax.annotation.Nullable;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelWriter;
+import org.apache.calcite.rel.core.CorrelationId;
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.hint.RelHint;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.sql.validate.SqlValidatorUtil;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+
+public class PinotLogicalEnrichedJoin extends Join {
+
+ private final RelDataType _outputRowType;
+ private final List<FilterProjectRexNode> _filterProjectRexNodes;
+ /**
+ * Currently variableSet of Project Rel is ignored since
+ * we don't support nested expressions in execution
+ */
+ private final Set<CorrelationId> _projectVariableSet;
+ @Nullable
+ private final List<RexNode> _squashedProjects;
+ @Nullable
+ private final RexNode _fetch;
+ @Nullable
+ private final RexNode _offset;
+
+ public PinotLogicalEnrichedJoin(RelOptCluster cluster, RelTraitSet traitSet,
+ List<RelHint> hints, RelNode left, RelNode right, RexNode joinCondition,
+ Set<CorrelationId> variablesSet, JoinRelType joinType,
+ List<FilterProjectRexNode> filterProjectRexNodes,
+ @Nullable RelDataType outputRowType, @Nullable Set<CorrelationId> projectVariableSet,
+ @Nullable RexNode fetch, @Nullable RexNode offset) {
+ super(cluster, traitSet, hints, left, right, joinCondition, variablesSet, joinType);
+ _filterProjectRexNodes = filterProjectRexNodes;
+ _squashedProjects = squashProjects();
+ RelDataType joinRowType = getJoinRowType();
+ // if there's projection, getRowType() should return the final projected row type as output row type
+ // otherwise it's the same as _joinRowType
+ _outputRowType = outputRowType == null ? joinRowType : outputRowType;
+ _projectVariableSet = projectVariableSet;
+ _offset = offset;
+ _fetch = fetch;
+ }
+
+ @Override
+ public PinotLogicalEnrichedJoin copy(RelTraitSet traitSet, RexNode conditionExpr,
+ RelNode left, RelNode right, JoinRelType joinType, boolean semiJoinDone) {
+ return new PinotLogicalEnrichedJoin(getCluster(), traitSet, getHints(), left, right,
+ conditionExpr, getVariablesSet(), getJoinType(),
+ _filterProjectRexNodes, _outputRowType, _projectVariableSet,
+ _fetch, _offset);
+ }
+
+ public PinotLogicalEnrichedJoin withNewProject(FilterProjectRexNode project, RelDataType outputRowType,
+ Set<CorrelationId> projectVariableSet) {
+ List<FilterProjectRexNode> filterProjectRexNodes = new ArrayList<>(_filterProjectRexNodes.size() + 1);
+ filterProjectRexNodes.addAll(_filterProjectRexNodes);
+ filterProjectRexNodes.add(project);
+ return new PinotLogicalEnrichedJoin(getCluster(), getTraitSet(), getHints(), left, right,
+ getCondition(), getVariablesSet(), getJoinType(),
+ filterProjectRexNodes, outputRowType, projectVariableSet,
+ _fetch, _offset);
+ }
+
+ public PinotLogicalEnrichedJoin withNewFilter(FilterProjectRexNode filter) {
+ List<FilterProjectRexNode> filterProjectRexNodes = new ArrayList<>(_filterProjectRexNodes.size() + 1);
+ filterProjectRexNodes.addAll(_filterProjectRexNodes);
+ filterProjectRexNodes.add(filter);
+ return new PinotLogicalEnrichedJoin(getCluster(), getTraitSet(), getHints(), left, right,
+ getCondition(), getVariablesSet(), getJoinType(),
+ filterProjectRexNodes, _outputRowType, _projectVariableSet,
+ _fetch, _offset);
+ }
+
+ public PinotLogicalEnrichedJoin withNewFetchOffset(@Nullable RexNode fetch, @Nullable RexNode offset) {
+ return new PinotLogicalEnrichedJoin(getCluster(), getTraitSet(), getHints(), left, right,
+ getCondition(), getVariablesSet(), getJoinType(),
+ _filterProjectRexNodes, _outputRowType, _projectVariableSet,
+ fetch, offset);
+ }
+
+ @Override
+ protected RelDataType deriveRowType() {
+ return checkNotNull(_outputRowType);
+ }
+
+ public final RelDataType getJoinRowType() {
+ return SqlValidatorUtil.deriveJoinRowType(left.getRowType(),
+ right.getRowType(), joinType, getCluster().getTypeFactory(), null,
+ getSystemFieldList());
+ }
+
+ public List<FilterProjectRexNode> getFilterProjectRexNodes() {
+ return _filterProjectRexNodes;
+ }
+
+ public List<RexNode> getProjects() {
+ return _squashedProjects == null ? Collections.emptyList() : _squashedProjects;
+ }
+
+ /// Combine all projects in _filterProjectRexNodes into a single project
+ @Nullable
+ private List<RexNode> squashProjects() {
+ List<RexNode> prevProject = null;
+ for (FilterProjectRexNode node : _filterProjectRexNodes) {
+ if (node.getType() == FilterProjectRexNodeType.FILTER) {
+ continue;
+ }
+ assert node.getProjectAndResultRowType() != null;
+ List<RexNode> project = node.getProjectAndResultRowType().getProject();
+ if (prevProject == null) {
+ prevProject = project;
+ continue;
+ }
+ // combine project
+ prevProject = combineProjects(project, prevProject);
+ }
+ return prevProject;
+ }
+
+ /// Adopted from {@link org.apache.calcite.plan.RelOptUtil#pushPastProject}
+ private static List<RexNode> combineProjects(List<RexNode> upper, List<RexNode> lower) {
+ return new RexShuttle() {
+ @Override
+ public RexNode visitInputRef(RexInputRef ref) {
+ return lower.get(ref.getIndex());
+ }
+ }.visitList(upper);
+ }
+
+ @Override
+ public RelWriter explainTerms(RelWriter pw) {
+ return super.explainTerms(pw)
+ .item("filterProjectRex", _filterProjectRexNodes)
+ .itemIf("limit", _fetch, _fetch != null)
+ .itemIf("offset", _offset, _offset != null);
+ }
+
+ public enum FilterProjectRexNodeType {
+ FILTER,
+ PROJECT
+ }
+
+ public static class ProjectAndResultRowType {
+ private final List<RexNode> _project;
+ private final RelDataType _dataType;
+
+ public ProjectAndResultRowType(List<RexNode> project, RelDataType resultRowType) {
+ _project = project;
+ _dataType = resultRowType;
+ }
+
+ public List<RexNode> getProject() {
+ return _project;
+ }
+
+ public RelDataType getDataType() {
+ return _dataType;
+ }
+ }
+
+ public static class FilterProjectRexNode {
+ private final FilterProjectRexNodeType _type;
+ @Nullable
+ private final RexNode _filter;
+ @Nullable
+ private final ProjectAndResultRowType _projectAndResultRowType;
+
+ @Override
+ public String toString() {
+ if (_type == FilterProjectRexNodeType.FILTER) {
+ assert _filter != null;
+ return "Filter: " + _filter;
+ } else {
+ assert _projectAndResultRowType != null;
+ return "Project: " + _projectAndResultRowType.getProject().toString();
+ }
+ }
+
+ public FilterProjectRexNode(RexNode filter) {
+ _type = FilterProjectRexNodeType.FILTER;
+ _filter = filter;
+ _projectAndResultRowType = null;
+ }
+
+ public FilterProjectRexNode(List<RexNode> project, RelDataType resultDataType) {
+ _type = FilterProjectRexNodeType.PROJECT;
+ _filter = null;
+ _projectAndResultRowType = new ProjectAndResultRowType(project, resultDataType);
+ }
+
+ public FilterProjectRexNodeType getType() {
+ return _type;
+ }
+
+ @Nullable
+ public RexNode getFilter() {
+ return _filter;
+ }
+
+ @Nullable
+ public ProjectAndResultRowType getProjectAndResultRowType() {
+ return _projectAndResultRowType;
+ }
+ }
+
+ @Nullable
+ public RexNode getFetch() {
+ return _fetch;
+ }
+
+ @Nullable
+ public RexNode getOffset() {
+ return _offset;
+ }
+}
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
index a367c95e569..639f825993b 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotAggregateExchangeNodeInsertRule.java
@@ -61,6 +61,7 @@ import org.apache.calcite.util.mapping.Mappings;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.calcite.rel.hint.PinotHintStrategyTable;
import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange;
import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange;
import org.apache.pinot.common.function.sql.PinotSqlAggFunction;
@@ -472,6 +473,8 @@ public class PinotAggregateExchangeNodeInsertRule {
return ((Project) relNode).getProjects();
} else if (relNode instanceof Union) {
return findImmediateProjects(relNode.getInput(0));
+ } else if (relNode instanceof PinotLogicalEnrichedJoin) {
+ return ((PinotLogicalEnrichedJoin) relNode).getProjects();
}
return null;
}
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRule.java
new file mode 100644
index 00000000000..5a140e223ee
--- /dev/null
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRule.java
@@ -0,0 +1,396 @@
+/**
+ * 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.pinot.calcite.rel.rules;
+
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.immutables.value.Value;
+
+
+/**
+ * Rules for fusing filter, projection, and limit operators into join and enriched joins.
+ * This collection of rules will be fired bottom-up and fuse operators into the enriched join greedily.
+ */
+@Value.Enclosing
+public class PinotEnrichedJoinRule {
+
+ private PinotEnrichedJoinRule() {
+ }
+
+ /**
+ * Rule that fuses a LogicalProject into a PinotLogicalEnrichedJoin
+ */
+ public static class ProjectEnrichedJoin extends RelRule<ProjectEnrichedJoin.ProjectEnrichedJoinConfig> {
+
+ @Value.Immutable
+ public interface ProjectEnrichedJoinConfig extends RelRule.Config {
+ ProjectEnrichedJoinConfig DEFAULT =
+ ImmutablePinotEnrichedJoinRule.ProjectEnrichedJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalProject.class).oneInput(b1 ->
+ b1.operand(PinotLogicalEnrichedJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default ProjectEnrichedJoin toRule() {
+ return new ProjectEnrichedJoin(this);
+ }
+ }
+
+ private ProjectEnrichedJoin(ProjectEnrichedJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalProject project = call.rel(0);
+ PinotLogicalEnrichedJoin enrichedJoin = call.rel(1);
+
+ // Add projection to the enriched join
+ PinotLogicalEnrichedJoin newEnrichedJoin = enrichedJoin.withNewProject(
+ new PinotLogicalEnrichedJoin.FilterProjectRexNode(project.getProjects(), project.getRowType()),
+ project.getRowType(),
+ project.getVariablesSet());
+
+ call.transformTo(newEnrichedJoin);
+ }
+ }
+
+ /**
+ * Rule that fuses a LogicalFilter into a PinotLogicalEnrichedJoin
+ */
+ public static class FilterEnrichedJoin extends RelRule<FilterEnrichedJoin.FilterEnrichedJoinConfig> {
+
+ @Value.Immutable
+ public interface FilterEnrichedJoinConfig extends RelRule.Config {
+ FilterEnrichedJoinConfig DEFAULT =
+ ImmutablePinotEnrichedJoinRule.FilterEnrichedJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalFilter.class).oneInput(b1 ->
+ b1.operand(PinotLogicalEnrichedJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default FilterEnrichedJoin toRule() {
+ return new FilterEnrichedJoin(this);
+ }
+ }
+
+ private FilterEnrichedJoin(FilterEnrichedJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalFilter filter = call.rel(0);
+ PinotLogicalEnrichedJoin enrichedJoin = call.rel(1);
+
+ // Add filter to the enriched join
+ PinotLogicalEnrichedJoin newEnrichedJoin = enrichedJoin.withNewFilter(
+ new PinotLogicalEnrichedJoin.FilterProjectRexNode(filter.getCondition()));
+
+ call.transformTo(newEnrichedJoin);
+ }
+ }
+
+ /**
+ * Rule that fuses a LogicalSort into a PinotLogicalEnrichedJoin
+ */
+ public static class SortEnrichedJoin extends RelRule<SortEnrichedJoin.SortEnrichedJoinConfig> {
+
+ @Value.Immutable
+ public interface SortEnrichedJoinConfig extends RelRule.Config {
+ SortEnrichedJoinConfig DEFAULT = ImmutablePinotEnrichedJoinRule.SortEnrichedJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalSort.class).oneInput(b1 ->
+ b1.operand(PinotLogicalEnrichedJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default SortEnrichedJoin toRule() {
+ return new SortEnrichedJoin(this);
+ }
+ }
+
+ private SortEnrichedJoin(SortEnrichedJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalSort sort = call.rel(0);
+ PinotLogicalEnrichedJoin enrichedJoin = call.rel(1);
+
+ // If enriched join already had a sort operator merged, return
+ if (enrichedJoin.getOffset() != null || enrichedJoin.getFetch() != null) {
+ return;
+ }
+
+ // Enriched join does not support sort collation, only fetch and offset
+ if (sort.getCollation() != null && !sort.getCollation().equals(RelCollations.EMPTY)) {
+ return;
+ }
+
+ // Add sort limit to the enriched join
+ PinotLogicalEnrichedJoin newEnrichedJoin = enrichedJoin.withNewFetchOffset(sort.fetch, sort.offset);
+ call.transformTo(newEnrichedJoin);
+ }
+ }
+
+ /**
+ * Rule that converts LogicalProject + LogicalJoin into PinotLogicalEnrichedJoin
+ */
+ public static class ProjectJoin extends RelRule<ProjectJoin.ProjectJoinConfig> {
+
+ @Value.Immutable
+ public interface ProjectJoinConfig extends RelRule.Config {
+ ProjectJoinConfig DEFAULT = ImmutablePinotEnrichedJoinRule.ProjectJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalProject.class).oneInput(b1 ->
+ b1.operand(LogicalJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default ProjectJoin toRule() {
+ return new ProjectJoin(this);
+ }
+ }
+
+ private ProjectJoin(ProjectJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalProject project = call.rel(0);
+ LogicalJoin join = call.rel(1);
+
+ if (!canConvertJoin(join)) {
+ return;
+ }
+
+ List<PinotLogicalEnrichedJoin.FilterProjectRexNode> filterProjectRexNodes = new ArrayList<>();
+ filterProjectRexNodes.add(
+ new PinotLogicalEnrichedJoin.FilterProjectRexNode(project.getProjects(), project.getRowType()));
+
+ PinotLogicalEnrichedJoin enrichedJoin = new PinotLogicalEnrichedJoin(
+ join.getCluster(),
+ join.getTraitSet(),
+ join.getHints(),
+ join.getInput(0),
+ join.getInput(1),
+ join.getCondition(),
+ join.getVariablesSet(),
+ join.getJoinType(),
+ filterProjectRexNodes,
+ project.getRowType(),
+ project.getVariablesSet(),
+ null,
+ null);
+
+ call.transformTo(enrichedJoin);
+ }
+ }
+
+ /**
+ * Rule that converts LogicalFilter + LogicalJoin into PinotLogicalEnrichedJoin
+ */
+ public static class FilterJoin extends RelRule<FilterJoin.FilterJoinConfig> {
+
+ @Value.Immutable
+ public interface FilterJoinConfig extends RelRule.Config {
+ FilterJoinConfig DEFAULT = ImmutablePinotEnrichedJoinRule.FilterJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalFilter.class).oneInput(b1 ->
+ b1.operand(LogicalJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default FilterJoin toRule() {
+ return new FilterJoin(this);
+ }
+ }
+
+ private FilterJoin(FilterJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalFilter filter = call.rel(0);
+ LogicalJoin join = call.rel(1);
+
+ if (!canConvertJoin(join)) {
+ return;
+ }
+
+ List<PinotLogicalEnrichedJoin.FilterProjectRexNode> filterProjectRexNodes = new ArrayList<>();
+ filterProjectRexNodes.add(new PinotLogicalEnrichedJoin.FilterProjectRexNode(filter.getCondition()));
+
+ PinotLogicalEnrichedJoin enrichedJoin = new PinotLogicalEnrichedJoin(
+ join.getCluster(),
+ join.getTraitSet(),
+ join.getHints(),
+ join.getInput(0),
+ join.getInput(1),
+ join.getCondition(),
+ join.getVariablesSet(),
+ join.getJoinType(),
+ filterProjectRexNodes,
+ null, // No projection, final output row type is same as join output
+ null, // No projection variable set
+ null,
+ null);
+
+ call.transformTo(enrichedJoin);
+ }
+ }
+
+ /**
+ * Rule that converts LogicalSort + LogicalJoin into PinotLogicalEnrichedJoin
+ */
+ public static class SortJoin extends RelRule<SortJoin.SortJoinConfig> {
+
+ @Value.Immutable
+ public interface SortJoinConfig extends RelRule.Config {
+ SortJoinConfig DEFAULT = ImmutablePinotEnrichedJoinRule.SortJoinConfig.builder()
+ .operandSupplier(b0 ->
+ b0.operand(LogicalSort.class).oneInput(b1 ->
+ b1.operand(LogicalJoin.class).anyInputs()
+ )
+ ).build();
+
+ @Override
+ default SortJoin toRule() {
+ return new SortJoin(this);
+ }
+ }
+
+ private SortJoin(SortJoinConfig config) {
+ super(config);
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ LogicalSort sort = call.rel(0);
+ LogicalJoin join = call.rel(1);
+
+ // Enriched join does not support sort collation, only fetch and offset
+ if (sort.getCollation() != null && !sort.getCollation().equals(RelCollations.EMPTY)) {
+ return;
+ }
+
+ if (!canConvertJoin(join)) {
+ return;
+ }
+
+ List<PinotLogicalEnrichedJoin.FilterProjectRexNode> filterProjectRexNodes = new ArrayList<>();
+
+ PinotLogicalEnrichedJoin enrichedJoin = new PinotLogicalEnrichedJoin(
+ join.getCluster(),
+ join.getTraitSet(),
+ join.getHints(),
+ join.getInput(0),
+ join.getInput(1),
+ join.getCondition(),
+ join.getVariablesSet(),
+ join.getJoinType(),
+ filterProjectRexNodes,
+ null, // No projection, final output row type is same as join output
+ null, // No projection variable set
+ sort.fetch,
+ sort.offset);
+
+ call.transformTo(enrichedJoin);
+ }
+ }
+
+ /**
+ * Temporarily disable conversion of lookup join and non-equijoin to EnrichedJoin
+ */
+ private static boolean canConvertJoin(LogicalJoin join) {
+ // Disable lookup join for now
+ if (PinotHintOptions.JoinHintOptions.useLookupJoinStrategy(join)) {
+ return false;
+ }
+
+ // Disable non-equijoin for now
+ if (join.analyzeCondition().leftKeys.isEmpty()) {
+ return false;
+ }
+
+ return true;
+ }
+
+ // Rule instances with descriptions
+ public static final ProjectEnrichedJoin PROJECT_ENRICHED_JOIN =
+ (ProjectEnrichedJoin) ProjectEnrichedJoin.ProjectEnrichedJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final FilterEnrichedJoin FILTER_ENRICHED_JOIN =
+ (FilterEnrichedJoin) FilterEnrichedJoin.FilterEnrichedJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final SortEnrichedJoin SORT_ENRICHED_JOIN =
+ (SortEnrichedJoin) SortEnrichedJoin.SortEnrichedJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final ProjectJoin PROJECT_JOIN =
+ (ProjectJoin) ProjectJoin.ProjectJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final FilterJoin FILTER_JOIN =
+ (FilterJoin) FilterJoin.FilterJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final SortJoin SORT_JOIN =
+ (SortJoin) SortJoin.SortJoinConfig.DEFAULT
+ .withDescription(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN)
+ .toRule();
+
+ public static final List<RelOptRule> PINOT_ENRICHED_JOIN_RULES = List.of(
+ PROJECT_ENRICHED_JOIN,
+ FILTER_ENRICHED_JOIN,
+ SORT_ENRICHED_JOIN,
+ PROJECT_JOIN,
+ FILTER_JOIN,
+ SORT_JOIN
+ );
+}
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
index 187125bc236..225e1dc4f6a 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java
@@ -56,6 +56,7 @@ import org.apache.calcite.tools.Frameworks;
import org.apache.calcite.tools.RelBuilder;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pinot.calcite.rel.rules.PinotEnrichedJoinRule;
import org.apache.pinot.calcite.rel.rules.PinotImplicitTableHintRule;
import org.apache.pinot.calcite.rel.rules.PinotJoinToDynamicBroadcastRule;
import org.apache.pinot.calcite.rel.rules.PinotQueryRuleSets;
@@ -170,9 +171,9 @@ public class QueryEnvironment {
WorkerManager workerManager = getWorkerManager(sqlNodeAndOptions);
Map<String, String> options = sqlNodeAndOptions.getOptions();
HepProgram optProgram = _optProgram;
+ Set<String> useRuleSet = QueryOptionsUtils.getUsePlannerRules(options);
if (MapUtils.isNotEmpty(options)) {
Set<String> skipRuleSet = QueryOptionsUtils.getSkipPlannerRules(options);
- Set<String> useRuleSet = QueryOptionsUtils.getUsePlannerRules(options);
if (!skipRuleSet.isEmpty() || !useRuleSet.isEmpty()) {
// dynamically create optProgram according to rule options
optProgram = getOptProgram(skipRuleSet, useRuleSet);
@@ -180,7 +181,7 @@ public class QueryEnvironment {
}
boolean usePhysicalOptimizer = QueryOptionsUtils.isUsePhysicalOptimizer(sqlNodeAndOptions.getOptions(),
_envConfig.defaultUsePhysicalOptimizer());
- HepProgram traitProgram = getTraitProgram(workerManager, _envConfig, usePhysicalOptimizer);
+ HepProgram traitProgram = getTraitProgram(workerManager, _envConfig, usePhysicalOptimizer, useRuleSet);
SqlExplainFormat format = SqlExplainFormat.DOT;
if (sqlNodeAndOptions.getSqlNode().getKind().equals(SqlKind.EXPLAIN)) {
SqlExplain explain = (SqlExplain) sqlNodeAndOptions.getSqlNode();
@@ -533,6 +534,7 @@ public class QueryEnvironment {
// Prune duplicate/unnecessary nodes using a single HepInstruction.
// TODO: We can consider using HepMatchOrder.TOP_DOWN if we find cases where it would help.
hepProgramBuilder.addRuleCollection(pruneRules);
+
return hepProgramBuilder.build();
}
@@ -577,7 +579,7 @@ public class QueryEnvironment {
}
private static HepProgram getTraitProgram(@Nullable WorkerManager workerManager, Config config,
- boolean usePhysicalOptimizer) {
+ boolean usePhysicalOptimizer, Set<String> useRuleSet) {
HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
// Set the match order as BOTTOM_UP.
@@ -591,6 +593,10 @@ public class QueryEnvironment {
hepProgramBuilder.addRuleInstance(relOptRule);
}
}
+ if (!isRuleSkipped(CommonConstants.Broker.PlannerRuleNames.JOIN_TO_ENRICHED_JOIN, Set.of(), useRuleSet)) {
+ // push filter and project above join to enrichedJoin, does not work with physical optimizer
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ }
} else {
for (RelOptRule relOptRule : PinotQueryRuleSets.PINOT_POST_RULES_V2) {
if (isEligibleQueryPostRule(relOptRule, config)) {
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/ExplainNodeSimplifier.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/ExplainNodeSimplifier.java
index 17b0b02fe1e..b8fcaaab442 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/ExplainNodeSimplifier.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/ExplainNodeSimplifier.java
@@ -24,6 +24,7 @@ import java.util.Collections;
import java.util.List;
import org.apache.pinot.core.query.reduce.ExplainPlanDataTableReducer;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -89,6 +90,11 @@ public class ExplainNodeSimplifier {
return defaultNode(node);
}
+ @Override
+ public PlanNode visitEnrichedJoin(EnrichedJoinNode node, Void context) {
+ return defaultNode(node);
+ }
+
@Override
public PlanNode visitMailboxReceive(MailboxReceiveNode node, Void context) {
return defaultNode(node);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PhysicalExplainPlanVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PhysicalExplainPlanVisitor.java
index 615ee3ba05b..c3b6066a5b7 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PhysicalExplainPlanVisitor.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PhysicalExplainPlanVisitor.java
@@ -29,6 +29,7 @@ import java.util.stream.Stream;
import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
import org.apache.pinot.query.planner.physical.DispatchableSubPlan;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -167,6 +168,14 @@ public class PhysicalExplainPlanVisitor implements PlanNodeVisitor<StringBuilder
return context._builder;
}
+ @Override
+ public StringBuilder visitEnrichedJoin(EnrichedJoinNode node, Context context) {
+ appendInfo(node, context).append('\n');
+ node.getInputs().get(0).visit(this, context.next(true, context._host, context._workerId));
+ node.getInputs().get(1).visit(this, context.next(false, context._host, context._workerId));
+ return context._builder;
+ }
+
@Override
public StringBuilder visitMailboxReceive(MailboxReceiveNode node, Context context) {
appendInfo(node, context).append('\n');
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java
index 6ae02da45fc..250cefea623 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeMerger.java
@@ -31,6 +31,7 @@ import org.apache.pinot.common.proto.Plan;
import org.apache.pinot.core.operator.ExplainAttributeBuilder;
import org.apache.pinot.core.query.reduce.ExplainPlanDataTableReducer;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -203,6 +204,38 @@ class PlanNodeMerger {
return node.withInputs(children);
}
+ @Nullable
+ @Override
+ public PlanNode visitEnrichedJoin(EnrichedJoinNode node, PlanNode context) {
+ if (context.getClass() != EnrichedJoinNode.class) {
+ return null;
+ }
+ EnrichedJoinNode otherNode = (EnrichedJoinNode) context;
+ if (!node.getJoinType().equals(otherNode.getJoinType())) {
+ return null;
+ }
+ if (!node.getLeftKeys().equals(otherNode.getLeftKeys())) {
+ return null;
+ }
+ if (!node.getRightKeys().equals(otherNode.getRightKeys())) {
+ return null;
+ }
+ if (!node.getNonEquiConditions().equals(otherNode.getNonEquiConditions())) {
+ return null;
+ }
+ if (!Objects.equals(node.getFilterProjectRexes(), otherNode.getFilterProjectRexes())) {
+ return null;
+ }
+ if (node.getFetch() != otherNode.getFetch() || node.getOffset() != otherNode.getOffset()) {
+ return null;
+ }
+ List<PlanNode> children = mergeChildren(node, context);
+ if (children == null) {
+ return null;
+ }
+ return node.withInputs(children);
+ }
+
@Nullable
@Override
public PlanNode visitMailboxReceive(MailboxReceiveNode node, PlanNode context) {
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeSorter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeSorter.java
index 95404c5b7d7..afcfc07abc1 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeSorter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/explain/PlanNodeSorter.java
@@ -25,6 +25,7 @@ import java.util.List;
import java.util.TreeSet;
import org.apache.pinot.common.proto.Plan;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -91,6 +92,11 @@ public class PlanNodeSorter {
return defaultNode(node, comparator);
}
+ @Override
+ public PlanNode visitEnrichedJoin(EnrichedJoinNode node, Comparator<PlanNode> comparator) {
+ return visitJoin(node, comparator);
+ }
+
@Override
public PlanNode visitMailboxReceive(MailboxReceiveNode node, Comparator<PlanNode> comparator) {
return defaultNode(node, comparator);
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 33e10cd22b0..87e63c29b03 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
@@ -22,6 +22,7 @@ import com.google.common.base.Preconditions;
import java.util.List;
import java.util.Objects;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -258,6 +259,21 @@ public class EquivalentStagesFinder {
&& node1.getJoinStrategy() == that.getJoinStrategy();
}
+ @Override
+ public Boolean visitEnrichedJoin(EnrichedJoinNode node1, PlanNode node2) {
+ if (!(node2 instanceof EnrichedJoinNode)) {
+ return false;
+ }
+ if (!visitJoin(node1, node2)) {
+ return false;
+ }
+ EnrichedJoinNode that = (EnrichedJoinNode) node2;
+ return Objects.equals(node1.getFilterProjectRexes(), that.getFilterProjectRexes())
+ && node1.getJoinStrategy() == that.getJoinStrategy()
+ && node1.getFetch() == that.getFetch()
+ && node1.getOffset() == that.getOffset();
+ }
+
@Override
public Boolean visitProject(ProjectNode node1, PlanNode node2) {
if (!(node2 instanceof ProjectNode)) {
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
index 420cf505bf2..233714c3577 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanFragmenter.java
@@ -29,6 +29,7 @@ import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.query.planner.PlanFragment;
import org.apache.pinot.query.planner.SubPlan;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -126,6 +127,11 @@ public class PlanFragmenter implements PlanNodeVisitor<PlanNode, PlanFragmenter.
return process(node, context);
}
+ @Override
+ public PlanNode visitEnrichedJoin(EnrichedJoinNode node, Context context) {
+ return visitJoin(node, context);
+ }
+
@Override
public PlanNode visitMailboxReceive(MailboxReceiveNode node, Context context) {
throw new UnsupportedOperationException("MailboxReceiveNode should not be visited by PlanNodeFragmenter");
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
index 9ead0c3c621..7d3aa6aabb3 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/PlanNodeToRelConverter.java
@@ -21,7 +21,9 @@ package org.apache.pinot.query.planner.logical;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import java.util.stream.Collectors;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelTraitSet;
@@ -47,10 +49,12 @@ import org.apache.calcite.rex.RexWindowExclusion;
import org.apache.calcite.sql.SqlAggFunction;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.pinot.common.proto.Plan;
import org.apache.pinot.common.utils.DatabaseUtils;
import org.apache.pinot.core.operator.ExplainAttributeBuilder;
import org.apache.pinot.core.plan.PinotExplainedRelNode;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -170,6 +174,40 @@ public final class PlanNodeToRelConverter {
return null;
}
+ @Override
+ public Void visitEnrichedJoin(EnrichedJoinNode node, Void context) {
+ visitChildren(node);
+
+ try {
+ List<RexNode> conditions = new ArrayList<>(
+ node.getLeftKeys().size() + node.getRightKeys().size() + node.getNonEquiConditions().size());
+ for (Integer leftKey : node.getLeftKeys()) {
+ conditions.add(_builder.field(2, 0, leftKey));
+ }
+ for (Integer rightKey : node.getRightKeys()) {
+ conditions.add(_builder.field(2, 1, rightKey));
+ }
+ for (RexExpression nonEquiCondition : node.getNonEquiConditions()) {
+ conditions.add(RexExpressionUtils.toRexNode(_builder, nonEquiCondition));
+ }
+
+ if (node.getJoinType() == JoinRelType.ASOF || node.getJoinType() == JoinRelType.LEFT_ASOF) {
+ _builder.push(new PinotExplainedRelNode(_builder.getCluster(), "EnrichedASOFJoin", Collections.emptyMap(),
+ node.getDataSchema(), readAlreadyPushedChildren(node)));
+ } else {
+ Map<String, Plan.ExplainNode.AttributeValue> attributes = new HashMap<>();
+ _builder.push(new PinotExplainedRelNode(_builder.getCluster(), "EnrichedJoin", attributes,
+ node.getDataSchema(), readAlreadyPushedChildren(node)));
+ }
+ } catch (RuntimeException e) {
+ LOGGER.warn("Failed to convert join node: {}", node, e);
+ _builder.push(new PinotExplainedRelNode(_builder.getCluster(), "UnknownJoin", Collections.emptyMap(),
+ node.getDataSchema(), readAlreadyPushedChildren(node)));
+ }
+
+ return null;
+ }
+
@Override
public Void visitMailboxReceive(MailboxReceiveNode node, Void context) {
visitChildren(node);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
index abaf4704287..2212a8e1526 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java
@@ -55,6 +55,7 @@ import org.apache.calcite.sql.SqlLiteral;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.pinot.calcite.rel.hint.PinotHintOptions;
import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
import org.apache.pinot.calcite.rel.logical.PinotLogicalExchange;
import org.apache.pinot.calcite.rel.logical.PinotLogicalSortExchange;
import org.apache.pinot.calcite.rel.logical.PinotLogicalTableScan;
@@ -67,6 +68,7 @@ import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.DatabaseUtils;
import org.apache.pinot.common.utils.request.RequestUtils;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
import org.apache.pinot.query.planner.plannode.JoinNode;
@@ -78,6 +80,7 @@ import org.apache.pinot.query.planner.plannode.SortNode;
import org.apache.pinot.query.planner.plannode.TableScanNode;
import org.apache.pinot.query.planner.plannode.ValueNode;
import org.apache.pinot.query.planner.plannode.WindowNode;
+import org.codehaus.commons.nullanalysis.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -120,6 +123,8 @@ public final class RelToPlanNodeConverter {
result = convertLogicalSort((LogicalSort) node);
} else if (node instanceof Exchange) {
result = convertLogicalExchange((Exchange) node);
+ } else if (node instanceof PinotLogicalEnrichedJoin) {
+ result = convertLogicalEnrichedJoin((PinotLogicalEnrichedJoin) node);
} else if (node instanceof LogicalJoin) {
_brokerMetrics.addMeteredGlobalValue(BrokerMeter.JOIN_COUNT, 1);
if (!_joinFound) {
@@ -370,6 +375,90 @@ public final class RelToPlanNodeConverter {
joinStrategy);
}
+ private EnrichedJoinNode convertLogicalEnrichedJoin(PinotLogicalEnrichedJoin rel) {
+ JoinInfo joinInfo = rel.analyzeCondition();
+ DataSchema joinResultSchema = toDataSchema(rel.getJoinRowType());
+ DataSchema projectedSchema = toDataSchema(rel.getRowType());
+ List<PlanNode> inputs = convertInputs(rel.getInputs());
+ JoinRelType joinType = rel.getJoinType();
+
+ // Run some validations for join
+ Preconditions.checkState(inputs.size() == 2, "Join should have exactly 2 inputs, got: %s", inputs.size());
+ PlanNode left = inputs.get(0);
+ PlanNode right = inputs.get(1);
+ int numLeftColumns = left.getDataSchema().size();
+ int numJoinResultColumns = joinResultSchema.size();
+ if (joinType.projectsRight()) {
+ int numRightColumns = right.getDataSchema().size();
+ Preconditions.checkState(numLeftColumns + numRightColumns == numJoinResultColumns,
+ "Invalid number of columns for join type: %s, left: %s, right: %s, result: %s", joinType, numLeftColumns,
+ numRightColumns, numJoinResultColumns);
+ } else {
+ Preconditions.checkState(numLeftColumns == numJoinResultColumns,
+ "Invalid number of columns for join type: %s, left: %s, result: %s", joinType, numLeftColumns,
+ numJoinResultColumns);
+ }
+
+ // Check if the join hint specifies the join strategy
+ JoinNode.JoinStrategy joinStrategy;
+ if (PinotHintOptions.JoinHintOptions.useLookupJoinStrategy(rel)) {
+ joinStrategy = JoinNode.JoinStrategy.LOOKUP;
+
+ // Run some validations for lookup join
+ Preconditions.checkArgument(!joinInfo.leftKeys.isEmpty(), "Lookup join requires join keys");
+ // Right table should be a dimension table, and the right input should be an identifier only ProjectNode over
+ // TableScanNode.
+ RelNode rightInput = PinotRuleUtils.unboxRel(rel.getRight());
+ Preconditions.checkState(rightInput instanceof Project, "Right input for lookup join must be a Project, got: %s",
+ rightInput.getClass().getSimpleName());
+ Project project = (Project) rightInput;
+ for (RexNode node : project.getProjects()) {
+ Preconditions.checkState(node instanceof RexInputRef,
+ "Right input for lookup join must be an identifier (RexInputRef) only Project, got: %s in project",
+ node.getClass().getSimpleName());
+ }
+ RelNode projectInput = PinotRuleUtils.unboxRel(project.getInput());
+ Preconditions.checkState(projectInput instanceof TableScan,
+ "Right input for lookup join must be a Project over TableScan, got Project over: %s",
+ projectInput.getClass().getSimpleName());
+ } else {
+ // TODO: Consider adding DYNAMIC_BROADCAST as a separate join strategy
+ joinStrategy = JoinNode.JoinStrategy.HASH;
+ }
+
+ // convert filter and project RexNode into RexExpression
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = getFilterProjectRexes(rel);
+
+ int fetch = RexExpressionUtils.getValueAsInt(rel.getFetch());
+ int offset = RexExpressionUtils.getValueAsInt(rel.getOffset());
+
+ return new EnrichedJoinNode(DEFAULT_STAGE_ID, joinResultSchema, projectedSchema,
+ NodeHint.fromRelHints(rel.getHints()), inputs, joinType,
+ joinInfo.leftKeys, joinInfo.rightKeys, RexExpressionUtils.fromRexNodes(joinInfo.nonEquiConditions),
+ joinStrategy,
+ null,
+ filterProjectRexes,
+ fetch, offset);
+ }
+
+ @NotNull
+ private static List<EnrichedJoinNode.FilterProjectRex> getFilterProjectRexes(PinotLogicalEnrichedJoin rel) {
+ List<PinotLogicalEnrichedJoin.FilterProjectRexNode> filterProjectRexNode = rel.getFilterProjectRexNodes();
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = new ArrayList<>();
+ filterProjectRexNode.forEach((node) -> {
+ if (node.getType() == PinotLogicalEnrichedJoin.FilterProjectRexNodeType.FILTER) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(RexExpressionUtils.fromRexNode(node.getFilter())));
+ } else {
+ filterProjectRexes.add(
+ new EnrichedJoinNode.FilterProjectRex(
+ RexExpressionUtils.fromRexNodes(node.getProjectAndResultRowType().getProject()),
+ toDataSchema(node.getProjectAndResultRowType().getDataType())
+ ));
+ }
+ });
+ return filterProjectRexes;
+ }
+
private JoinNode convertLogicalAsofJoin(LogicalAsofJoin join) {
JoinInfo joinInfo = join.analyzeCondition();
DataSchema dataSchema = toDataSchema(join.getRowType());
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/SubPlanFragmenter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/SubPlanFragmenter.java
index ee0f20f2890..5e9b4dcd625 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/SubPlanFragmenter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/SubPlanFragmenter.java
@@ -26,6 +26,7 @@ import org.apache.calcite.runtime.ImmutablePairList;
import org.apache.pinot.calcite.rel.logical.PinotRelExchangeType;
import org.apache.pinot.query.planner.SubPlanMetadata;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -79,6 +80,11 @@ public class SubPlanFragmenter implements PlanNodeVisitor<PlanNode, SubPlanFragm
return process(node, context);
}
+ @Override
+ public PlanNode visitEnrichedJoin(EnrichedJoinNode node, Context context) {
+ return visitJoin(node, context);
+ }
+
@Override
public PlanNode visitMailboxReceive(MailboxReceiveNode node, Context context) {
throw new UnsupportedOperationException("MailboxReceiveNode should not be visited by StageFragmenter");
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
index 77ca8a68e80..d89e18f7ff3 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java
@@ -27,6 +27,7 @@ import org.apache.pinot.common.config.provider.TableCache;
import org.apache.pinot.core.routing.LogicalTableRouteInfo;
import org.apache.pinot.core.routing.LogicalTableRouteProvider;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -104,6 +105,12 @@ public class DispatchablePlanVisitor implements PlanNodeVisitor<Void, Dispatchab
return null;
}
+ @Override
+ public Void visitEnrichedJoin(EnrichedJoinNode node, DispatchablePlanContext context) {
+ visitJoin(node, context);
+ return null;
+ }
+
@Override
public Void visitMailboxReceive(MailboxReceiveNode node, DispatchablePlanContext context) {
node.getSender().visit(this, context);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/RelToPRelConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/RelToPRelConverter.java
index 09323341b1e..59f2df42e96 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/RelToPRelConverter.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/RelToPRelConverter.java
@@ -35,6 +35,7 @@ import org.apache.calcite.rel.core.Union;
import org.apache.calcite.rel.core.Values;
import org.apache.calcite.rel.core.Window;
import org.apache.pinot.calcite.rel.logical.PinotLogicalAggregate;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
import org.apache.pinot.calcite.rel.rules.PinotRuleUtils;
import org.apache.pinot.calcite.rel.traits.TraitAssignment;
import org.apache.pinot.common.config.provider.TableCache;
@@ -111,6 +112,9 @@ public class RelToPRelConverter {
return new PhysicalAsOfJoin(asofJoin.getCluster(), asofJoin.getTraitSet(), asofJoin.getHints(),
asofJoin.getCondition(), asofJoin.getMatchCondition(), asofJoin.getVariablesSet(), asofJoin.getJoinType(),
nodeIdGenerator.get(), inputs.get(0), inputs.get(1), null);
+ } else if (relNode instanceof PinotLogicalEnrichedJoin) {
+ // this should be unreachable
+ throw new IllegalStateException("EnrichedJoin is not supported in physical optimizer");
} else if (relNode instanceof Join) {
Preconditions.checkState(relNode.getInputs().size() == 2, "Expected exactly 2 inputs to join. Found: %s", inputs);
Join join = (Join) relNode;
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalJoin.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalJoin.java
index 74f5d9dc9d6..940ed88ac52 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalJoin.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/nodes/PhysicalJoin.java
@@ -24,11 +24,13 @@ import javax.annotation.Nullable;
import org.apache.calcite.plan.RelOptCluster;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelWriter;
import org.apache.calcite.rel.core.CorrelationId;
import org.apache.calcite.rel.core.Join;
import org.apache.calcite.rel.core.JoinRelType;
import org.apache.calcite.rel.hint.RelHint;
import org.apache.calcite.rex.RexNode;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
import org.apache.pinot.query.planner.physical.v2.PRelNode;
import org.apache.pinot.query.planner.physical.v2.PinotDataDistribution;
@@ -38,6 +40,8 @@ public class PhysicalJoin extends Join implements PRelNode {
private final List<PRelNode> _pRelInputs;
@Nullable
private final PinotDataDistribution _pinotDataDistribution;
+ @Nullable
+ private final List<PinotLogicalEnrichedJoin.FilterProjectRexNode> _filterProjectRexNodes;
public PhysicalJoin(RelOptCluster cluster, RelTraitSet traitSet, List<RelHint> hints,
RexNode condition, Set<CorrelationId> variablesSet, JoinRelType joinType, int nodeId, PRelNode left,
@@ -46,13 +50,25 @@ public class PhysicalJoin extends Join implements PRelNode {
_nodeId = nodeId;
_pRelInputs = List.of(left, right);
_pinotDataDistribution = pinotDataDistribution;
+ _filterProjectRexNodes = null;
+ }
+
+ public PhysicalJoin(RelOptCluster cluster, RelTraitSet traitSet, List<RelHint> hints,
+ RexNode condition, Set<CorrelationId> variablesSet, JoinRelType joinType, int nodeId, PRelNode left,
+ PRelNode right, @Nullable PinotDataDistribution pinotDataDistribution,
+ List<PinotLogicalEnrichedJoin.FilterProjectRexNode> filterProjectRexNodes) {
+ super(cluster, traitSet, hints, left.unwrap(), right.unwrap(), condition, variablesSet, joinType);
+ _nodeId = nodeId;
+ _pRelInputs = List.of(left, right);
+ _pinotDataDistribution = pinotDataDistribution;
+ _filterProjectRexNodes = filterProjectRexNodes;
}
@Override
public Join copy(RelTraitSet traitSet, RexNode conditionExpr, RelNode left, RelNode right, JoinRelType joinType,
boolean semiJoinDone) {
return new PhysicalJoin(getCluster(), traitSet, getHints(), conditionExpr, getVariablesSet(), joinType, _nodeId,
- (PRelNode) left, (PRelNode) right, _pinotDataDistribution);
+ (PRelNode) left, (PRelNode) right, _pinotDataDistribution, _filterProjectRexNodes);
}
@Override
@@ -84,6 +100,12 @@ public class PhysicalJoin extends Join implements PRelNode {
@Override
public PRelNode with(int newNodeId, List<PRelNode> newInputs, PinotDataDistribution newDistribution) {
return new PhysicalJoin(getCluster(), getTraitSet(), getHints(), getCondition(), getVariablesSet(), getJoinType(),
- newNodeId, newInputs.get(0), newInputs.get(1), newDistribution);
+ newNodeId, newInputs.get(0), newInputs.get(1), newDistribution, _filterProjectRexNodes);
+ }
+
+ @Override
+ public RelWriter explainTerms(RelWriter pw) {
+ return super.explainTerms(pw)
+ .itemIf("filterProjectRex", _filterProjectRexNodes, _filterProjectRexNodes != null);
}
}
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/DefaultPostOrderTraversalVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/DefaultPostOrderTraversalVisitor.java
index 63ba2fa3e90..225a94660ae 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/DefaultPostOrderTraversalVisitor.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/DefaultPostOrderTraversalVisitor.java
@@ -45,6 +45,11 @@ public abstract class DefaultPostOrderTraversalVisitor<T, C> implements PlanNode
return process(node, context);
}
+ @Override
+ public T visitEnrichedJoin(EnrichedJoinNode node, C context) {
+ return visitJoin(node, context);
+ }
+
@Override
public T visitMailboxReceive(MailboxReceiveNode node, C context) {
node.getSender().visit(this, context);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/EnrichedJoinNode.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/EnrichedJoinNode.java
new file mode 100644
index 00000000000..7e44f3d9a21
--- /dev/null
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/EnrichedJoinNode.java
@@ -0,0 +1,160 @@
+/**
+ * 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.pinot.query.planner.plannode;
+
+import java.util.List;
+import javax.annotation.Nullable;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.planner.logical.RexExpression;
+
+
+public class EnrichedJoinNode extends JoinNode {
+ private final List<FilterProjectRex> _filterProjectRexes;
+ /// Output schema of the join
+ private final DataSchema _joinResultSchema;
+ /// Output schema after projection, same as _joinResultSchema if no projection
+ private final DataSchema _projectResultSchema;
+ private final int _fetch;
+ private final int _offset;
+
+
+ public EnrichedJoinNode(int stageId, DataSchema joinResultSchema, DataSchema projectResultSchema,
+ NodeHint nodeHint, List<PlanNode> inputs,
+ JoinRelType joinType, List<Integer> leftKeys, List<Integer> rightKeys,
+ List<RexExpression> nonEquiConditions, JoinStrategy joinStrategy, RexExpression matchCondition,
+ List<FilterProjectRex> filterProjectRexes,
+ int fetch, int offset) {
+ super(stageId, projectResultSchema, nodeHint, inputs, joinType, leftKeys, rightKeys,
+ nonEquiConditions, joinStrategy, matchCondition);
+ _joinResultSchema = joinResultSchema;
+ _projectResultSchema = projectResultSchema;
+ _filterProjectRexes = filterProjectRexes;
+ _fetch = fetch;
+ _offset = offset;
+ }
+
+ public DataSchema getJoinResultSchema() {
+ return _joinResultSchema;
+ }
+
+ /// The final output schema is project output schema since only project and join alters schema
+ @Override
+ public DataSchema getDataSchema() {
+ return _projectResultSchema;
+ }
+
+ public List<FilterProjectRex> getFilterProjectRexes() {
+ return _filterProjectRexes;
+ }
+
+ @Override
+ public <T, C> T visit(PlanNodeVisitor<T, C> visitor, C context) {
+ return visitor.visitEnrichedJoin(this, context);
+ }
+
+ @Override
+ public String explain() {
+ return "ENRICHED_JOIN";
+ }
+
+ @Override
+ public PlanNode withInputs(List<PlanNode> inputs) {
+ return new EnrichedJoinNode(_stageId, _joinResultSchema, _projectResultSchema, _nodeHint,
+ inputs, getJoinType(), getLeftKeys(), getRightKeys(), getNonEquiConditions(),
+ getJoinStrategy(), getMatchCondition(), _filterProjectRexes, _fetch, _offset);
+ }
+
+ public int getFetch() {
+ return _fetch;
+ }
+
+ public int getOffset() {
+ return _offset;
+ }
+
+ public enum FilterProjectRexType {
+ FILTER,
+ PROJECT
+ }
+
+ public static class FilterProjectRex {
+
+ public static class ProjectAndResultSchema {
+ private final List<RexExpression> _project;
+ private final DataSchema _schema;
+
+ private ProjectAndResultSchema(List<RexExpression> project, DataSchema resultSchema) {
+ _project = project;
+ _schema = resultSchema;
+ }
+
+ public List<RexExpression> getProject() {
+ return _project;
+ }
+
+ public DataSchema getSchema() {
+ return _schema;
+ }
+ }
+
+ private final FilterProjectRexType _type;
+ @Nullable
+ private final RexExpression _filter;
+ @Nullable
+ private final ProjectAndResultSchema _projectAndResultSchema;
+
+ public FilterProjectRex(RexExpression filter) {
+ _type = FilterProjectRexType.FILTER;
+ _filter = filter;
+ _projectAndResultSchema = null;
+ }
+
+ public FilterProjectRex(List<RexExpression> projects, DataSchema resultSchema) {
+ _type = FilterProjectRexType.PROJECT;
+ _filter = null;
+ _projectAndResultSchema = new ProjectAndResultSchema(projects, resultSchema);
+ }
+
+ @Nullable
+ public RexExpression getFilter() {
+ return _filter;
+ }
+
+ @Nullable
+ public ProjectAndResultSchema getProjectAndResultSchema() {
+ return _projectAndResultSchema;
+ }
+
+ public FilterProjectRexType getType() {
+ return _type;
+ }
+
+ @Override
+ public String toString() {
+ if (_type == FilterProjectRexType.FILTER) {
+ assert _filter != null;
+ return "Filter: " + _filter;
+ } else {
+ assert _projectAndResultSchema != null;
+ return "Project: " + _projectAndResultSchema.getProject().toString();
+ }
+ }
+ }
+}
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java
index 49494f8df65..2d922d53048 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/plannode/PlanNodeVisitor.java
@@ -44,6 +44,8 @@ public interface PlanNodeVisitor<T, C> {
T visitJoin(JoinNode node, C context);
+ T visitEnrichedJoin(EnrichedJoinNode node, C context);
+
T visitMailboxReceive(MailboxReceiveNode node, C context);
T visitMailboxSend(MailboxSendNode node, C context);
@@ -160,6 +162,13 @@ public interface PlanNodeVisitor<T, C> {
return postChildren(node, context);
}
+ @Override
+ public T visitEnrichedJoin(EnrichedJoinNode node, C context) {
+ preChildren(node, context);
+ visitChildren(node, context);
+ return postChildren(node, context);
+ }
+
@Override
public T visitMailboxReceive(MailboxReceiveNode node, C context) {
preChildren(node, context);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
index b9701323471..8abfd62572e 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeDeserializer.java
@@ -34,6 +34,7 @@ import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.partitioning.KeySelector;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
import org.apache.pinot.query.planner.plannode.JoinNode;
@@ -78,6 +79,8 @@ public class PlanNodeDeserializer {
return deserializeWindowNode(protoNode);
case EXPLAINNODE:
return deserializeExplainedNode(protoNode);
+ case ENRICHEDJOINNODE:
+ return deserializeEnrichedJoinNode(protoNode);
default:
throw new IllegalStateException("Unsupported PlanNode type: " + protoNode.getNodeCase());
}
@@ -108,6 +111,34 @@ public class PlanNodeDeserializer {
protoJoinNode.getMatchCondition()) : null);
}
+ private static EnrichedJoinNode deserializeEnrichedJoinNode(Plan.PlanNode protoNode) {
+ Plan.EnrichedJoinNode protoEnrichedJoinNode = protoNode.getEnrichedJoinNode();
+ // reconstruct filterProjectRex
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = new ArrayList<>();
+ for (Plan.FilterProjectRex rex : protoEnrichedJoinNode.getFilterProjectRexList()) {
+ if (rex.getType() == Plan.FilterProjectRexType.FILTER) {
+ filterProjectRexes.add(
+ new EnrichedJoinNode.FilterProjectRex(ProtoExpressionToRexExpression.convertExpression(rex.getFilter())));
+ } else {
+ filterProjectRexes.add(
+ new EnrichedJoinNode.FilterProjectRex(convertExpressions(rex.getProjectAndResultSchema().getProjectList()),
+ extractDataSchema(rex.getProjectAndResultSchema().getSchema())));
+ }
+ }
+ return new EnrichedJoinNode(protoNode.getStageId(),
+ extractDataSchema(protoEnrichedJoinNode.getJoinResultDataSchema()), extractDataSchema(protoNode),
+ extractNodeHint(protoNode), extractInputs(protoNode), convertJoinType(protoEnrichedJoinNode.getJoinType()),
+ protoEnrichedJoinNode.getLeftKeysList(),
+ protoEnrichedJoinNode.getRightKeysList(), convertExpressions(protoEnrichedJoinNode.getNonEquiConditionsList()),
+ convertJoinStrategy(protoEnrichedJoinNode.getJoinStrategy()),
+ protoEnrichedJoinNode.hasMatchCondition()
+ ? ProtoExpressionToRexExpression.convertExpression(protoEnrichedJoinNode.getMatchCondition()) : null,
+ filterProjectRexes,
+ protoEnrichedJoinNode.getFetch(),
+ protoEnrichedJoinNode.getOffset()
+ );
+ }
+
private static MailboxReceiveNode deserializeMailboxReceiveNode(Plan.PlanNode protoNode) {
List<PlanNode> planNodes = extractInputs(protoNode);
Preconditions.checkState(planNodes.isEmpty(), "MailboxReceiveNode should not have inputs but has: %s", planNodes);
@@ -188,6 +219,17 @@ public class PlanNodeDeserializer {
extractInputs(protoNode), protoExplainNode.getTitle(), protoExplainNode.getAttributesMap());
}
+ private static DataSchema extractDataSchema(Plan.DataSchema protoDataSchema) {
+ String[] columnNames = protoDataSchema.getColumnNamesList().toArray(new String[0]);
+ int numColumns = columnNames.length;
+ List<Expressions.ColumnDataType> protoColumnDataTypes = protoDataSchema.getColumnDataTypesList();
+ ColumnDataType[] columnDataTypes = new ColumnDataType[numColumns];
+ for (int i = 0; i < numColumns; i++) {
+ columnDataTypes[i] = ProtoExpressionToRexExpression.convertColumnDataType(protoColumnDataTypes.get(i));
+ }
+ return new DataSchema(columnNames, columnDataTypes);
+ }
+
private static DataSchema extractDataSchema(Plan.PlanNode protoNode) {
Plan.DataSchema protoDataSchema = protoNode.getDataSchema();
String[] columnNames = protoDataSchema.getColumnNamesList().toArray(new String[0]);
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
index c23f3bfbd4d..6124609fbb0 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/PlanNodeSerializer.java
@@ -31,6 +31,7 @@ import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -130,6 +131,44 @@ public class PlanNodeSerializer {
return null;
}
+ @Override
+ public Void visitEnrichedJoin(EnrichedJoinNode node, Plan.PlanNode.Builder builder) {
+ Plan.EnrichedJoinNode.Builder enrichedJoinNode = Plan.EnrichedJoinNode.newBuilder()
+ .setJoinType(convertJoinType(node.getJoinType()))
+ .addAllLeftKeys(node.getLeftKeys())
+ .addAllRightKeys(node.getRightKeys())
+ .addAllNonEquiConditions(convertExpressions(node.getNonEquiConditions()))
+ .setJoinStrategy(convertJoinStrategy(node.getJoinStrategy()))
+ .setJoinResultDataSchema(convertDataSchema(node.getJoinResultSchema()));
+
+ for (EnrichedJoinNode.FilterProjectRex rex : node.getFilterProjectRexes()) {
+ Plan.FilterProjectRex.Builder rexBuilder = Plan.FilterProjectRex.newBuilder();
+ if (rex.getType() == EnrichedJoinNode.FilterProjectRexType.FILTER) {
+ rexBuilder
+ .setFilter(RexExpressionToProtoExpression.convertExpression(rex.getFilter()))
+ .setType(Plan.FilterProjectRexType.FILTER);
+ } else {
+ rexBuilder
+ .setProjectAndResultSchema(
+ Plan.ProjectAndResultSchema.newBuilder()
+ .addAllProject(convertExpressions(rex.getProjectAndResultSchema().getProject()))
+ .setSchema(convertDataSchema(rex.getProjectAndResultSchema().getSchema()))
+ .build())
+ .setType(Plan.FilterProjectRexType.PROJECT);
+ }
+ enrichedJoinNode.addFilterProjectRex(rexBuilder.build());
+ }
+
+ enrichedJoinNode.setFetch(node.getFetch());
+ enrichedJoinNode.setOffset(node.getOffset());
+
+ if (node.getMatchCondition() != null) {
+ enrichedJoinNode.setMatchCondition(RexExpressionToProtoExpression.convertExpression(node.getMatchCondition()));
+ }
+ builder.setEnrichedJoinNode(enrichedJoinNode.build());
+ return null;
+ }
+
@Override
public Void visitMailboxReceive(MailboxReceiveNode node, Plan.PlanNode.Builder builder) {
Plan.MailboxReceiveNode mailboxReceiveNode = Plan.MailboxReceiveNode.newBuilder()
diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java
index a31c446c20f..445520bd7d8 100644
--- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java
+++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/validation/ArrayToMvValidationVisitor.java
@@ -21,6 +21,7 @@ package org.apache.pinot.query.planner.validation;
import java.util.List;
import org.apache.pinot.query.planner.logical.RexExpression;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -65,6 +66,12 @@ public class ArrayToMvValidationVisitor implements PlanNodeVisitor<Void, Boolean
return null;
}
+ @Override
+ public Void visitEnrichedJoin(EnrichedJoinNode node, Boolean isIntermediateStage) {
+ visitJoin(node, isIntermediateStage);
+ return null;
+ }
+
@Override
public Void visitMailboxReceive(MailboxReceiveNode node, Boolean isIntermediateStage) {
node.getInputs().forEach(e -> e.visit(this, isIntermediateStage));
diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRuleTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRuleTest.java
new file mode 100644
index 00000000000..fedcfba89cf
--- /dev/null
+++ b/pinot-query-planner/src/test/java/org/apache/pinot/calcite/rel/rules/PinotEnrichedJoinRuleTest.java
@@ -0,0 +1,638 @@
+/**
+ * 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.pinot.calcite.rel.rules;
+
+import com.google.common.collect.ImmutableList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.plan.hep.HepMatchOrder;
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.logical.LogicalSort;
+import org.apache.calcite.rel.logical.LogicalValues;
+import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
+import org.apache.calcite.rel.rules.CoreRules;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.sql.validate.SqlValidatorUtil;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
+import org.apache.pinot.query.type.TypeFactory;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+
+
+public class PinotEnrichedJoinRuleTest {
+ private static final TypeFactory TYPE_FACTORY = TypeFactory.INSTANCE;
+ private static final RexBuilder REX_BUILDER = RexBuilder.DEFAULT;
+
+ private final RelDataType _intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ private LogicalValues _input;
+
+ @BeforeMethod
+ public void setUp() {
+ RelTraitSet traits = RelTraitSet.createEmpty();
+
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+ _input = new LogicalValues(cluster, traits, _intType, ImmutableList.of());
+ }
+
+ @Test
+ public void testRuleWithPlannerFilterJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // filter condition col2 = 1
+ RexNode filterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 2),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+
+ planner.setRoot(originalFilter);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerFilterEnrichedJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ RexNode prevFilterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeLiteral(2, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ PinotLogicalEnrichedJoin originalJoin =
+ new PinotLogicalEnrichedJoin(cluster, cluster.traitSet(), List.of(), _input, _input, joinCondition,
+ Collections.emptySet(),
+ JoinRelType.INNER, List.of(new PinotLogicalEnrichedJoin.FilterProjectRexNode(prevFilterCondition)), null,
+ null, null, null);
+ // filter condition col2 = 1
+ RexNode filterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 2),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+
+ planner.setRoot(originalFilter);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().size(), 2);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getFilter(), prevFilterCondition);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ assertEquals(enrichedJoin.getRowType(), enrichedJoin.getJoinRowType());
+ }
+
+ @Test
+ public void testRuleWithPlannerProjectJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // project one column
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 1));
+ LogicalProject project =
+ LogicalProject.create(originalJoin, Collections.emptyList(), projects, List.of("projectCol1"));
+
+ planner.setRoot(project);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getProjects(), projects);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerProjectEnrichedJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ RexNode prevFilterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeLiteral(2, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+
+ // old project
+ List<RexNode> oldProjects = List.of(REX_BUILDER.makeInputRef(intType, 1), REX_BUILDER.makeInputRef(intType, 0));
+ RelDataType oldRowType =
+ RexUtil.createStructType(cluster.getTypeFactory(), oldProjects,
+ List.of("projectCol2", "projectCol1"), SqlValidatorUtil.F_SUGGESTER);
+
+ // one project and one filter
+ PinotLogicalEnrichedJoin originalJoin =
+ new PinotLogicalEnrichedJoin(cluster, cluster.traitSet(), List.of(), _input, _input, joinCondition,
+ Collections.emptySet(),
+ JoinRelType.INNER, List.of(new PinotLogicalEnrichedJoin.FilterProjectRexNode(oldProjects, oldRowType),
+ new PinotLogicalEnrichedJoin.FilterProjectRexNode(prevFilterCondition)), oldRowType,
+ null, null, null);
+
+ // new project
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 1));
+ LogicalProject project =
+ LogicalProject.create(originalJoin, Collections.emptyList(), projects, List.of("projectCol1"));
+ RelDataType rowType =
+ RexUtil.createStructType(cluster.getTypeFactory(), projects,
+ List.of("projectCol1"), SqlValidatorUtil.F_SUGGESTER);
+
+ planner.setRoot(project);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().size(), 3);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getProject(), oldProjects);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getDataType(), oldRowType);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getFilter(), prevFilterCondition);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(2).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(2).getProjectAndResultRowType().getDataType(), rowType);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ assertEquals(enrichedJoin.getRowType(), rowType);
+ }
+
+ @Test
+ public void testRuleWithPlannerLimitOffsetJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+
+ RelCollation collation = RelCollations.EMPTY;
+ RexNode limit = literal(1);
+ RexNode offset = literal(1);
+ LogicalSort sort = LogicalSort.create(originalJoin, collation, offset, limit);
+
+ planner.setRoot(sort);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getOffset(), offset);
+ assertEquals(enrichedJoin.getFetch(), limit);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerLimitOffsetEnrichedJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+
+ // old filter
+ RexNode prevFilterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeLiteral(2, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+
+ // old project
+ List<RexNode> oldProjects = List.of(REX_BUILDER.makeInputRef(intType, 1), REX_BUILDER.makeInputRef(intType, 0));
+ RelDataType oldRowType =
+ RexUtil.createStructType(cluster.getTypeFactory(), oldProjects,
+ List.of("projectCol2", "projectCol1"), SqlValidatorUtil.F_SUGGESTER);
+
+ // one project and one filter
+ PinotLogicalEnrichedJoin originalJoin =
+ new PinotLogicalEnrichedJoin(cluster, cluster.traitSet(), List.of(), _input, _input, joinCondition,
+ Collections.emptySet(),
+ JoinRelType.INNER, List.of(new PinotLogicalEnrichedJoin.FilterProjectRexNode(oldProjects, oldRowType),
+ new PinotLogicalEnrichedJoin.FilterProjectRexNode(prevFilterCondition)), oldRowType,
+ null, null, null);
+
+ RelCollation collation = RelCollations.EMPTY;
+ RexNode limit = literal(1);
+ RexNode offset = literal(1);
+ LogicalSort sort = LogicalSort.create(originalJoin, collation, offset, limit);
+
+ planner.setRoot(sort);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().size(), 2);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getProject(), oldProjects);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getDataType(), oldRowType);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getFilter(), prevFilterCondition);
+ assertEquals(enrichedJoin.getOffset(), offset);
+ assertEquals(enrichedJoin.getFetch(), limit);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ assertEquals(enrichedJoin.getRowType(), oldRowType);
+ }
+
+ @Test
+ public void testRuleWithPlannerProjectFilterJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // filter condition col2 = 1
+ RexNode filterCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 2),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+ // project above filter
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 1));
+ LogicalProject project =
+ LogicalProject.create(originalFilter, Collections.emptyList(), projects, List.of("projectCol1"));
+
+ planner.setRoot(project);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerFilterProjectJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // project above filter with width of 5
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 0), REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeInputRef(intType, 2), REX_BUILDER.makeInputRef(intType, 3),
+ REX_BUILDER.makeInputRef(intType, 4), REX_BUILDER.makeInputRef(intType, 4),
+ REX_BUILDER.makeInputRef(intType, 0));
+ LogicalProject project =
+ LogicalProject.create(originalJoin, Collections.emptyList(), projects, (List<? extends String>) null);
+ // filter condition col2 = 1
+ RexNode filterCondition =
+ REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(project.getRowType(), 6),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(project, filterCondition);
+
+ planner.setRoot(originalFilter);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerLimitFilterProjectJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // project above filter with width of 5
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 0), REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeInputRef(intType, 2), REX_BUILDER.makeInputRef(intType, 3),
+ REX_BUILDER.makeInputRef(intType, 4), REX_BUILDER.makeInputRef(intType, 4),
+ REX_BUILDER.makeInputRef(intType, 0));
+ LogicalProject project =
+ LogicalProject.create(originalJoin, Collections.emptyList(), projects, (List<? extends String>) null);
+ // filter condition col2 = 1
+ RexNode filterCondition =
+ REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(project.getRowType(), 6),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(project, filterCondition);
+
+ RelCollation collation = RelCollations.EMPTY;
+ RexNode limit = literal(1);
+ RexNode offset = literal(1);
+ LogicalSort sort = LogicalSort.create(originalFilter, collation, offset, limit);
+
+ planner.setRoot(sort);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getOffset(), offset);
+ assertEquals(enrichedJoin.getFetch(), limit);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleWithPlannerLimitProjectFilterJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // filter condition col2 = 1
+ RexNode filterCondition =
+ REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(originalJoin.getRowType(), 0),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+ // project above filter with width of 5
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 0), REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeInputRef(intType, 2), REX_BUILDER.makeInputRef(intType, 3),
+ REX_BUILDER.makeInputRef(intType, 4), REX_BUILDER.makeInputRef(intType, 4),
+ REX_BUILDER.makeInputRef(intType, 0));
+ LogicalProject project =
+ LogicalProject.create(originalFilter, Collections.emptyList(), projects, (List<? extends String>) null);
+
+ RelCollation collation = RelCollations.EMPTY;
+ RexNode limit = literal(1);
+ RexNode offset = literal(1);
+ LogicalSort sort = LogicalSort.create(project, collation, offset, limit);
+
+ planner.setRoot(sort);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getOffset(), offset);
+ assertEquals(enrichedJoin.getFetch(), limit);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ }
+
+ @Test
+ public void testRuleRejectSortWithCollationProjectFilterJoinCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col0 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // filter condition col2 = 1
+ RexNode filterCondition =
+ REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(originalJoin.getRowType(), 0),
+ REX_BUILDER.makeLiteral(1, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+ // project above filter with width of 5
+ List<RexNode> projects = List.of(REX_BUILDER.makeInputRef(intType, 0), REX_BUILDER.makeInputRef(intType, 1),
+ REX_BUILDER.makeInputRef(intType, 2), REX_BUILDER.makeInputRef(intType, 3),
+ REX_BUILDER.makeInputRef(intType, 4), REX_BUILDER.makeInputRef(intType, 4),
+ REX_BUILDER.makeInputRef(intType, 0));
+ LogicalProject project =
+ LogicalProject.create(originalFilter, Collections.emptyList(), projects, (List<? extends String>) null);
+
+ RelCollation collation = RelCollations.of(0);
+ RexNode limit = literal(1);
+ RexNode offset = literal(1);
+ LogicalSort sort = LogicalSort.create(project, collation, offset, limit);
+
+ planner.setRoot(sort);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof LogicalSort);
+ assert (newRoot.getInput(0) instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot.getInput(0);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(0).getFilter(), filterCondition);
+ assertEquals(enrichedJoin.getFilterProjectRexNodes().get(1).getProjectAndResultRowType().getProject(), projects);
+ assertEquals(enrichedJoin.getCondition(), joinCondition);
+ assertNull(enrichedJoin.getOffset());
+ assertNull(enrichedJoin.getFetch());
+ }
+
+ @Test
+ public void testRuleWithPlannerProjectJoinAfterProjectPushdownCase() {
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ // add projection pushdown rule
+ hepProgramBuilder.addRuleInstance(CoreRules.PROJECT_JOIN_TRANSPOSE);
+ // add enriched join rule
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ hepProgramBuilder.addMatchOrder(HepMatchOrder.BOTTOM_UP);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, REX_BUILDER);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = TYPE_FACTORY.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ // join condition col1 = col1
+ RexNode joinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS, REX_BUILDER.makeInputRef(intType, 0),
+ REX_BUILDER.makeInputRef(intType, 3));
+ LogicalJoin originalJoin =
+ LogicalJoin.create(_input, _input, Collections.emptyList(), joinCondition, Collections.emptySet(),
+ JoinRelType.INNER);
+ // project that touches both left and right relations
+ List<RexNode> projects = List.of(REX_BUILDER.makeCall(SqlStdOperatorTable.PLUS,
+ REX_BUILDER.makeInputRef(intType, 1), REX_BUILDER.makeInputRef(intType, 4)));
+ LogicalProject project =
+ LogicalProject.create(originalJoin, Collections.emptyList(), projects, List.of("projectCol1"));
+
+ planner.setRoot(project);
+ RelNode newRoot = planner.findBestExp();
+
+ assert (newRoot instanceof PinotLogicalEnrichedJoin);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) newRoot;
+ // projection pushdown should trim both inputs to first two cols only
+ assert (newRoot.getInput(0) instanceof LogicalProject);
+ assert (newRoot.getInput(1) instanceof LogicalProject);
+ LogicalProject leftChild = (LogicalProject) newRoot.getInput(0);
+ LogicalProject rightChild = (LogicalProject) newRoot.getInput(1);
+ assertEquals(leftChild.getProjects().size(), 2);
+ assertEquals(rightChild.getProjects().size(), 2);
+ // join output layout [t1.col1, t1.col2, t2.col1, t2.col2]
+ List<RexNode> expectedProjects = List.of(REX_BUILDER.makeCall(SqlStdOperatorTable.PLUS,
+ REX_BUILDER.makeInputRef(intType, 1), REX_BUILDER.makeInputRef(intType, 3)));
+ assertEquals(enrichedJoin.getProjects(), expectedProjects);
+ // join condition col1 = col1
+ RexNode expectedJoinCondition = REX_BUILDER.makeCall(SqlStdOperatorTable.EQUALS,
+ REX_BUILDER.makeInputRef(intType, 0), REX_BUILDER.makeInputRef(intType, 2));
+ assertEquals(enrichedJoin.getCondition(), expectedJoinCondition);
+ }
+
+ private static RexNode literal(int i) {
+ return REX_BUILDER.makeLiteral(i, TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER));
+ }
+}
diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java
index 846465da728..8a5848576ed 100644
--- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java
+++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java
@@ -18,13 +18,36 @@
*/
package org.apache.pinot.query.planner.logical;
+import com.google.common.collect.ImmutableList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.calcite.plan.hep.HepProgramBuilder;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rel.logical.LogicalProject;
+import org.apache.calcite.rel.logical.LogicalValues;
+import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider;
+import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeSystem;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlIdentifier;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.type.ArraySqlType;
import org.apache.calcite.sql.type.BasicSqlType;
import org.apache.calcite.sql.type.ObjectSqlType;
import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.pinot.calcite.rel.logical.PinotLogicalEnrichedJoin;
+import org.apache.pinot.calcite.rel.rules.PinotEnrichedJoinRule;
import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.spi.utils.CommonConstants;
import org.testng.Assert;
import org.testng.annotations.Test;
@@ -130,4 +153,54 @@ public class RelToPlanNodeConverterTest {
new ArraySqlType(new ObjectSqlType(SqlTypeName.VARBINARY, SqlIdentifier.STAR, true, null, null), true)),
DataSchema.ColumnDataType.BYTES_ARRAY);
}
+
+ @Test
+ public void testConvertEnrichedJoinNodeTest() {
+ final TypeFactory typeFactory = TypeFactory.INSTANCE;
+ final RexBuilder rexBuilder = RexBuilder.DEFAULT;
+ RelTraitSet traits = RelTraitSet.createEmpty();
+
+ HepProgramBuilder hepProgramBuilder = new HepProgramBuilder();
+ hepProgramBuilder.addRuleCollection(PinotEnrichedJoinRule.PINOT_ENRICHED_JOIN_RULES);
+ HepPlanner planner = new HepPlanner(hepProgramBuilder.build());
+ RelOptCluster cluster = RelOptCluster.create(planner, rexBuilder);
+ cluster.setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
+
+ RelDataType intType = typeFactory.builder()
+ .add("col1", SqlTypeName.INTEGER)
+ .add("col2", SqlTypeName.INTEGER)
+ .add("col3", SqlTypeName.INTEGER)
+ .build();
+
+ LogicalValues input = new LogicalValues(cluster, traits, intType, ImmutableList.of());
+
+ // join condition col0 = col1
+ RexNode joinCondition = rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
+ rexBuilder.makeInputRef(intType, 0), rexBuilder.makeInputRef(intType, 3));
+ LogicalJoin originalJoin = LogicalJoin.create(input, input, Collections.emptyList(),
+ joinCondition, Collections.emptySet(), JoinRelType.INNER);
+
+ // filter condition col2 = 1
+ RexNode filterCondition = rexBuilder.makeCall(
+ SqlStdOperatorTable.EQUALS, rexBuilder.makeInputRef(intType, 2),
+ rexBuilder.makeLiteral(1, typeFactory.createSqlType(SqlTypeName.INTEGER)));
+ LogicalFilter originalFilter = LogicalFilter.create(originalJoin, filterCondition);
+
+ // project above filter
+ List<RexNode> projects = List.of(rexBuilder.makeInputRef(intType, 1));
+ LogicalProject project = LogicalProject.create(
+ originalFilter, Collections.emptyList(), projects, List.of("projectCol1"));
+
+ planner.setRoot(project);
+ PinotLogicalEnrichedJoin enrichedJoin = (PinotLogicalEnrichedJoin) planner.findBestExp();
+
+ RelToPlanNodeConverter relToPlanNodeConverter = new RelToPlanNodeConverter(null,
+ CommonConstants.Broker.DEFAULT_BROKER_DEFAULT_HASH_FUNCTION);
+
+ PlanNode node = relToPlanNodeConverter.toPlanNode(enrichedJoin);
+ assert (node instanceof EnrichedJoinNode);
+
+ EnrichedJoinNode enrichedJoinNode = (EnrichedJoinNode) node;
+ Assert.assertEquals(enrichedJoinNode.getFilterProjectRexes().size(), 2);
+ }
}
diff --git a/pinot-query-planner/src/test/resources/queries/JoinPlans.json b/pinot-query-planner/src/test/resources/queries/JoinPlans.json
index c1fb4b1e52b..09b5881702b 100644
--- a/pinot-query-planner/src/test/resources/queries/JoinPlans.json
+++ b/pinot-query-planner/src/test/resources/queries/JoinPlans.json
@@ -861,6 +861,103 @@
"\n PinotLogicalTableScan(table=[[default, b]])",
"\n"
]
+ },
+ {
+ "description": "Enriched join with one projection",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT a.col1 + b.col1 FROM a JOIN b ON a.col1 = b.col1",
+ "output": [
+ "Execution Plan",
+ "\nPinotLogicalEnrichedJoin(condition=[=($0, $2)], joinType=[inner], filterProjectRex=[[Project: [+($1, $3)]]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
+ },
+ {
+ "description": "Enriched right join with one projection",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT a.col1 + b.col1 FROM a RIGHT JOIN b ON a.col1 = b.col1",
+ "output": [
+ "Execution Plan",
+ "\nPinotLogicalEnrichedJoin(condition=[=($0, $2)], joinType=[right], filterProjectRex=[[Project: [+($1, $3)]]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000)])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
+ },
+ {
+ "description": "Enriched anti join with one limit",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT a.col1 FROM a WHERE NOT EXISTS (SELECT 1 FROM b WHERE a.col1 = b.col1) LIMIT 10",
+ "output": [
+ "Execution Plan",
+ "\nLogicalSort(offset=[0], fetch=[10])",
+ "\n PinotLogicalSortExchange(distribution=[hash], collation=[[]], isSortOnSender=[false], isSortOnReceiver=[false])",
+ "\n PinotLogicalEnrichedJoin(condition=[=($0, $1)], joinType=[left], filterProjectRex=[[Filter: IS NULL($2), Project: [$0]]], limit=[10])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])",
+ "\n LogicalProject(col1=[$0], $f0=[true])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
+ },
+ {
+ "description": "Enriched join with one projection and limit",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT a.col1 + b.col1 FROM a JOIN b ON a.col1 = b.col1 LIMIT 100",
+ "output": [
+ "Execution Plan",
+ "\nLogicalSort(offset=[0], fetch=[100])",
+ "\n PinotLogicalSortExchange(distribution=[hash], collation=[[]], isSortOnSender=[false], isSortOnReceiver=[false])",
+ "\n PinotLogicalEnrichedJoin(condition=[=($0, $2)], joinType=[inner], filterProjectRex=[[Project: [+($1, $3)]]], limit=[100])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
+ },
+ {
+ "description": "Enriched join with two projections",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 FROM (SELECT a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 AS sum2 FROM a JOIN b ON a.col1 = b.col1) t;\n",
+ "output": [
+ "Execution Plan",
+ "\nPinotLogicalEnrichedJoin(condition=[=($0, $2)], joinType=[inner], filterProjectRex=[[Project: [+(+(+(+(+(+(+(+(+($1, $3), $1), $3), $1), $3), $1), $3), $1), $3)], Project: [+(+(+(+(+(+(+(+(+(+(+(+($0, $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0)]]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
+ },
+ {
+ "description": "Enriched join with two projections and one filter",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; EXPLAIN PLAN FOR SELECT * FROM (SELECT sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 + sum2 AS sum3 FROM (SELECT a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 + a.col1 + b.col1 AS sum2 FROM a JOIN b ON a.col1 = b.col1)) WHERE sum3 = 1;\n",
+ "output": [
+ "Execution Plan",
+ "\nPinotLogicalEnrichedJoin(condition=[=($0, $2)], joinType=[inner], filterProjectRex=[[Project: [+(+(+(+(+(+(+(+(+($1, $3), $1), $3), $1), $3), $1), $3), $1), $3)], Filter: =(+(+(+(+(+(+(+(+(+(+(+(+($0, $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), 1.0), Project: [+(+(+(+(+(+(+(+(+(+(+(+($0, $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0), $0)]]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, a]])",
+ "\n PinotLogicalExchange(distribution=[hash[0]])",
+ "\n LogicalProject(col1=[$0], EXPR$0=[CAST($0):DECIMAL(2000, 1000) NOT NULL])",
+ "\n PinotLogicalTableScan(table=[[default, b]])",
+ "\n"
+ ]
}
]
},
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
index e2140308053..8e51b84f6ec 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/InStageStatsTreeBuilder.java
@@ -27,6 +27,7 @@ import java.util.function.IntFunction;
import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.planner.plannode.AggregateNode;
import org.apache.pinot.query.planner.plannode.BasePlanNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -154,6 +155,11 @@ public class InStageStatsTreeBuilder implements PlanNodeVisitor<ObjectNode, InSt
}
}
+ @Override
+ public ObjectNode visitEnrichedJoin(EnrichedJoinNode node, Context context) {
+ return visitJoin(node, context);
+ }
+
@Override
public ObjectNode visitMailboxReceive(MailboxReceiveNode node, Context context) {
ObjectNode json = selfNode(MultiStageOperator.Type.MAILBOX_RECEIVE, context);
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
index 52d46fa9864..ce8b959ff2e 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/BaseJoinOperator.java
@@ -108,6 +108,27 @@ public abstract class BaseJoinOperator extends MultiStageOperator {
_joinOverflowMode = getJoinOverflowMode(metadata, nodeHint);
}
+ /// Constructor that takes the schema for NonEquiEvaluator as an argument
+ public BaseJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema,
+ MultiStageOperator rightInput, JoinNode node, DataSchema nonEquiEvaluationSchema) {
+ super(context);
+ _leftInput = leftInput;
+ _rightInput = rightInput;
+ _joinType = node.getJoinType();
+ _leftColumnSize = leftSchema.size();
+ _resultSchema = node.getDataSchema();
+ _resultColumnSize = _resultSchema.size();
+ List<RexExpression> nonEquiConditions = node.getNonEquiConditions();
+ _nonEquiEvaluators = new ArrayList<>(nonEquiConditions.size());
+ for (RexExpression nonEquiCondition : nonEquiConditions) {
+ _nonEquiEvaluators.add(TransformOperandFactory.getTransformOperand(nonEquiCondition, nonEquiEvaluationSchema));
+ }
+ Map<String, String> metadata = context.getOpChainMetadata();
+ PlanNode.NodeHint nodeHint = node.getNodeHint();
+ _maxRowsInJoin = getMaxRowsInJoin(metadata, nodeHint);
+ _joinOverflowMode = getJoinOverflowMode(metadata, nodeHint);
+ }
+
protected static int getMaxRowsInJoin(Map<String, String> opChainMetadata, @Nullable PlanNode.NodeHint nodeHint) {
if (nodeHint != null) {
Map<String, String> joinOptions = nodeHint.getHintOptions().get(PinotHintOptions.JOIN_HINT_OPTIONS);
@@ -383,13 +404,11 @@ public abstract class BaseJoinOperator extends MultiStageOperator {
}
/**
- * This util class is a view over the left and right row joined together
- * currently this is used for filtering and input of projection. So if the joined
- * tuple doesn't pass the predicate, the join result is not materialized into Object[].
- *
- * It is debatable whether we always want to use this instead of copying the tuple
+ * This util class is a view over the left and right row joined together.
+ * Currently, this is used for filtering and input of projection. So if the joined
+ * tuple doesn't pass the predicate, the join result is not materialized into {@code Object[]}.
*/
- private abstract static class JoinedRowView extends AbstractList<Object> implements List<Object> {
+ protected abstract static class JoinedRowView extends AbstractList<Object> implements List<Object> {
protected final int _leftSize;
protected final int _size;
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperator.java
new file mode 100644
index 00000000000..371cd6e3952
--- /dev/null
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperator.java
@@ -0,0 +1,417 @@
+/**
+ * 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.pinot.query.runtime.operator;
+
+import java.util.ArrayList;
+import java.util.BitSet;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.apache.pinot.query.runtime.operator.operands.TransformOperand;
+import org.apache.pinot.query.runtime.operator.operands.TransformOperandFactory;
+import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
+import org.apache.pinot.spi.utils.BooleanUtils;
+import org.apache.pinot.spi.utils.CommonConstants;
+
+@SuppressWarnings({"unchecked", "rawType"})
+public class EnrichedHashJoinOperator extends HashJoinOperator {
+ private static final String EXPLAIN_NAME = "ENRICHED_JOIN";
+ private final List<FilterProjectOperand> _filterProjectOperands;
+ private final int _resultColumnSize;
+ private final int _numRowsToKeep;
+ private int _rowsSeen = 0;
+ private int _numRowsToOffset;
+
+ public EnrichedHashJoinOperator(OpChainExecutionContext context,
+ MultiStageOperator leftInput, DataSchema leftSchema, MultiStageOperator rightInput,
+ EnrichedJoinNode node) {
+ super(context, leftInput, leftSchema, rightInput, node, node.getJoinResultSchema());
+
+ DataSchema joinResultSchema = node.getJoinResultSchema();
+
+ _resultColumnSize = joinResultSchema.size();
+
+ // a chain of filter and project operands
+ _filterProjectOperands = new ArrayList<>();
+ DataSchema currentSchema = joinResultSchema;
+ // iterate and convert filter and project rexes into operands with the correctly chained schema
+ for (EnrichedJoinNode.FilterProjectRex rex : node.getFilterProjectRexes()) {
+ _filterProjectOperands.add(new FilterProjectOperand(rex, currentSchema));
+ currentSchema = rex.getProjectAndResultSchema() == null
+ ? currentSchema : rex.getProjectAndResultSchema().getSchema();
+ }
+
+ int offset = Math.max(node.getOffset(), 0);
+ int fetch = node.getFetch();
+
+ // TODO: see if this need to be converted to input args
+ int defaultResponseLimit = CommonConstants.Broker.DEFAULT_BROKER_QUERY_RESPONSE_LIMIT;
+ _numRowsToKeep = fetch > 0 ? fetch + offset : defaultResponseLimit;
+ _numRowsToOffset = offset;
+ }
+
+ @Override
+ public String toExplainString() {
+ return EXPLAIN_NAME;
+ }
+
+ /// Filter, project on a joined row view
+ private void filterProjectLimit(List<Object> rowView, List<Object[]> rows) {
+ Object[] row = null;
+
+ for (FilterProjectOperand filterProjectOperand : _filterProjectOperands) {
+ if (filterProjectOperand.getType() == FilterProjectOperandsType.FILTER) {
+ // use rowView before reaching a project
+ assert (filterProjectOperand.getFilter() != null);
+ if (row == null
+ ? filterDiscardRow(rowView, filterProjectOperand.getFilter())
+ : filterDiscardRow(row, filterProjectOperand.getFilter())
+ ) {
+ return;
+ }
+ } else {
+ assert (filterProjectOperand.getProject() != null);
+ row = row == null
+ ? projectRow(rowView, filterProjectOperand.getProject())
+ : projectRow(row, filterProjectOperand.getProject());
+ }
+ }
+
+ if (rowNotNeeded()) {
+ return;
+ }
+
+ // if filter only, materialize rowView at the end
+ rows.add(row == null ? rowView.toArray() : row);
+ }
+
+ /// Filter, project, limit on the left and right row by creating a view
+ private void filterProjectLimit(Object[] leftRow, Object[] rightRow, List<Object[]> rows,
+ int resultColumnSize, int leftColumnSize) {
+ List<Object> rowView = JoinedRowView.of(leftRow, rightRow, resultColumnSize, leftColumnSize);
+ Object[] row = null;
+
+ for (FilterProjectOperand filterProjectOperand : _filterProjectOperands) {
+ if (filterProjectOperand.getType() == FilterProjectOperandsType.FILTER) {
+ // use rowView before reaching a project
+ assert (filterProjectOperand.getFilter() != null);
+ if (row == null
+ ? filterDiscardRow(rowView, filterProjectOperand.getFilter())
+ : filterDiscardRow(row, filterProjectOperand.getFilter())
+ ) {
+ return;
+ }
+ } else {
+ assert (filterProjectOperand.getProject() != null);
+ row = row == null
+ ? projectRow(rowView, filterProjectOperand.getProject())
+ : projectRow(row, filterProjectOperand.getProject());
+ }
+ }
+
+ if (rowNotNeeded()) {
+ return;
+ }
+
+ // if filter only, materialize rowView at the end
+ rows.add(row == null ? rowView.toArray() : row);
+ }
+
+ /// Limit on a row, return true if the limit reached before adding this row
+ private boolean rowNotNeeded() {
+ // limit only, terminate if enough rows
+ if (_rowsSeen++ == _numRowsToKeep) {
+ earlyTerminate();
+ logger().debug("EnrichedHashJoinOperator: seen enough rows with no sort, early terminating");
+ return true;
+ }
+ return false;
+ }
+
+ /// Filter a row by left and right child, return whether the row is discarded
+ private boolean filterDiscardRow(List<Object> rowView, TransformOperand filter) {
+ Object filterResult = filter.apply(rowView);
+ return !BooleanUtils.isTrueInternalValue(filterResult);
+ }
+
+ private boolean filterDiscardRow(Object[] row, TransformOperand filter) {
+ Object filterResult = filter.apply(row);
+ return !BooleanUtils.isTrueInternalValue(filterResult);
+ }
+
+ /// Return the projected row
+ private Object[] projectRow(List<Object> rowView, List<TransformOperand> project) {
+ Object[] resultRow = new Object[project.size()];
+ for (int i = 0; i < project.size(); i++) {
+ resultRow[i] = project.get(i).apply(rowView);
+ }
+ return resultRow;
+ }
+
+ /// Return the projected row from input row
+ private Object[] projectRow(Object[] row, List<TransformOperand> project) {
+ Object[] resultRow = new Object[project.size()];
+ for (int i = 0; i < project.size(); i++) {
+ resultRow[i] = project.get(i).apply(row);
+ }
+ return resultRow;
+ }
+
+ /// Read result from _priorityQueue if sort needed, else return rows
+ private List<Object[]> getOutputRows(List<Object[]> rows) {
+ if (_numRowsToOffset <= 0) {
+ return rows;
+ }
+ if (rows.size() > _numRowsToOffset) {
+ int rowSize = rows.size();
+ rows = rows.subList(_numRowsToOffset, rows.size());
+ _numRowsToOffset -= rowSize;
+ return rows;
+ }
+ _numRowsToOffset -= rows.size();
+ return Collections.emptyList();
+ }
+
+ /**
+ * Enriched version of buildNonMatchedRightRows that filter, project, sort-limit it
+ * @return filter, projected, sort-limited rows
+ */
+ @Override
+ protected List<Object[]> buildNonMatchRightRows() {
+ assert _rightTable != null : "Right table should not be null when building non-matched right rows";
+ assert _matchedRightRows != null : "Matched right rows should not be null when building non-matched right rows";
+ List<Object[]> rows = new ArrayList<>();
+ if (_rightTable.isKeysUnique()) {
+ for (Map.Entry<Object, Object> entry : _rightTable.entrySet()) {
+ Object[] rightRow = (Object[]) entry.getValue();
+ if (_matchedRightRows.containsKey(entry.getKey())) {
+ continue;
+ }
+ // join row with null, then project-merge-sort-limit
+ filterProjectLimit(null, rightRow, rows, _resultColumnSize, _leftColumnSize);
+ }
+ } else {
+ for (Map.Entry<Object, Object> entry : _rightTable.entrySet()) {
+ List<Object[]> rightRows = ((List<Object[]>) entry.getValue());
+ BitSet matchedIndices = _matchedRightRows.get(entry.getKey());
+ if (matchedIndices == null) {
+ for (Object[] rightRow : rightRows) {
+ filterProjectLimit(null, rightRow, rows, _resultColumnSize, _leftColumnSize);
+ }
+ } else {
+ int numRightRows = rightRows.size();
+ int unmatchedIndex = 0;
+ while ((unmatchedIndex = matchedIndices.nextClearBit(unmatchedIndex)) < numRightRows) {
+ filterProjectLimit(null, rightRows.get(unmatchedIndex++), rows, _resultColumnSize, _leftColumnSize);
+ }
+ }
+ }
+ }
+ // return the result, fetch from pq if there's sort-limit
+ return getOutputRows(rows);
+ }
+
+ private void handleUnmatchedLeftRow(Object[] leftRow, List<Object[]> rows) {
+ if (needUnmatchedLeftRows()) {
+ if (isMaxRowsLimitReached(rows.size())) {
+ return;
+ }
+ filterProjectLimit(leftRow, null, rows, _resultColumnSize, _leftColumnSize);
+ }
+ }
+
+ @Override
+ protected List<Object[]> buildJoinedRows(MseBlock.Data leftBlock) {
+ assert _rightTable != null : "Right table should not be null when building joined rows";
+ switch (_joinType) {
+ case SEMI:
+ return buildJoinedDataBlockSemi(leftBlock);
+ case ANTI:
+ return buildJoinedDataBlockAnti(leftBlock);
+ default: { // INNER, LEFT, RIGHT, FULL
+ if (_rightTable.isKeysUnique()) {
+ return buildJoinedDataBlockUniqueKeys(leftBlock);
+ } else {
+ return buildJoinedDataBlockDuplicateKeys(leftBlock);
+ }
+ }
+ }
+ }
+
+ /// MatchNonEquiConditions that takes the row view
+ protected final boolean matchNonEquiConditions(List<Object> rowView) {
+ if (_nonEquiEvaluators.isEmpty()) {
+ return true;
+ }
+ for (TransformOperand evaluator : _nonEquiEvaluators) {
+ if (!BooleanUtils.isTrueInternalValue(evaluator.apply(rowView))) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private List<Object[]> buildJoinedDataBlockUniqueKeys(MseBlock.Data leftBlock) {
+ assert _rightTable != null : "Right table should not be null when building joined rows";
+ List<Object[]> leftRows = leftBlock.asRowHeap().getRows();
+ ArrayList<Object[]> rows = new ArrayList<>(leftRows.size());
+
+ for (Object[] leftRow : leftRows) {
+ Object key = _leftKeySelector.getKey(leftRow);
+ Object[] rightRow = (Object[]) _rightTable.lookup(key);
+ if (rightRow == null) {
+ handleUnmatchedLeftRow(leftRow, rows);
+ } else {
+ List<Object> resultRowView = JoinedRowView.of(leftRow, rightRow, _resultColumnSize, _leftColumnSize);
+ if (matchNonEquiConditions(resultRowView)) {
+ if (isMaxRowsLimitReached(rows.size())) {
+ break;
+ }
+ // filter project sortLimit on the produced row
+ filterProjectLimit(resultRowView, rows);
+ if (_matchedRightRows != null) {
+ _matchedRightRows.put(key, BIT_SET_PLACEHOLDER);
+ }
+ } else {
+ handleUnmatchedLeftRow(leftRow, rows);
+ }
+ }
+ }
+
+ return getOutputRows(rows);
+ }
+
+ private List<Object[]> buildJoinedDataBlockDuplicateKeys(MseBlock.Data leftBlock) {
+ assert _rightTable != null : "Right table should not be null when building joined rows";
+ List<Object[]> leftRows = leftBlock.asRowHeap().getRows();
+ List<Object[]> rows = new ArrayList<>(leftRows.size());
+
+ for (Object[] leftRow : leftRows) {
+ Object key = _leftKeySelector.getKey(leftRow);
+ List<Object[]> rightRows = (List<Object[]>) _rightTable.lookup(key);
+ if (rightRows == null) {
+ handleUnmatchedLeftRow(leftRow, rows);
+ } else {
+ boolean maxRowsLimitReached = false;
+ boolean hasMatchForLeftRow = false;
+ int numRightRows = rightRows.size();
+ for (int i = 0; i < numRightRows; i++) {
+ List<Object> resultRowView = JoinedRowView.of(leftRow, rightRows.get(i), _resultColumnSize, _leftColumnSize);
+ if (matchNonEquiConditions(resultRowView)) {
+ if (isMaxRowsLimitReached(rows.size())) {
+ maxRowsLimitReached = true;
+ break;
+ }
+ // filter project sortLimit on the produced row
+ filterProjectLimit(resultRowView, rows);
+ hasMatchForLeftRow = true;
+ if (_matchedRightRows != null) {
+ _matchedRightRows.computeIfAbsent(key, k -> new BitSet(numRightRows)).set(i);
+ }
+ }
+ }
+ if (maxRowsLimitReached) {
+ break;
+ }
+ if (!hasMatchForLeftRow) {
+ handleUnmatchedLeftRow(leftRow, rows);
+ }
+ }
+ }
+
+ return getOutputRows(rows);
+ }
+
+ private List<Object[]> buildJoinedDataBlockAnti(MseBlock.Data leftBlock) {
+ assert _rightTable != null : "Right table should not be null when building joined rows";
+ List<Object[]> leftRows = leftBlock.asRowHeap().getRows();
+ List<Object[]> rows = new ArrayList<>(leftRows.size());
+
+ for (Object[] leftRow : leftRows) {
+ Object key = _leftKeySelector.getKey(leftRow);
+ // ANTI-JOIN only checks non-existence of the key
+ if (!_rightTable.containsKey(key)) {
+ filterProjectLimit(leftRow, null, rows, _leftColumnSize, _leftColumnSize);
+ }
+ }
+
+ return getOutputRows(rows);
+ }
+
+ private List<Object[]> buildJoinedDataBlockSemi(MseBlock.Data leftBlock) {
+ assert _rightTable != null : "Right table should not be null when building joined rows";
+ List<Object[]> leftRows = leftBlock.asRowHeap().getRows();
+ List<Object[]> rows = new ArrayList<>(leftRows.size());
+
+ for (Object[] leftRow : leftRows) {
+ Object key = _leftKeySelector.getKey(leftRow);
+ // SEMI-JOIN only checks existence of the key
+ if (_rightTable.containsKey(key)) {
+ filterProjectLimit(leftRow, null, rows, _leftColumnSize, _leftColumnSize);
+ }
+ }
+
+ return getOutputRows(rows);
+ }
+
+ public enum FilterProjectOperandsType {
+ FILTER,
+ PROJECT
+ }
+
+ public static class FilterProjectOperand {
+ private final FilterProjectOperandsType _type;
+ @Nullable
+ final TransformOperand _filter;
+ @Nullable
+ final List<TransformOperand> _project;
+
+ public FilterProjectOperand(EnrichedJoinNode.FilterProjectRex rex, DataSchema inputSchema) {
+ if (rex.getType() == EnrichedJoinNode.FilterProjectRexType.FILTER) {
+ _type = FilterProjectOperandsType.FILTER;
+ _filter = TransformOperandFactory.getTransformOperand(rex.getFilter(), inputSchema);
+ _project = null;
+ } else {
+ _type = FilterProjectOperandsType.PROJECT;
+ _filter = null;
+ List<TransformOperand> projects = new ArrayList<>();
+ assert (rex.getProjectAndResultSchema() != null);
+ rex.getProjectAndResultSchema().getProject().forEach((x) ->
+ projects.add(TransformOperandFactory.getTransformOperand(x, inputSchema)));
+ _project = projects;
+ }
+ }
+
+ public TransformOperand getFilter() {
+ return _filter;
+ }
+
+ public List<TransformOperand> getProject() {
+ return _project;
+ }
+
+ public FilterProjectOperandsType getType() {
+ return _type;
+ }
+ }
+}
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
index 0a33fcd12ac..dfec8c7e44c 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java
@@ -51,17 +51,17 @@ public class HashJoinOperator extends BaseJoinOperator {
private static final String EXPLAIN_NAME = "HASH_JOIN";
// Placeholder for BitSet in _matchedRightRows when all keys are unique in the right table.
- private static final BitSet BIT_SET_PLACEHOLDER = new BitSet(0);
+ protected static final BitSet BIT_SET_PLACEHOLDER = new BitSet(0);
- private final KeySelector<?> _leftKeySelector;
- private final KeySelector<?> _rightKeySelector;
+ protected final KeySelector<?> _leftKeySelector;
+ protected final KeySelector<?> _rightKeySelector;
@Nullable
- private LookupTable _rightTable;
+ protected LookupTable _rightTable;
// Track matched right rows for right join and full join to output non-matched right rows.
// TODO: Revisit whether we should use IntList or RoaringBitmap for smaller memory footprint.
// TODO: Optimize this
@Nullable
- private Map<Object, BitSet> _matchedRightRows;
+ protected Map<Object, BitSet> _matchedRightRows;
// Store null key rows separately for RIGHT and FULL JOINs
@Nullable
private List<Object[]> _nullKeyRightRows;
@@ -79,6 +79,18 @@ public class HashJoinOperator extends BaseJoinOperator {
_nullKeyRightRows = needUnmatchedRightRows() ? new ArrayList<>() : null;
}
+ /// Constructor that takes the schema for NonEquiEvaluator as an argument
+ public HashJoinOperator(OpChainExecutionContext context, MultiStageOperator leftInput, DataSchema leftSchema,
+ MultiStageOperator rightInput, JoinNode node, DataSchema nonEquiEvaluationSchema) {
+ super(context, leftInput, leftSchema, rightInput, node, nonEquiEvaluationSchema);
+ List<Integer> leftKeys = node.getLeftKeys();
+ Preconditions.checkState(!leftKeys.isEmpty(), "Hash join operator requires join keys");
+ _leftKeySelector = KeySelectorFactory.getKeySelector(leftKeys);
+ _rightKeySelector = KeySelectorFactory.getKeySelector(node.getRightKeys());
+ _rightTable = createLookupTable(leftKeys, leftSchema);
+ _matchedRightRows = needUnmatchedRightRows() ? new HashMap<>() : null;
+ }
+
private static LookupTable createLookupTable(List<Integer> joinKeys, DataSchema schema) {
if (joinKeys.size() > 1) {
return new ObjectLookupTable();
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java
index 53d413d1eaa..01ca2314cb3 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/PlanNodeToOpChain.java
@@ -25,6 +25,7 @@ import java.util.function.BiConsumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -41,6 +42,7 @@ import org.apache.pinot.query.planner.plannode.ValueNode;
import org.apache.pinot.query.planner.plannode.WindowNode;
import org.apache.pinot.query.runtime.operator.AggregateOperator;
import org.apache.pinot.query.runtime.operator.AsofJoinOperator;
+import org.apache.pinot.query.runtime.operator.EnrichedHashJoinOperator;
import org.apache.pinot.query.runtime.operator.ErrorOperator;
import org.apache.pinot.query.runtime.operator.FilterOperator;
import org.apache.pinot.query.runtime.operator.HashJoinOperator;
@@ -243,6 +245,39 @@ public class PlanNodeToOpChain {
}
}
+ @Override
+ public MultiStageOperator visitEnrichedJoin(EnrichedJoinNode node, OpChainExecutionContext context) {
+ MultiStageOperator leftOperator = null;
+ MultiStageOperator rightOperator = null;
+ try {
+ List<PlanNode> inputs = node.getInputs();
+ PlanNode left = inputs.get(0);
+ leftOperator = visit(left, context);
+ PlanNode right = inputs.get(1);
+ rightOperator = visit(right, context);
+ JoinNode.JoinStrategy joinStrategy = node.getJoinStrategy();
+ switch (joinStrategy) {
+ case HASH:
+ if (node.getLeftKeys().isEmpty()) {
+ throw new UnsupportedOperationException("NonEquiJoin yet to be supported for EnrichedJoin");
+ } else {
+ return new EnrichedHashJoinOperator(context, leftOperator, left.getDataSchema(), rightOperator, node);
+ }
+ case LOOKUP:
+ throw new UnsupportedOperationException("LookupJoin yet to be supported for EnrichedJoin");
+ case ASOF:
+ throw new UnsupportedOperationException("AsOfJoin yet to be supported for EnrichedJoin");
+ default:
+ throw new IllegalStateException("Unsupported JoinStrategy for EnrichedJoin: " + joinStrategy);
+ }
+ } catch (Exception e) {
+ List<MultiStageOperator> children = Stream.of(leftOperator, rightOperator)
+ .filter(Objects::nonNull)
+ .collect(Collectors.toList());
+ return new ErrorOperator(context, QueryErrorCode.QUERY_EXECUTION, e.getMessage(), children);
+ }
+ }
+
@Override
public MultiStageOperator visitProject(ProjectNode node, OpChainExecutionContext context) {
MultiStageOperator child = null;
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
index 74671dcf882..683cacf6a61 100644
--- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
+++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestVisitor.java
@@ -33,6 +33,7 @@ import org.apache.pinot.common.utils.request.RequestUtils;
import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
import org.apache.pinot.query.parser.CalciteRexExpressionParser;
import org.apache.pinot.query.planner.plannode.AggregateNode;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
import org.apache.pinot.query.planner.plannode.ExchangeNode;
import org.apache.pinot.query.planner.plannode.ExplainedNode;
import org.apache.pinot.query.planner.plannode.FilterNode;
@@ -181,6 +182,71 @@ public class ServerPlanRequestVisitor implements PlanNodeVisitor<Void, ServerPla
return null;
}
+ @Override
+ public Void visitEnrichedJoin(EnrichedJoinNode node, ServerPlanRequestContext context) {
+ // We can reach here for dynamic broadcast SEMI join and lookup join.
+ List<PlanNode> inputs = node.getInputs();
+ PlanNode left = inputs.get(0);
+ PlanNode right = inputs.get(1);
+
+ if (right instanceof MailboxReceiveNode
+ && ((MailboxReceiveNode) right).getExchangeType() == PinotRelExchangeType.PIPELINE_BREAKER) {
+ // For dynamic broadcast SEMI join, right child should be a PIPELINE_BREAKER exchange. Visit the left child and
+ // attach the dynamic filter to the query.
+ if (visit(left, context)) {
+ // semi join to dynamic filter logic
+ PipelineBreakerResult pipelineBreakerResult = context.getPipelineBreakerResult();
+ int resultMapId = pipelineBreakerResult.getNodeIdMap().get(right);
+ List<MseBlock> blocks = pipelineBreakerResult.getResultMap().getOrDefault(resultMapId, Collections.emptyList());
+ List<Object[]> resultDataContainer = new ArrayList<>();
+ DataSchema dataSchema = right.getDataSchema();
+ for (MseBlock block : blocks) {
+ if (block.isData()) {
+ resultDataContainer.addAll(((MseBlock.Data) block).asRowHeap().getRows());
+ }
+ }
+ // TODO: we should keep query stats here as well
+ ServerPlanRequestUtils.attachDynamicFilter(context.getPinotQuery(), node.getLeftKeys(), node.getRightKeys(),
+ resultDataContainer, dataSchema);
+
+ PinotQuery pinotQuery = context.getPinotQuery();
+ for (EnrichedJoinNode.FilterProjectRex rex : node.getFilterProjectRexes()) {
+ if (rex.getType() == EnrichedJoinNode.FilterProjectRexType.FILTER) {
+ // filter logic here
+ if (pinotQuery.getFilterExpression() == null) {
+ Expression expression = CalciteRexExpressionParser.toExpression(rex.getFilter(),
+ pinotQuery.getSelectList());
+ applyTimestampIndex(expression, pinotQuery);
+ pinotQuery.setFilterExpression(expression);
+ } else {
+ // if filter is already applied then it cannot have another one on leaf.
+ context.setLeafStageBoundaryNode(node.getInputs().get(0));
+ }
+ } else {
+ // project logic here
+ List<Expression> selectList =
+ CalciteRexExpressionParser.convertRexNodes(
+ rex.getProjectAndResultSchema().getProject(),
+ pinotQuery.getSelectList());
+ for (Expression expression : selectList) {
+ applyTimestampIndex(expression, pinotQuery);
+ }
+ pinotQuery.setSelectList(selectList);
+ }
+ }
+ }
+ } else {
+ // For lookup join, visit the right child and set it as the leaf boundary.
+ Preconditions.checkState(node.getJoinStrategy() == JoinNode.JoinStrategy.LOOKUP,
+ "Leaf stage should not visit regular JoinNode");
+ if (visit(right, context)) {
+ context.setLeafStageBoundaryNode(right);
+ }
+ }
+
+ return null;
+ }
+
@Override
public Void visitMailboxReceive(MailboxReceiveNode node, ServerPlanRequestContext context) {
throw new UnsupportedOperationException("Leaf stage should not visit MailboxReceiveNode!");
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperatorTest.java
new file mode 100644
index 00000000000..f7da265910d
--- /dev/null
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/EnrichedHashJoinOperatorTest.java
@@ -0,0 +1,652 @@
+/**
+ * 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.pinot.query.runtime.operator;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.query.planner.logical.RexExpression;
+import org.apache.pinot.query.planner.plannode.EnrichedJoinNode;
+import org.apache.pinot.query.planner.plannode.JoinNode;
+import org.apache.pinot.query.planner.plannode.PlanNode;
+import org.apache.pinot.query.runtime.blocks.MseBlock;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+public class EnrichedHashJoinOperatorTest {
+ private AutoCloseable _mocks;
+ private MultiStageOperator _leftInput;
+ private MultiStageOperator _rightInput;
+ private static final DataSchema DEFAULT_CHILD_SCHEMA = new DataSchema(new String[]{"int_col", "string_col"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING});
+
+ @Test
+ public void shouldHandleBasicInnerJoin() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+ HashJoinOperator operator = getBasicOperator(resultSchema, JoinRelType.INNER, List.of(1), List.of(1), List.of());
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 3);
+ assertEquals(resultRows.get(0), new Object[]{1, "Aa", 2, "Aa"});
+ assertEquals(resultRows.get(1), new Object[]{2, "BB", 2, "BB"});
+ assertEquals(resultRows.get(2), new Object[]{2, "BB", 3, "BB"});
+ }
+
+ @Test
+ public void shouldHandleInnerJoinWithTrueFilter() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, RexExpression.Literal.TRUE, null);
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 3);
+ assertEquals(resultRows.get(0), new Object[]{1, "Aa", 2, "Aa"});
+ assertEquals(resultRows.get(1), new Object[]{2, "BB", 2, "BB"});
+ assertEquals(resultRows.get(2), new Object[]{2, "BB", 3, "BB"});
+ }
+
+ @Test
+ public void shouldHandleInnerJoinWithFalseFilter() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, RexExpression.Literal.FALSE, null);
+ assertTrue(operator.nextBlock().isSuccess());
+ }
+
+ @Test
+ public void shouldHandleInnerJoinWithFuncCallFilter() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ RexExpression.FunctionCall startsWith =
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN, SqlKind.STARTS_WITH.name(),
+ List.of(new RexExpression.InputRef(1), new RexExpression.Literal(DataSchema.ColumnDataType.STRING, "B")));
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, startsWith, null);
+
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 2);
+ assertEquals(resultRows.get(0), new Object[]{2, "BB", 2, "BB"});
+ assertEquals(resultRows.get(1), new Object[]{2, "BB", 3, "BB"});
+ }
+
+ // project tests ----
+ @Test
+ public void shouldHandleRefProject() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ DataSchema projectedSchema = new DataSchema(
+ new String[]{"int_col1", "int_col2"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.INT});
+
+ List<RexExpression> projects = List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(2));
+
+ HashJoinOperator operator =
+ getProjectedOperator(DEFAULT_CHILD_SCHEMA, resultSchema, projectedSchema, JoinRelType.INNER, List.of(1),
+ List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, null, projects);
+
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 3);
+ assertEquals(resultRows.get(0), new Object[]{1, 2});
+ assertEquals(resultRows.get(1), new Object[]{2, 2});
+ assertEquals(resultRows.get(2), new Object[]{2, 3});
+ }
+
+ @Test
+ public void shouldHandlePlusMinusTransform() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(3, "BB")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ DataSchema projectSchema = new DataSchema(
+ new String[]{"sum", "diff"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.INT});
+
+ List<RexExpression> operands = List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(2));
+
+ List<RexExpression> projects =
+ List.of(new RexExpression.FunctionCall(DataSchema.ColumnDataType.DOUBLE, SqlKind.PLUS.name(), operands),
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.DOUBLE, SqlKind.MINUS.name(), operands));
+
+ HashJoinOperator operator =
+ getProjectedOperator(DEFAULT_CHILD_SCHEMA, resultSchema, projectSchema, JoinRelType.INNER, List.of(1),
+ List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, null, projects);
+
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 3);
+ assertEquals(resultRows.get(0), new Object[]{3.0, -1.0});
+ assertEquals(resultRows.get(1), new Object[]{4.0, 0.0});
+ assertEquals(resultRows.get(2), new Object[]{5.0, -1.0});
+ }
+
+ @Test
+ public void shouldHandleInnerJoinFilteredProjected() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "BB")
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(2, "Bc")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ // filter
+ RexExpression.FunctionCall startsWith =
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN, SqlKind.STARTS_WITH.name(),
+ List.of(new RexExpression.InputRef(1), new RexExpression.Literal(DataSchema.ColumnDataType.STRING, "B")));
+
+ // project
+ List<RexExpression> projects = List.of(new RexExpression.InputRef(0), new RexExpression.InputRef(2));
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), List.of(),
+ PlanNode.NodeHint.EMPTY, null, startsWith, projects);
+
+ List<Object[]> resultRows = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows();
+ assertEquals(resultRows.size(), 3);
+ }
+
+ @Test
+ public void shouldHandleProjectFilterJoinWithNonEquiConditions() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(4, "Be")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "BB")
+ .addRow(4, "Bd")
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(2, "Bc")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ // NonEquiCondition
+ List<RexExpression> nonEquiConditions = List.of(
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.INT, SqlKind.GREATER_THAN_OR_EQUAL.name(), List.of(
+ new RexExpression.InputRef(0), new RexExpression.InputRef(2)
+ )));
+
+ // filter is before project
+ RexExpression.FunctionCall startsWith =
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.BOOLEAN, SqlKind.STARTS_WITH.name(),
+ List.of(new RexExpression.InputRef(1), new RexExpression.Literal(DataSchema.ColumnDataType.STRING, "Bc")));
+
+ // project
+ List<RexExpression> projects = List.of(new RexExpression.InputRef(0));
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), nonEquiConditions,
+ PlanNode.NodeHint.EMPTY, null, startsWith, projects);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 1);
+ assertEquals(resultRows.get(0), new Object[]{3});
+ }
+
+ @Test
+ public void shouldHandleBasicLimit() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(4, "Be")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "BB")
+ .addRow(4, "Bd")
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(2, "Bc")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ // NonEquiCondition
+ List<RexExpression> nonEquiConditions = List.of(
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.INT, SqlKind.LESS_THAN_OR_EQUAL.name(), List.of(
+ new RexExpression.InputRef(0), new RexExpression.InputRef(2)
+ )));
+
+ int offset = 1;
+ int fetch = 2;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), nonEquiConditions,
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 2);
+ }
+
+ @Test
+ public void shouldHandleOffsetLargerThanResult() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(4, "Be")
+ .addRow(2, "BB")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "BB")
+ .addRow(4, "Bd")
+ .addRow(2, "Aa")
+ .addRow(2, "BB")
+ .addRow(2, "Bc")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ // NonEquiCondition
+ List<RexExpression> nonEquiConditions = List.of(
+ new RexExpression.FunctionCall(DataSchema.ColumnDataType.INT, SqlKind.LESS_THAN_OR_EQUAL.name(), List.of(
+ new RexExpression.InputRef(0), new RexExpression.InputRef(2)
+ )));
+
+ int offset = 3;
+ int fetch = -1;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.INNER, List.of(1), List.of(1), nonEquiConditions,
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 0);
+ }
+
+ @Test
+ public void shouldHandleOffsetLeftJoinOneBlock() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .addRow(4, "Be")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Bc")
+ .addRow(2, "Aa")
+ .addRow(3, "BB")
+ .addRow(2, "BB")
+ .addRow(4, "Bd")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ int offset = 4;
+ int fetch = -1;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.LEFT, List.of(1), List.of(1),
+ Collections.emptyList(),
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 1);
+ assertEquals(resultRows.get(0), new Object[]{4, "Be", null, null});
+ }
+
+ @Test
+ public void shouldHandleOffsetLeftJoinTwoBlocks() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .addRow(4, "Be")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Bc")
+ .addRow(2, "Aa")
+ .addRow(3, "BB")
+ .addRow(2, "BB")
+ .addRow(4, "Bd")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ int offset = 3;
+ int fetch = -1;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.LEFT, List.of(1), List.of(1),
+ Collections.emptyList(),
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 2);
+ assertEquals(resultRows.get(0)[1], "BB");
+ assertEquals(resultRows.get(1), new Object[]{4, "Be", null, null});
+ }
+
+ @Test
+ public void shouldHandleOffsetRightJoinOneBlock() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .addRow(4, "Be")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Bc")
+ .addRow(2, "Aa")
+ .addRow(3, "BB")
+ .addRow(2, "BB")
+ .addRow(4, "Bd")
+ .addRow(5, "Bf")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ int offset = 4;
+ int fetch = -1;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.RIGHT, List.of(1), List.of(1),
+ Collections.emptyList(),
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 2);
+ assertEquals(resultRows.get(0), new Object[]{null, null, 4, "Bd"});
+ assertEquals(resultRows.get(1), new Object[]{null, null, 5, "Bf"});
+ }
+
+ @Test
+ public void shouldHandleOffsetRightJoinTwoBlocks() {
+ _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(3, "Bc")
+ .addRow(1, "Aa")
+ .addRow(2, "BB")
+ .addRow(4, "Be")
+ .buildWithEos();
+
+ _rightInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA)
+ .addRow(2, "Bc")
+ .addRow(2, "Aa")
+ .addRow(3, "BB")
+ .addRow(2, "BB")
+ .addRow(4, "Bd")
+ .addRow(5, "Bf")
+ .buildWithEos();
+ DataSchema resultSchema = new DataSchema(
+ new String[]{"int_col1", "string_col1", "int_col2", "string_col2"},
+ new DataSchema.ColumnDataType[]{
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING,
+ DataSchema.ColumnDataType.INT, DataSchema.ColumnDataType.STRING
+ });
+
+ int offset = 3;
+ int fetch = 2;
+
+ HashJoinOperator operator =
+ getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, JoinRelType.RIGHT, List.of(1), List.of(1),
+ Collections.emptyList(),
+ PlanNode.NodeHint.EMPTY, null, null, null, fetch, offset);
+
+ List<Object[]> resultRows = new ArrayList<>();
+ MseBlock resultBlock = operator.nextBlock();
+ while (!resultBlock.isEos()) {
+ resultRows.addAll(((MseBlock.Data) resultBlock).asRowHeap().getRows());
+ resultBlock = operator.nextBlock();
+ }
+ assertEquals(resultRows.size(), 2);
+ assertEquals(resultRows.get(0)[1], "BB");
+ assertEquals(resultRows.get(1), new Object[]{null, null, 4, "Bd"});
+ }
+
+ // utils ----
+ private EnrichedHashJoinOperator getOperator(DataSchema leftSchema, DataSchema resultSchema, JoinRelType joinType,
+ List<Integer> leftKeys, List<Integer> rightKeys, List<RexExpression> nonEquiConditions,
+ PlanNode.NodeHint nodeHint,
+ RexExpression matchCondition, RexExpression filterCondition, List<RexExpression> projects
+ ) {
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = new ArrayList<>();
+ if (filterCondition != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(filterCondition));
+ }
+ if (projects != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(projects, resultSchema));
+ }
+ return new EnrichedHashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, leftSchema, _rightInput,
+ new EnrichedJoinNode(-1, resultSchema, resultSchema, nodeHint, List.of(), joinType, leftKeys, rightKeys,
+ nonEquiConditions,
+ JoinNode.JoinStrategy.HASH, matchCondition,
+ filterProjectRexes,
+ -1, -1));
+ }
+
+ private EnrichedHashJoinOperator getOperator(DataSchema leftSchema, DataSchema resultSchema, JoinRelType joinType,
+ List<Integer> leftKeys, List<Integer> rightKeys, List<RexExpression> nonEquiConditions,
+ PlanNode.NodeHint nodeHint,
+ RexExpression matchCondition, RexExpression filterCondition, List<RexExpression> projects,
+ int fetch, int offset
+ ) {
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = new ArrayList<>();
+ if (filterCondition != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(filterCondition));
+ }
+ if (projects != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(projects, resultSchema));
+ }
+ return new EnrichedHashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, leftSchema, _rightInput,
+ new EnrichedJoinNode(-1, resultSchema, resultSchema, nodeHint, List.of(), joinType, leftKeys, rightKeys,
+ nonEquiConditions,
+ JoinNode.JoinStrategy.HASH, matchCondition,
+ filterProjectRexes,
+ fetch, offset));
+ }
+
+ private EnrichedHashJoinOperator getProjectedOperator(DataSchema leftSchema, DataSchema joinSchema,
+ DataSchema projectSchema, JoinRelType joinType,
+ List<Integer> leftKeys, List<Integer> rightKeys, List<RexExpression> nonEquiConditions,
+ PlanNode.NodeHint nodeHint,
+ RexExpression matchCondition, RexExpression filterCondition, List<RexExpression> projects
+ ) {
+ List<EnrichedJoinNode.FilterProjectRex> filterProjectRexes = new ArrayList<>();
+ if (filterCondition != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(filterCondition));
+ }
+ if (projects != null) {
+ filterProjectRexes.add(new EnrichedJoinNode.FilterProjectRex(projects, projectSchema));
+ }
+ return new EnrichedHashJoinOperator(OperatorTestUtil.getTracingContext(), _leftInput, leftSchema, _rightInput,
+ new EnrichedJoinNode(-1, joinSchema, projectSchema, nodeHint, List.of(), joinType, leftKeys, rightKeys,
+ nonEquiConditions,
+ JoinNode.JoinStrategy.HASH, matchCondition,
+ filterProjectRexes,
+ -1, -1));
+ }
+
+ private HashJoinOperator getBasicOperator(DataSchema resultSchema, JoinRelType joinType,
+ List<Integer> leftKeys, List<Integer> rightKeys, List<RexExpression> nonEquiConditions) {
+ return getOperator(DEFAULT_CHILD_SCHEMA, resultSchema, joinType, leftKeys, rightKeys, nonEquiConditions,
+ PlanNode.NodeHint.EMPTY, null, null, null);
+ }
+}
diff --git a/pinot-query-runtime/src/test/resources/queries/EnrichedJoin.json b/pinot-query-runtime/src/test/resources/queries/EnrichedJoin.json
new file mode 100644
index 00000000000..fe83ba2c1b7
--- /dev/null
+++ b/pinot-query-runtime/src/test/resources/queries/EnrichedJoin.json
@@ -0,0 +1,93 @@
+{
+ "enriched_join_test": {
+ "tables": {
+ "t1": {
+ "schema": [
+ {"name": "col1", "type": "STRING"},
+ {"name": "col2", "type": "INT"}
+ ],
+ "inputs": [
+ ["foo", 1],
+ ["bar", 2],
+ ["qux", 6],
+ ["qux", 6]
+ ]
+ },
+ "t2": {
+ "schema": [
+ {"name": "col1", "type": "STRING"},
+ {"name": "col2", "type": "INT"}
+ ],
+ "inputs": [
+ ["foo", 3],
+ ["foo", 3],
+ ["bar", 4],
+ ["baz", 5]
+ ]
+ }
+ },
+ "queries": [
+ {
+ "description": "enriched join with one projection",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT {t1}.col2 + {t2}.col2 AS result FROM {t1} JOIN {t2} ON {t1}.col1 = {t2}.col1 ORDER BY result",
+ "outputs": [
+ [4],
+ [4],
+ [6]
+ ]
+ },
+ {
+ "description": "enriched left join with one projection",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT {t1}.col2 + CASE WHEN {t2}.col2 IS NULL THEN 0 ELSE {t2}.col2 END AS result FROM {t1} LEFT JOIN {t2} ON {t1}.col1 = {t2}.col1 ORDER BY result",
+ "outputs": [
+ [4],
+ [4],
+ [6],
+ [6],
+ [6]
+ ]
+ },
+ {
+ "description": "enriched right join with one projection",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT CASE WHEN {t1}.col2 IS NULL THEN 0 ELSE {t1}.col2 END + {t2}.col2 AS result FROM {t1} RIGHT JOIN {t2} ON {t1}.col1 = {t2}.col1 ORDER BY result",
+ "outputs": [
+ [4],
+ [4],
+ [5],
+ [6]
+ ]
+ },
+ {
+ "description": "enriched anti join with limit",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT {t1}.col2 AS result FROM {t1} WHERE NOT EXISTS (SELECT 1 FROM {t2} WHERE {t1}.col1 = {t2}.col1) LIMIT 1",
+ "outputs": [
+ [6]
+ ]
+ },
+ {
+ "description": "enriched join with one projection and limit",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT {t1}.col2 + {t2}.col2 AS result FROM {t1} JOIN {t2} ON {t1}.col1 = {t2}.col1 WHERE {t1}.col2=1 LIMIT 1",
+ "outputs": [
+ [4]
+ ]
+ },
+ {
+ "description": "enriched join with two projects",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 AS result FROM (SELECT {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 AS sum5 FROM {t1} JOIN {t2} ON {t1}.col1 = {t2}.col1) ORDER BY sum5",
+ "outputs": [
+ [220],
+ [220],
+ [330]
+ ]
+ },
+ {
+ "description": "enriched join with two projects and a filter",
+ "sql": "SET usePlannerRules='JoinToEnrichedJoin'; SELECT * FROM (SELECT * FROM (SELECT {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 + {t1}.col2 + {t2}.col2 AS sum5 FROM {t1} JOIN {t2} ON {t1}.col1 = {t2}.col1)) WHERE sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 + sum5 < 300",
+ "outputs": [
+ [20],
+ [20]
+ ]
+ }
+ ]
+ }
+}
\ No newline at end of file
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 82b18210791..bfdb2f7559a 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -809,6 +809,7 @@ public class CommonConstants {
public static final String PRUNE_EMPTY_CORRELATE_RIGHT = "PruneEmptyCorrelateRight";
public static final String PRUNE_EMPTY_JOIN_LEFT = "PruneEmptyJoinLeft";
public static final String PRUNE_EMPTY_JOIN_RIGHT = "PruneEmptyJoinRight";
+ public static final String JOIN_TO_ENRICHED_JOIN = "JoinToEnrichedJoin";
}
/**
@@ -823,7 +824,8 @@ public class CommonConstants {
PlannerRuleNames.AGGREGATE_JOIN_TRANSPOSE_EXTENDED,
PlannerRuleNames.SORT_JOIN_TRANSPOSE,
PlannerRuleNames.SORT_JOIN_COPY,
- PlannerRuleNames.AGGREGATE_UNION_AGGREGATE
+ PlannerRuleNames.AGGREGATE_UNION_AGGREGATE,
+ PlannerRuleNames.JOIN_TO_ENRICHED_JOIN
);
public static class FailureDetector {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org