diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java index 67431e8986b4f..021d264f36091 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java @@ -34,6 +34,7 @@ import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef; import org.opensearch.analytics.planner.rules.ExtractLiteralAggRule; +import org.opensearch.analytics.planner.rules.OpenSearchAggLiteralArgProjectSplitRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateReduceRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateRule; import org.opensearch.analytics.planner.rules.OpenSearchAggregateSplitRule; @@ -53,6 +54,8 @@ import org.opensearch.analytics.planner.rules.OpenSearchUnionSplitRule; import org.opensearch.analytics.planner.rules.OpenSearchValuesRule; +import java.io.PrintWriter; +import java.io.StringWriter; import java.util.List; import java.util.Optional; @@ -101,6 +104,7 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con modifiedRelNode = decomposeAggregates(modifiedRelNode, listener); modifiedRelNode = mark(modifiedRelNode, context, listener); LOGGER.debug("After marking:\n{}", RelOptUtil.toString(modifiedRelNode)); + modifiedRelNode = splitAggLiteralArgProject(modifiedRelNode, listener); // TODO(combine-delegated-predicates): a post-marking HEP rule should fuse same-backend // AND-sibling AnnotatedPredicates into one combined predicate per group, collapsing N // FFM round-trips per RG into one. Blocked on two open design points: @@ -248,6 +252,20 @@ private static RelNode extractLiteralAgg(RelNode input, RuleProfilingListener li return HepPhase.named("literal-agg-extract").addRuleInstance(new ExtractLiteralAggRule()).run(input, listener); } + /** + * Phase 1c': duplicate an aggregate's literal-config-arg Project (e.g. percentile's {@code 50}) + * into a pinned upper copy (literal stays with the aggregate) over an unpinned physical-only + * lower copy (pushes below the ExchangeReducer). Runs AFTER marking — operates on + * {@code OpenSearch*} nodes and emits a pinned {@code OpenSearchProject} whose + * {@code computeSelfCost} forces the CBO-inserted ER below it, keeping the literal in the + * coordinator fragment for the DataFusion substrait converter. Placed after marking so the + * pre-marking {@code PROJECT_MERGE} cannot re-fuse the two copies. See + * {@link OpenSearchAggLiteralArgProjectSplitRule}. + */ + private static RelNode splitAggLiteralArgProject(RelNode input, RuleProfilingListener listener) { + return HepPhase.named("agg-literal-arg-split").addRuleInstance(new OpenSearchAggLiteralArgProjectSplitRule()).run(input, listener); + } + /** * Phase 1a: constant-expression reduction on Filter and Project predicates. Kept in * its own phase so {@code ProjectReduceExpressionsRule} cannot use a downstream @@ -277,15 +295,16 @@ private static RelNode reduceExpressions(RelNode input, RuleProfilingListener li private static RelNode pushdownRules(RelNode input, RuleProfilingListener listener) { return HepPhase.named("pushdown-rules") .bottomUp() - // Transposes (filter-into-* and sort-into-project) cascade together within - // one fixpoint, alongside PROJECT_MERGE which collapses the intermediate - // adjacent Projects that SORT_PROJECT_TRANSPOSE produces. FILTER_MERGE - // runs as its own instruction so it only fires after the transposes have - // settled — that way any auto-injected NOT NULL collapses with the user's - // WHERE on the post-pushdown filter, not on a half-pushed intermediate. - // SORT_PROJECT_TRANSPOSE + PROJECT_MERGE feed the QTF (late-materialization) - // rewriter by lifting Project above Sort so it sees a single Project layer - // above the anchor. + // Transposes (filter-into-*) cascade together within one fixpoint, alongside + // PROJECT_MERGE which collapses adjacent Projects. FILTER_MERGE runs as its own + // instruction so it only fires after the transposes have settled — that way any + // auto-injected NOT NULL collapses with the user's WHERE on the post-pushdown filter, + // not on a half-pushed intermediate. + // + // SORT_PROJECT_TRANSPOSE is intentionally omitted: lifting Project above Sort puts it + // above the Exchange, defeating projection pushdown. Keeping it below lets DataFusion + // prune the scan. QTF relocates the below-Sort Project above its wrapper itself. + // // SORT_REMOVE_REDUNDANT drops a Sort/LIMIT whose input is provably bounded to // within the limit (e.g. a collation-less `head N` or a sort over a scalar // aggregate, getMaxRowCount <= 1): a no-op that the marking rule must not have to @@ -297,7 +316,6 @@ private static RelNode pushdownRules(RelNode input, RuleProfilingListener listen CoreRules.FILTER_PROJECT_TRANSPOSE, CoreRules.FILTER_AGGREGATE_TRANSPOSE, CoreRules.FILTER_INTO_JOIN, - CoreRules.SORT_PROJECT_TRANSPOSE, CoreRules.PROJECT_MERGE, CoreRules.LIMIT_MERGE, CoreRules.SORT_REMOVE_REDUNDANT @@ -390,7 +408,13 @@ private static RelNode cbo(RelNode marked, RelNode rawRelNode, PlannerContext co if (!copied.getTraitSet().equals(desiredTraits)) { volcanoPlanner.setRoot(volcanoPlanner.changeTraits(copied, desiredTraits)); } - return volcanoPlanner.findBestExp(); + RelNode best = volcanoPlanner.findBestExp(); + if (LOGGER.isDebugEnabled()) { + StringWriter sw = new StringWriter(); + volcanoPlanner.dump(new PrintWriter(sw)); + LOGGER.debug("Volcano memo:\n{}", sw); + } + return best; } finally { if (listener != null) listener.endPhase("cbo"); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java index b7b5df0218007..dfedbf3f37309 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java @@ -283,16 +283,15 @@ private static RelNode cutAtExchange( new Stage(childStageId, childFragment, grandchildren, reducer.getExchangeInfo(), childSinkProvider, targetResolver) ); - // Use the reducer's OUTPUT rowType so QTF's appended ___ugsi (set on erRowType by the - // rewriter) flows into the parent stage's partition schema. No-op for non-QTF reducers. - OpenSearchRelNode reducerInput = (OpenSearchRelNode) reducer.getInput(); + // Source both rowType and FSI from the reducer so they stay aligned 1:1 (its input lacks + // QTF's ___ugsi entry). No-op for non-QTF reducers. OpenSearchStageInputScan stageInput = new OpenSearchStageInputScan( reducer.getCluster(), reducer.getTraitSet(), childStageId, reducer.getRowType(), reducer.getViableBackends(), - reducerInput.getOutputFieldStorage() + reducer.getOutputFieldStorage() ); return new OpenSearchExchangeReducer( reducer.getCluster(), diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java index e44445ed7eb54..57dba9731b0da 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java @@ -261,10 +261,6 @@ private static boolean containsEngineNativeAggregate(RelNode root, AggregateMode return false; } - private static boolean isPureReorderProject(org.apache.calcite.rel.core.Project project) { - return project.getProjects().stream().allMatch(e -> e instanceof org.apache.calcite.rex.RexInputRef); - } - /** * Accumulates serialized delegated query bytes during fragment conversion. * @@ -534,14 +530,10 @@ private static byte[] convertReduceNode( } // Single-input operator above the final-fragment boundary — convert child first, then attach. + // A pure-reorder Project above an engine-native-merge FINAL is emitted like any other operator; + // dropping it would strand operators above it (e.g. Sort) with the post-reorder schema over + // un-reordered data, corrupting the final column order. byte[] innerBytes = convertReduceNode(node.getInputs().getFirst(), convertor, false, delegationBytes); - // Skip pure-reorder Project above engine-native-merge FINAL — DataFusion's substrait - // consumer can't bind reordered field names against the FINAL aggregate's state columns. - if (node instanceof org.apache.calcite.rel.core.Project p - && isPureReorderProject(p) - && containsEngineNativeAggregate(node.getInputs().getFirst(), AggregateMode.FINAL)) { - return innerBytes; - } return convertor.attachFragmentOnTop(strippedNode, innerBytes); } throw new IllegalStateException("Unexpected reduce stage node: " + node.getClass().getSimpleName()); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java index 9b7a33a2f4bc3..761f0b14a32c0 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java @@ -43,6 +43,13 @@ public class OpenSearchExchangeReducer extends ConverterImpl implements OpenSear */ private final RelDataType overrideRowType; + /** + * Field storage matching {@link #overrideRowType} 1:1, supplied by QTF alongside the override + * so the {@code ___ugsi} column it declares has a corresponding {@link FieldStorageInfo}. Null + * whenever {@code overrideRowType} is null; non-null only on QTF-rebuilt ERs. + */ + private final List overrideStorage; + /** Convenience constructor — defaults to {@link ExchangeInfo#singleton()}. */ public OpenSearchExchangeReducer(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, List viableBackends) { this(cluster, traitSet, input, viableBackends, ExchangeInfo.singleton(), null); @@ -71,6 +78,23 @@ public OpenSearchExchangeReducer( List viableBackends, ExchangeInfo exchangeInfo, RelDataType overrideRowType + ) { + this(cluster, traitSet, input, viableBackends, exchangeInfo, overrideRowType, null); + } + + /** + * Full constructor. {@code overrideStorage}, when non-null, must align 1:1 with + * {@code overrideRowType} and is returned verbatim by {@link #getOutputFieldStorage()}. QTF + * uses it to pair the {@code ___ugsi} column it declares with a matching {@link FieldStorageInfo}. + */ + public OpenSearchExchangeReducer( + RelOptCluster cluster, + RelTraitSet traitSet, + RelNode input, + List viableBackends, + ExchangeInfo exchangeInfo, + RelDataType overrideRowType, + List overrideStorage ) { // ConverterImpl makes this a Calcite-recognized trait converter — inserted by // Volcano via OpenSearchDistributionTraitDef.convert when a downstream operator @@ -79,6 +103,7 @@ public OpenSearchExchangeReducer( this.viableBackends = viableBackends; this.exchangeInfo = exchangeInfo; this.overrideRowType = overrideRowType; + this.overrideStorage = overrideStorage; } @Override @@ -98,6 +123,11 @@ public ExchangeInfo getExchangeInfo() { @Override public List getOutputFieldStorage() { + // When QTF rebuilds this ER with an overrideRowType (declaring ___ugsi), it supplies the + // matching overrideStorage so rowType and FSI stay aligned 1:1. Otherwise delegate to input. + if (overrideStorage != null) { + return overrideStorage; + } RelNode input = RelNodeUtils.unwrapHep(getInput()); if (input instanceof OpenSearchRelNode openSearchInput) { return openSearchInput.getOutputFieldStorage(); @@ -107,7 +137,15 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends, exchangeInfo, overrideRowType); + return new OpenSearchExchangeReducer( + getCluster(), + traitSet, + sole(inputs), + viableBackends, + exchangeInfo, + overrideRowType, + overrideStorage + ); } /** @@ -122,7 +160,8 @@ public RelNode copy(RelTraitSet traitSet, List inputs) { @Override public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { double rows = mq.getRowCount(getInput()); - return planner.getCostFactory().makeCost(SETUP_COST + rows, SETUP_COST + rows, 0); + double widthFactor = getRowType().getFieldCount(); + return planner.getCostFactory().makeCost(SETUP_COST + rows * widthFactor, SETUP_COST + rows * widthFactor, 0); } @Override @@ -138,7 +177,8 @@ public RelNode copyResolved(String backend, List children, List strippedChildren) { strippedChildren.getFirst(), viableBackends, exchangeInfo, - overrideRowType + overrideRowType, + overrideStorage ); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java index 9df7455d0b467..8821664dadfc6 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java @@ -67,6 +67,16 @@ public class OpenSearchLateMaterialization extends SingleRel implements OpenSear /** Helper columns consumed internally by the Scatter-Gather stage (not in wrapper output). */ public static final Set RESERVED_LATE_MATERIALIZATION_FIELDS = Set.of(ROW_ID_FIELD, UGSI_FIELD); + /** + * Trailing-helper layout contract shared by the QTF rewriter (declares these columns) and the + * execution side (Stitcher / LM-stage drain read them by name). Helpers always TRAIL the + * reduce-set in this order: {@code [..reduce-set.., ___row_id, ___ugsi]}. {@code ___row_id} is + * added on the narrowed Scan; {@code ___ugsi} is appended by the ExchangeReducer. Any operator + * between the Scan and the wrapper must preserve every trailing helper present in its input, in + * this order — never assume a fixed position. + */ + public static final List TRAILING_HELPERS_IN_ORDER = List.of(ROW_ID_FIELD, UGSI_FIELD); + private final List aboveAnchorPhysicalFields; private final List aboveAnchorPhysicalFieldStorage; private final List viableBackends; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java index 6220b16e62909..2776d1d0e4ca7 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java @@ -45,6 +45,16 @@ public class OpenSearchProject extends Project implements OpenSearchRelNode { private final List viableBackends; + /** + * When true, this Project must stay ABOVE the ExchangeReducer (in the coordinator fragment) — + * {@link #computeSelfCost} returns infinite cost unless its input is already gathered + * (SINGLETON/ANY), forcing Volcano to place an ER below it. Used to keep an aggregate's literal + * config arg (e.g. percentile's {@code 50}) adjacent to the aggregate while a duplicate, + * unpinned, physical-only Project pushes below the gather for projection-pushdown. Mirrors the + * RexOver gate, which has the same coordinator-side requirement. + */ + private final boolean pinAboveExchange; + public OpenSearchProject( RelOptCluster cluster, RelTraitSet traitSet, @@ -52,9 +62,22 @@ public OpenSearchProject( List projects, RelDataType rowType, List viableBackends + ) { + this(cluster, traitSet, input, projects, rowType, viableBackends, false); + } + + public OpenSearchProject( + RelOptCluster cluster, + RelTraitSet traitSet, + RelNode input, + List projects, + RelDataType rowType, + List viableBackends, + boolean pinAboveExchange ) { super(cluster, traitSet, List.of(), input, projects, rowType); this.viableBackends = viableBackends; + this.pinAboveExchange = pinAboveExchange; } @Override @@ -62,6 +85,11 @@ public List getViableBackends() { return viableBackends; } + /** See {@link #pinAboveExchange}. */ + public boolean isPinAboveExchange() { + return pinAboveExchange; + } + @Override public List getOutputFieldStorage() { RelNode input = RelNodeUtils.unwrapHep(getInput()); @@ -86,19 +114,21 @@ public List getOutputFieldStorage() { @Override public Project copy(RelTraitSet traitSet, RelNode input, List projects, RelDataType rowType) { - return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends); + return new OpenSearchProject(getCluster(), traitSet, input, projects, rowType, viableBackends, pinAboveExchange); } /** * Projects containing {@code RexOver} (window functions) need fully-gathered input so the - * window's global frame semantics are correct — infinite cost unless input is SINGLETON. - * Volcano picks the plan where an ER sits under this project. + * window's global frame semantics are correct. Projects flagged {@link #pinAboveExchange} must + * likewise stay in the coordinator fragment (they carry an aggregate's literal config arg). Both + * return infinite cost unless input is SINGLETON/ANY — Volcano then picks the plan where an ER + * sits under this project. * - *

Plain projects (no RexOver) have no ordering requirement — tiny cost unconditionally. + *

Plain projects (neither) have no ordering requirement — tiny cost unconditionally. */ @Override public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { - if (!containsOver()) { + if (!containsOver() && !pinAboveExchange) { return planner.getCostFactory().makeTinyCost(); } // containsOver() is Calcite's own — inherited from Project. @@ -143,7 +173,15 @@ public RelNode copyResolved(String backend, List children, List + * Aggregate(percentile_approx($0,$1,$2)) + * OpenSearchProject(ParamPrice=$0, $f1=50, $f2=FLAG) ← UPPER, pinned: literal stays with agg + * [ExchangeReducer inserted here by CBO] + * OpenSearchProject(ParamPrice=$59) ← LOWER, unpinned: physical-only → pushdown + * Scan + * + * + *

The DataFusion converter re-inlines the literal from the directly-attached Project. Without the split, + * width-cost pushes the single literal-bearing Project below the gather, the converter can't reach the + * literal, and percentile errors "must be a literal" / take returns an empty array. The pinned upper copy + * ({@link OpenSearchProject#isPinAboveExchange()}, infinite cost unless input is gathered) keeps the literal + * in the coordinator fragment; the unpinned lower copy narrows the scan. {@code argList} is untouched. + * + *

Runs after marking, before CBO — {@code PROJECT_MERGE} (pre-marking) can't re-fuse the copies, and the + * ER lands between them. Scope = {@code PERCENTILE_APPROX} and {@code TAKE}, the only aggregates that + * materialize a literal config arg as a Project column (FIRST/LAST drop their optional N; others carry no + * literal). Fires only when such a call references a literal column and the Project also has a non-literal + * column to push down. + * + * @opensearch.internal + */ +public class OpenSearchAggLiteralArgProjectSplitRule extends RelOptRule { + + // Aggregates whose literal config arg the DataFusion converter re-inlines from the attached Project. + private static final java.util.Set LITERAL_ARG_AGGS = java.util.Set.of("PERCENTILE_APPROX", "TAKE"); + + public OpenSearchAggLiteralArgProjectSplitRule() { + super(operand(OpenSearchAggregate.class, operand(OpenSearchProject.class, none())), "OpenSearchAggLiteralArgProjectSplitRule"); + } + + @Override + public void onMatch(RelOptRuleCall call) { + OpenSearchAggregate aggregate = call.rel(0); + OpenSearchProject project = call.rel(1); + + // Already split — the upper copy is pinned; don't re-fire on it. + if (project.isPinAboveExchange()) { + return; + } + + List exprs = project.getProjects(); + + // Does any percentile call reference a literal column on this Project? + boolean referencesLiteralArg = false; + for (AggregateCall ac : aggregate.getAggCallList()) { + if (!LITERAL_ARG_AGGS.contains(ac.getAggregation().getName().toUpperCase(java.util.Locale.ROOT))) { + continue; + } + for (int arg : ac.getArgList()) { + if (arg < exprs.size() && exprs.get(arg) instanceof RexLiteral) { + referencesLiteralArg = true; + break; + } + } + } + if (!referencesLiteralArg) { + return; + } + + // Lower Project: keep only the non-literal (RexInputRef-style) columns — these are what must + // cross the gather. Pure literals carry no scan I/O and are re-materialized on the upper copy. + int[] remap = new int[exprs.size()]; + List lowerExprs = new ArrayList<>(); + List origNames = project.getRowType().getFieldNames(); + List lowerNames = new ArrayList<>(); + int next = 0; + for (int i = 0; i < exprs.size(); i++) { + if (exprs.get(i) instanceof RexLiteral) { + remap[i] = -1; + } else { + remap[i] = next++; + lowerExprs.add(exprs.get(i)); + lowerNames.add(origNames.get(i)); + } + } + // Nothing to push down (all columns are literals): the split saves nothing. + if (lowerExprs.isEmpty()) { + return; + } + + RelOptCluster cluster = aggregate.getCluster(); + RexBuilder rexBuilder = cluster.getRexBuilder(); + RelDataTypeFactory typeFactory = cluster.getTypeFactory(); + + RelDataType lowerRowType = buildRowType(typeFactory, project.getRowType(), remap); + OpenSearchProject lower = new OpenSearchProject( + cluster, + project.getTraitSet(), + project.getInput(), + lowerExprs, + lowerRowType, + project.getViableBackends() + ); + + // Upper Project: reconstruct the original row type 1:1. RexInputRefs reindex onto the lower + // Project's output; literals stay inline. The aggregate above references unchanged ordinals. + List upperExprs = new ArrayList<>(exprs.size()); + for (int i = 0; i < exprs.size(); i++) { + RexNode expr = exprs.get(i); + if (expr instanceof RexLiteral) { + upperExprs.add(expr); + } else { + // Non-literal kept on the lower Project at slot remap[i]; reference it by that slot. + int slot = remap[i]; + upperExprs.add(rexBuilder.makeInputRef(lowerRowType.getFieldList().get(slot).getType(), slot)); + } + } + OpenSearchProject upper = new OpenSearchProject( + cluster, + project.getTraitSet(), + lower, + upperExprs, + project.getRowType(), + project.getViableBackends(), + true + ); + + call.transformTo(aggregate.copy(aggregate.getTraitSet(), List.of(upper))); + } + + /** Row type over the columns kept by {@code remap} (entry >= 0), preserving order. */ + private static RelDataType buildRowType(RelDataTypeFactory typeFactory, RelDataType orig, int[] remap) { + RelDataTypeFactory.Builder builder = typeFactory.builder(); + List fields = orig.getFieldList(); + for (int i = 0; i < remap.length; i++) { + if (remap[i] >= 0) { + builder.add(fields.get(i).getName(), fields.get(i).getType()); + } + } + return builder.build(); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java index a839bc2fd6f90..46b2cc8c6ae60 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java @@ -16,6 +16,7 @@ import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexShuttle; import org.apache.calcite.sql.type.SqlTypeName; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -47,8 +48,8 @@ *

    *
  1. Detect ({@link #detect}) — read-only walk that runs the allow-list checks, * computes the {@link Detection} bundle ({@code aboveAnchorPhysicalFields}, - * {@code belowAnchorPhysicalFields}, {@code anchorSlotToPhysicalField}), and applies - * the skip predicate. Returns {@code null} when QTF doesn't apply.
    2. + * {@code belowAnchorPhysicalFields}), and applies the skip predicate. Returns + * {@code null} when QTF doesn't apply. *
  2. Rewrite ({@link #applyRewrite}) — pure transform consuming a {@link Detection}. * Builds the narrowed Scan, walks the below chain, declares {@code ___ugsi} on the * ExchangeReducer, swaps the wrapper in for the anchor, and remaps RexNodes in the @@ -132,7 +133,7 @@ private static Detection detect(RelNode root) { return null; } - BelowChain belowChain = analyzeBelow(anchorCtx.anchor.getInput()); + BelowChain belowChain = analyzeBelow(anchorCtx.anchor); if (belowChain == null) { LOGGER.debug("[QTF] below-anchor allow-list rejected; skipping rewrite"); return null; @@ -157,15 +158,12 @@ private static Detection detect(RelNode root) { return null; } - List anchorSlotToPhysicalField = buildAnchorSlotToPhysicalField(belowChain); - return new Detection( anchorCtx.anchor, anchorCtx.aboveAnchorOperators, belowChain, belowAnchorPhysicalFields, - aboveAnchorPhysicalFields, - anchorSlotToPhysicalField + aboveAnchorPhysicalFields ); } @@ -214,12 +212,12 @@ private static boolean isAboveAllowed(List chain) { * when an op falls outside the below-allow-list, when a below-Project is non-passthrough, * or when there's no Scan at the bottom. */ - private static BelowChain analyzeBelow(RelNode subtree) { + private static BelowChain analyzeBelow(OpenSearchSort anchor) { List chain = new ArrayList<>(); int[] belowProjOutToScan = null; OpenSearchTableScan scan = null; boolean hasExchangeReducer = false; - RelNode n = RelNodeUtils.unwrapHep(subtree); + RelNode n = RelNodeUtils.unwrapHep(anchor.getInput()); while (n != null) { if (n instanceof OpenSearchTableScan s) { scan = s; @@ -232,19 +230,21 @@ private static BelowChain analyzeBelow(RelNode subtree) { LOGGER.debug("[QTF] multiple Projects below anchor — skipping"); return null; } - int[] outToScan = passthroughMap(p); - if (outToScan == null) { - // TODO: derived below-Project lost-opportunity case. Today's algorithm - // declines plans like: - // SELECT description FROM hits ORDER BY UPPER(URL) LIMIT 10 - // → Project(description) ← Sort($1 ASC) ← Project(description, UPPER(URL)) ← Scan - // The derived col (UPPER(URL)) is consumed only by the anchor's collation; we - // *could* push the derived expression above the wrapper, narrow the Scan to - // {URL}, sort below, and fetch {description} for survivors. Skipping for now — - // adding requires threading the derived RexNode into Detection and emitting - // a synthesized above-Project during rewrite. Separate slice. The non-QTF - // path remains correct in the meantime. - LOGGER.debug("[QTF] expression below-Project — skipping (derived-pushup not yet implemented)"); + // A RexOver (window function) below the anchor cannot be relocated above the + // wrapper: window semantics need the full pre-Limit input, but above the wrapper + // only the K survivors remain. Decline (matches the above-anchor RexOver rule). + if (p.containsOver()) { + LOGGER.debug("[QTF] window function in below-Project — skipping (cannot relocate above wrapper)"); + return null; + } + // Per-slot scan mapping: passthrough refs resolve to a scan column index; + // expression slots map to -1 (display-only, reproduced above the wrapper during + // rewrite). A -1 at a slot the anchor's collation references is a DERIVED SORT + // KEY — still declines (sort can't order on a value computed from columns that + // QTF would defer; sorting on the raw column ≠ sorting on the expression). + int[] outToScan = belowProjectSlotMap(p); + if (collationReferencesDerivedSlot(anchor, outToScan)) { + LOGGER.debug("[QTF] anchor sorts on a derived below-Project column — skipping (derived sort key)"); return null; } belowProjOutToScan = outToScan; @@ -257,16 +257,28 @@ private static BelowChain analyzeBelow(RelNode subtree) { return new BelowChain(chain, belowProjOutToScan, scan, hasExchangeReducer); } - /** Returns output→scanIdx map iff every project is a {@link RexInputRef}; else null. */ - private static int[] passthroughMap(OpenSearchProject p) { + /** + * Per-output-slot scan-column index for a below-Project: {@code out[i]} = scan column index + * when slot {@code i} is a passthrough {@link RexInputRef}, else {@code -1} for an expression + * slot (display-only; its physical deps are fetched and the expression is reproduced above + * the wrapper during rewrite). + */ + private static int[] belowProjectSlotMap(OpenSearchProject p) { int[] out = new int[p.getProjects().size()]; for (int i = 0; i < p.getProjects().size(); i++) { - if (!(p.getProjects().get(i) instanceof RexInputRef ref)) return null; - out[i] = ref.getIndex(); + out[i] = (p.getProjects().get(i) instanceof RexInputRef ref) ? ref.getIndex() : -1; } return out; } + /** True iff any anchor collation index lands on an expression ({@code -1}) slot of the below-Project. */ + private static boolean collationReferencesDerivedSlot(OpenSearchSort anchor, int[] belowProjOutToScan) { + for (RelFieldCollation fc : anchor.getCollation().getFieldCollations()) { + if (belowProjOutToScan[fc.getFieldIndex()] < 0) return true; + } + return false; + } + /** * {@code BelowAnchorPhysicalFields} = anchor sort cols ∪ below-filter cols, expressed * as physical (Scan-level) field names. The narrowed Scan's rowType is built from this @@ -312,24 +324,6 @@ private static LinkedHashSet computeAboveAnchorPhysicalFields(List buildAnchorSlotToPhysicalField(BelowChain belowChain) { - // anchor.rowType inherits from belowChain.chain[0]'s rowType (which inherits ER → Filter → Scan). - RelNode topBelowOp = belowChain.chain.isEmpty() ? belowChain.scan : belowChain.chain.get(0); - int slotCount = topBelowOp.getRowType().getFieldCount(); - List scanFields = belowChain.scan.getRowType().getFieldList(); - List out = new ArrayList<>(slotCount); - for (int slot = 0; slot < slotCount; slot++) { - int scanIdx = (belowChain.belowProjOutToScan == null) ? slot : belowChain.belowProjOutToScan[slot]; - out.add(scanFields.get(scanIdx).getName()); - } - return out; - } - // ── Phase 2 — Rewrite ────────────────────────────────────────────── /** @@ -352,11 +346,122 @@ private static RelNode applyRewrite(RelNode root, Detection detection) { // 2d. Wrapper. Output rowType = aboveAnchorPhysicalFields in iteration order. OpenSearchLateMaterialization wrapper = buildWrapper(newAnchor, detection.aboveAnchorPhysicalFields, origScan, detection.anchor); + // 2f. SELECT projection above the wrapper. Without SORT_PROJECT_TRANSPOSE the SELECT-list + // Project sits BELOW the anchor; its display columns (URL, Title, …) are deferred to the + // fetch phase, so the projection — expressions and aliases included — is reproduced ABOVE + // the wrapper with RexInputRefs rebased from scan space to wrapper-output space. Passthrough + // slots of sort-only columns (not fetched, absent from the wrapper output) are dropped; the + // above-chain never references them. When there's no below-Project, the wrapper stands alone. + RelNode aboveWrapper = buildProjectAboveWrapper(wrapper, detection); + // 2e. Above chain: every op's RexInputRefs remapped by column name from its origChild's // rowType to its newChild's rowType. Pass-through ops (Filter, Sort) leak the narrowed // rowType upward, so a single immediate-parent remap is insufficient — every above op // needs the same treatment, recursively. - return rebuildAboveChain(RelNodeUtils.unwrapHep(root), detection.anchor, wrapper); + RelNode rewritten = rebuildAboveChain(RelNodeUtils.unwrapHep(root), detection.anchor, aboveWrapper); + + // 2g. Collapse adjacent Projects at the top. The reproduced below-Project (2f) is a bridge + // exposing the anchor's output schema so the above-chain can graft on by name; when the + // immediate above-op is itself a Project (the SELECT prune), the two stack. Compose them + // into one — substituting the bridge's expressions into the outer Project's refs — so the + // SELECT prune drops to a single Project and no bridge-only column (e.g. a refetched sort + // key the outer doesn't select) leaks through an intermediate node. + return mergeAdjacentTopProjects(rewritten); + } + + /** While the top node is a Project over a Project, compose the two into one. */ + private static RelNode mergeAdjacentTopProjects(RelNode top) { + while (top instanceof OpenSearchProject outer && RelNodeUtils.unwrapHep(outer.getInput()) instanceof OpenSearchProject inner) { + List innerExprs = inner.getProjects(); + RexShuttle substitute = new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + return innerExprs.get(ref.getIndex()); + } + }; + List composed = new ArrayList<>(outer.getProjects().size()); + for (RexNode expr : outer.getProjects()) { + composed.add(expr.accept(substitute)); + } + top = new OpenSearchProject( + outer.getCluster(), + outer.getTraitSet(), + inner.getInput(), + composed, + outer.getRowType(), + outer.getViableBackends() + ); + } + return top; + } + + /** + * Reproduces the below-anchor SELECT Project ABOVE the wrapper. Each below-Project output slot + * becomes one output here, with the SAME name (preserving aliases) and its RexNode rebased from + * scan-column space to wrapper-output (physical-field) space: + *
      + *
    • passthrough ref to a fetched column → ref to that column's wrapper-output index;
      • + *
    • passthrough ref to a sort-only column (not fetched, absent from the wrapper) → dropped;
      • + *
    • expression → same RexCall with operands rebased (all its physical deps are fetched, so + * every operand resolves into the wrapper output).
      • + *
    + * Returns the bare {@code wrapper} when there is no below-Project to reproduce. + */ + private static RelNode buildProjectAboveWrapper(OpenSearchLateMaterialization wrapper, Detection detection) { + OpenSearchProject innerProject = null; + for (RelNode op : detection.belowChain.chain) { + if (op instanceof OpenSearchProject p) { + innerProject = p; + break; + } + } + if (innerProject == null) return wrapper; + + RelDataTypeFactory typeFactory = wrapper.getCluster().getTypeFactory(); + List scanFields = detection.belowChain.scan.getRowType().getFieldList(); + List wrapperFields = wrapper.getRowType().getFieldList(); + + // scan-column index → wrapper-output index, by physical name; -1 when that column was not fetched. + Map wrapperIdxByName = new HashMap<>(wrapperFields.size()); + for (int i = 0; i < wrapperFields.size(); i++) { + wrapperIdxByName.put(wrapperFields.get(i).getName(), i); + } + int[] scanToWrapper = new int[scanFields.size()]; + for (int k = 0; k < scanFields.size(); k++) { + Integer w = wrapperIdxByName.get(scanFields.get(k).getName()); + scanToWrapper[k] = (w == null) ? -1 : w; + } + IndexRemapShuttle toWrapper = new IndexRemapShuttle(scanToWrapper, wrapper.getRowType()); + + List projects = new ArrayList<>(); + List names = new ArrayList<>(); + List innerFields = innerProject.getRowType().getFieldList(); + for (int slot = 0; slot < innerProject.getProjects().size(); slot++) { + RexNode expr = innerProject.getProjects().get(slot); + RexNode rebased; + if (expr instanceof RexInputRef ref) { + int w = scanToWrapper[ref.getIndex()]; + if (w < 0) continue; // sort-only passthrough column — not fetched, drop + rebased = new RexInputRef(w, wrapperFields.get(w).getType()); + } else { + rebased = expr.accept(toWrapper); // expression: all physical deps are fetched + } + projects.add(rebased); + names.add(innerFields.get(slot).getName()); + } + + RelDataTypeFactory.Builder rowType = typeFactory.builder(); + for (int i = 0; i < projects.size(); i++) { + rowType.add(names.get(i), projects.get(i).getType()); + } + return new OpenSearchProject( + wrapper.getCluster(), + wrapper.getTraitSet(), + wrapper, + projects, + rowType.build(), + innerProject.getViableBackends() + ); } // ── 2a. Narrowed Scan ───────────────────────────────────────────── @@ -415,20 +520,25 @@ private static BelowRebuild rebuildBelowChain(BelowChain belowChain, RelNode new } case OpenSearchExchangeReducer er -> { // Invariant 4: declare ___ugsi on the ER's output rowType (materialized at runtime - // by OrdinalAppendingSink before DataFusion's reduce sees the batch). + // by OrdinalAppendingSink before DataFusion's reduce sees the batch). Supply the + // matching FieldStorageInfo so the ER's rowType and FSI stay aligned 1:1 — ___ugsi + // is a derived (coord-side) column with no physical storage. RelDataType erRowType = RelNodeUtils.appendField( typeFactory, rebuilt.getRowType(), OpenSearchLateMaterialization.UGSI_FIELD, typeFactory.createSqlType(SqlTypeName.INTEGER) ); + List erStorage = new ArrayList<>(((OpenSearchRelNode) rebuilt).getOutputFieldStorage()); + erStorage.add(FieldStorageInfo.derivedColumn(OpenSearchLateMaterialization.UGSI_FIELD, SqlTypeName.INTEGER)); rebuilt = new OpenSearchExchangeReducer( er.getCluster(), er.getTraitSet(), rebuilt, er.getViableBackends(), er.getExchangeInfo(), - erRowType + erRowType, + erStorage ); } default -> throw new IllegalStateException("Unexpected below-anchor operator: " + orig.getClass().getSimpleName()); @@ -443,18 +553,18 @@ private static BelowProjectRebuild rebuildBelowProject( int[] scanIdxRemap, RelDataTypeFactory typeFactory ) { - // Below-Project is passthrough; its output→scan map was captured during analyzeBelow, - // but we recompute here from p's projects since each project is a RexInputRef. - int[] origOutToScan = new int[p.getProjects().size()]; - for (int i = 0; i < p.getProjects().size(); i++) { - origOutToScan[i] = ((RexInputRef) p.getProjects().get(i)).getIndex(); - } + // Below-Project output→scan map: passthrough slots give a scan index; expression slots + // give -1 (display-only; reproduced above the wrapper, never rebuilt below). Both kinds + // are dropped from the rebuilt below-Project — only sort/filter passthrough cols that + // survive into the narrowed Scan are kept. + int[] origOutToScan = belowProjectSlotMap(p); List newProjects = new ArrayList<>(); List newNames = new ArrayList<>(); int[] outputRemap = new int[origOutToScan.length]; Arrays.fill(outputRemap, -1); for (int origOut = 0; origOut < origOutToScan.length; origOut++) { + if (origOutToScan[origOut] < 0) continue; // expression slot — pulled up above the wrapper int newScanIdx = scanIdxRemap[origOutToScan[origOut]]; if (newScanIdx < 0) continue; // source dropped (now fetched, not in narrowed Scan) RelDataTypeField field = newChild.getRowType().getFieldList().get(newScanIdx); @@ -462,11 +572,20 @@ private static BelowProjectRebuild rebuildBelowProject( newProjects.add(new RexInputRef(newScanIdx, field.getType())); newNames.add(p.getRowType().getFieldList().get(origOut).getName()); } - // Pass through ___row_id (always last in narrowed Scan rowType). - int rowIdIdx = newChild.getRowType().getFieldCount() - 1; - RelDataTypeField rowIdField = newChild.getRowType().getFieldList().get(rowIdIdx); - newProjects.add(new RexInputRef(rowIdIdx, rowIdField.getType())); - newNames.add(OpenSearchLateMaterialization.ROW_ID_FIELD); + // Pass through every trailing helper present in the child, by NAME, in the shared layout + // order. When this below-Project sits above the ExchangeReducer the child already carries + // ___ugsi (appended by the ER), so a positional "last column" lookup would grab ___ugsi + // instead of ___row_id and drop the helper that the LM-stage drain / Stitcher read by name. + List childFields = newChild.getRowType().getFieldList(); + for (String helper : OpenSearchLateMaterialization.TRAILING_HELPERS_IN_ORDER) { + for (int i = 0; i < childFields.size(); i++) { + if (helper.equals(childFields.get(i).getName())) { + newProjects.add(new RexInputRef(i, childFields.get(i).getType())); + newNames.add(helper); + break; + } + } + } RelDataTypeFactory.Builder pb = typeFactory.builder(); for (int j = 0; j < newProjects.size(); j++) { @@ -553,19 +672,20 @@ private static OpenSearchLateMaterialization buildWrapper( // ── 2e. Above chain rebuild ─────────────────────────────────────── /** - * Walks down the above chain, swapping the anchor's slot with {@code wrapper}, then on - * the way up rewrites every op's RexInputRefs via a by-name remap from {@code origChild}'s - * rowType to {@code newChild}'s rowType. Names are the stable identity that survives - * narrowing — they exist in both rowTypes verbatim for kept columns, and resolve to -1 - * (rejected by {@link IndexRemapShuttle}) for dropped ones. + * Walks down the above chain, swapping the anchor's slot with {@code wrapperOrProject} + * (the wrapper, or the SELECT Project sitting above it), then on the way up rewrites every + * op's RexInputRefs via a by-name remap from {@code origChild}'s rowType to {@code newChild}'s + * rowType. Names are the stable identity that survives narrowing — they exist in both rowTypes + * verbatim for kept columns, and resolve to -1 (rejected by {@link IndexRemapShuttle}) for + * dropped ones. */ - private static RelNode rebuildAboveChain(RelNode current, OpenSearchSort origAnchor, OpenSearchLateMaterialization wrapper) { - if (current == origAnchor) return wrapper; + private static RelNode rebuildAboveChain(RelNode current, OpenSearchSort origAnchor, RelNode wrapperOrProject) { + if (current == origAnchor) return wrapperOrProject; if (current.getInputs().size() != 1) { throw new IllegalStateException("Multi-input parent in QTF chain: " + current.getClass().getSimpleName()); } RelNode origChild = RelNodeUtils.unwrapHep(current.getInput(0)); - RelNode newChild = rebuildAboveChain(origChild, origAnchor, wrapper); + RelNode newChild = rebuildAboveChain(origChild, origAnchor, wrapperOrProject); int[] remap = buildByNameRemap(origChild.getRowType(), newChild.getRowType()); IndexRemapShuttle shuttle = new IndexRemapShuttle(remap, newChild.getRowType()); @@ -621,7 +741,7 @@ private static int[] buildByNameRemap(RelDataType origType, RelDataType newType) * the original plan to recover any of these fields. */ private record Detection(OpenSearchSort anchor, List aboveAnchorOperators, BelowChain belowChain, Set< - String> belowAnchorPhysicalFields, LinkedHashSet aboveAnchorPhysicalFields, List anchorSlotToPhysicalField) { + String> belowAnchorPhysicalFields, LinkedHashSet aboveAnchorPhysicalFields) { } private record AnchorContext(OpenSearchSort anchor, List aboveAnchorOperators) { diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggLiteralArgProjectSplitPlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggLiteralArgProjectSplitPlanShapeTests.java new file mode 100644 index 0000000000000..e67a670b7adb4 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggLiteralArgProjectSplitPlanShapeTests.java @@ -0,0 +1,129 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.planner; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.apache.calcite.util.Optionality; + +import java.math.BigDecimal; +import java.util.List; + +/** + * Plan-shape tests for {@link org.opensearch.analytics.planner.rules.OpenSearchAggLiteralArgProjectSplitRule}. + * + *

    Multi-shard: the literal-bearing Project is duplicated into a pinned upper copy (keeping the literal + * with the aggregate) over an unpinned lower copy (physical-only), with the CBO-inserted ExchangeReducer + * landing between them — so the literal stays coordinator-side while the scan narrows. + * + *

    1-shard: no exchange is needed; the planner leaves a (harmless) stacked Project — DataFusion's physical + * optimizer folds it at execution (verified separately on a live node). + */ +public class AggLiteralArgProjectSplitPlanShapeTests extends PlanShapeTestBase { + + /** {@code Aggregate(($0,$1)) over Project(status=$0, $f1=50) over scan(status,size)}. */ + private RelNode aggOverLiteralProject(String aggName) { + RelNode scan = stubScan(mockTable("test_index", "status", "size")); + RexNode statusRef = rexBuilder.makeInputRef(scan, 0); + RexNode fifty = rexBuilder.makeLiteral(BigDecimal.valueOf(50), typeFactory.createSqlType(SqlTypeName.INTEGER), true); + RelDataType projectType = typeFactory.builder() + .add("status", typeFactory.createSqlType(SqlTypeName.INTEGER)) + .add("$f1", typeFactory.createSqlType(SqlTypeName.INTEGER)) + .build(); + LogicalProject project = (LogicalProject) LogicalProject.create(scan, List.of(), List.of(statusRef, fifty), projectType); + + AggregateCall call = AggregateCall.create( + udaf(aggName), + false, + false, + false, + List.of(), + List.of(0, 1), + -1, + null, + org.apache.calcite.rel.RelCollations.EMPTY, + 0, + project, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + aggName + ); + return LogicalAggregate.create(project, List.of(), ImmutableBitSet.of(), null, List.of(call)); + } + + public void testPercentileLiteralArg_2shard_splitsWithErBetween() { + RelNode result = runPlanner(aggOverLiteralProject("PERCENTILE_APPROX"), multiShardContext()); + // SINGLE percentile over a PINNED upper Project (literal $f1=50) over the ER over the + // physical-only lower Project (status) over the scan. + assertPlanShape( + """ + OpenSearchAggregate(group=[{}], PERCENTILE_APPROX=[PERCENTILE_APPROX($0, $1)], mode=[SINGLE], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], $f1=[50], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchProject(status=[$0], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + """, + result + ); + } + + public void testPercentileLiteralArg_1shard_noExchangeNoSplit() { + RelNode result = runPlanner(aggOverLiteralProject("PERCENTILE_APPROX"), singleShardContext()); + // 1-shard: SINGLETON already satisfied, no ER. The rule still duplicates the Project (the pinned + // upper copy's SINGLETON requirement is trivially met), leaving a stacked Project that DataFusion + // folds at execution. + assertPlanShape(""" + OpenSearchAggregate(group=[{}], PERCENTILE_APPROX=[PERCENTILE_APPROX($0, $1)], mode=[SINGLE], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], $f1=[50], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + """, result); + } + + public void testTakeLiteralArg_2shard_splitsWithErBetween() { + RelNode result = runPlanner(aggOverLiteralProject("TAKE"), multiShardContext()); + assertPlanShape( + """ + OpenSearchAggregate(group=[{}], TAKE=[TAKE($0, $1)], mode=[SINGLE], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], $f1=[50], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchProject(status=[$0], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + """, + result + ); + } + + /** A minimal user-defined agg function named {@code name} (two ANY operands), resolved by the SPI by name. */ + private static SqlAggFunction udaf(String name) { + return new SqlAggFunction( + name, + null, + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0_FORCE_NULLABLE, + null, + OperandTypes.ANY_ANY, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ) { + }; + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java index 9b7c1a31e07fd..d0aa1ed0e2bb0 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java @@ -8,17 +8,28 @@ package org.opensearch.analytics.planner; +import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchFilter; import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchProject; +import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.planner.rel.OpenSearchSort; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.FieldType; import org.opensearch.cluster.ClusterState; import java.util.ArrayList; @@ -78,6 +89,173 @@ public void testQtfFires_withWhere() { ); } + public void testQtfFires_aliasPreserved() { + // SELECT URL AS u, EventDate FROM hits ORDER BY EventDate LIMIT 10 + // Wrapper output is physical-named [URL, EventDate]; the Project above the wrapper + // re-applies the SELECT alias so the result column is `u`, not `URL`. + assertQtfFired( + "SELECT URL AS u, EventDate FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate"), + Expect.outerProjectNames("u", "EventDate") + ); + } + + public void testQtfFires_multiKeySortMixedDirection() { + // SELECT URL, EventDate, CounterID FROM hits ORDER BY EventDate ASC, CounterID DESC LIMIT 10 + // Two collation indices, mixed direction — exercises collationReferencesDerivedSlot over + // multiple keys and the collation rebuild preserving per-key direction. + assertQtfFired( + "SELECT URL, EventDate, CounterID FROM hits ORDER BY EventDate ASC, CounterID DESC LIMIT 10", + 2, + Expect.scanCols("CounterID", "EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate", "CounterID"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate", "CounterID"), + Expect.collationDirections(RelFieldCollation.Direction.ASCENDING, RelFieldCollation.Direction.DESCENDING) + ); + } + + public void testQtfFires_expressionOperandIsSortKey() { + // URL is BOTH the sort key (BelowAnchor → narrowed scan) AND an expression operand + // (refetched into the wrapper, v2 no-passthrough) — exercises the scan→wrapper rebase for a + // col in both roles, and the merged top Project dropping URL so only `combined` surfaces. + // AboveAnchor deps = [Title, URL]; BelowAnchor = {URL}; FetchOnly = {Title} → fire. + assertQtfFired( + "SELECT Title || '-' || URL AS combined FROM hits ORDER BY URL LIMIT 10", + 2, + Expect.scanCols("URL"), + Expect.aboveAnchorPhysicalFields("Title", "URL"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("Title", "URL"), + Expect.outerProjectNames("combined") + ); + } + + public void testQtfFires_expressionOnNeverDisplayedCol() { + // Title appears ONLY inside the expression — never selected alone, not a sort/filter col. + // AboveAnchor = [Title]; BelowAnchor = {EventDate}; FetchOnly = {Title} → fire. + assertQtfFired( + "SELECT UPPER(Title) AS tup FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("Title"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("Title"), + Expect.outerProjectNames("tup") + ); + } + + /** + * Repro of IndexSortPropagationIT.testSortLimit. QTF's post-CBO input (captured) is + * {@code Sort → Project → ExchangeReducer → Filter → Scan} — the below-Project sits ABOVE + * the ER (unlike the simpler {@code Sort → ER → Project → Scan} cases). After rewrite the + * ER carries ___ugsi as its last column, so the below-Project's positional "last = ___row_id" + * lookup grabs ___ugsi instead — surfacing downstream as the DAGBuilder FSI walk throwing + * "RexInputRef[N] has no matching FieldStorageInfo entry". This builds that exact input + * directly (no SQL/CBO/cluster) and asserts the rewritten plan's FSI walk does not throw. + */ + public void testQtfFires_belowProjectAboveExchangeReducer_fieldStorageResolves() { + // Scan: [CounterID, EventDate, URL] all physical (datafusion-backed), matching the IT. + RelOptTable table = mockTable( + "hits", + new String[] { "CounterID", "EventDate", "URL" }, + new SqlTypeName[] { SqlTypeName.INTEGER, SqlTypeName.INTEGER, SqlTypeName.VARCHAR } + ); + List scanStorage = List.of( + physicalFsi("CounterID", "integer", FieldType.INTEGER), + physicalFsi("EventDate", "integer", FieldType.INTEGER), + physicalFsi("URL", "keyword", FieldType.KEYWORD) + ); + List df = List.of("datafusion"); + OpenSearchTableScan scan = new OpenSearchTableScan(cluster, cluster.traitSet(), table, df, scanStorage); + + // Filter CounterID($0) > 0 + RexNode gtZero = rexBuilder.makeCall( + SqlStdOperatorTable.GREATER_THAN, + rexBuilder.makeInputRef(scan, 0), + rexBuilder.makeLiteral(0, typeFactory.createSqlType(SqlTypeName.INTEGER), false) + ); + OpenSearchFilter filter = new OpenSearchFilter(cluster, cluster.traitSet(), scan, gtZero, df); + + // ExchangeReducer (RANDOM → SINGLETON gather) + OpenSearchExchangeReducer er = new OpenSearchExchangeReducer(cluster, cluster.traitSet(), filter, df); + + // Below-Project ABOVE the ER: SELECT URL($2), EventDate($1), CounterID($0) + RelDataType projRowType = typeFactory.builder() + .add("URL", scan.getRowType().getFieldList().get(2).getType()) + .add("EventDate", scan.getRowType().getFieldList().get(1).getType()) + .add("CounterID", scan.getRowType().getFieldList().get(0).getType()) + .build(); + OpenSearchProject project = new OpenSearchProject( + cluster, + cluster.traitSet(), + er, + List.of(rexBuilder.makeInputRef(er, 2), rexBuilder.makeInputRef(er, 1), rexBuilder.makeInputRef(er, 0)), + projRowType, + df + ); + + // Anchor Sort: ORDER BY CounterID($2) DESC, EventDate($1) DESC LIMIT 10 + OpenSearchSort sort = new OpenSearchSort( + cluster, + cluster.traitSet(), + project, + RelCollations.of( + new RelFieldCollation(2, RelFieldCollation.Direction.DESCENDING), + new RelFieldCollation(1, RelFieldCollation.Direction.DESCENDING) + ), + null, + rexBuilder.makeLiteral(10, typeFactory.createSqlType(SqlTypeName.INTEGER), false), + df + ); + + java.util.Optional rewritten = org.opensearch.analytics.planner.rules.OpenSearchLateMaterializationRewriter.rewrite(sort); + assertTrue("QTF must fire for Sort→Project→ER→Filter→Scan", rewritten.isPresent()); + + // Walk FSI over every OpenSearchRelNode, mirroring DAGBuilder.cutAtLateMaterialization — must not throw. + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (node instanceof OpenSearchRelNode os) { + os.getOutputFieldStorage(); + } + super.visit(node, ordinal, parent); + } + }.go(rewritten.get()); + + // The ___ugsi FSI override must survive the copy paths FragmentConversionDriver / DAGBuilder + // use (copy, copyResolved, stripAnnotations) — otherwise rowType (4 cols) and FSI (3) diverge + // again and the reduce-fragment FSI walk throws. Assert each rebuilt ER keeps aligned FSI. + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (node instanceof OpenSearchExchangeReducer er) { + for (RelNode copy : List.of( + er.copy(er.getTraitSet(), er.getInputs()), + er.copyResolved("datafusion", er.getInputs(), List.of()), + er.stripAnnotations(er.getInputs()) + )) { + OpenSearchExchangeReducer c = (OpenSearchExchangeReducer) copy; + assertEquals( + "ER rowType and FSI must stay aligned 1:1 after " + copy, + c.getRowType().getFieldCount(), + c.getOutputFieldStorage().size() + ); + } + } + super.visit(node, ordinal, parent); + } + }.go(rewritten.get()); + } + + private static FieldStorageInfo physicalFsi(String name, String mappingType, FieldType type) { + return new FieldStorageInfo(name, mappingType, type, List.of("parquet"), List.of(), List.of(), false); + } + public void testQtfFires_sortColAlsoProjected() { // SELECT EventDate, URL FROM hits ORDER BY EventDate LIMIT 10 // Wrapper output is in topmost-op (SELECT) order, NOT scan order. @@ -164,7 +342,8 @@ public void testQtfFires_compositeExpressionWithDedup() { Expect.scanCols("EventDate"), Expect.aboveAnchorPhysicalFields("URL"), Expect.erHasUgsi(true), - Expect.wrapperOutput("URL") + Expect.wrapperOutput("URL"), + Expect.outerProjectNames("combined", "upper_url") ); } @@ -178,7 +357,8 @@ public void testQtfFires_compositeExpressionMultiCol() { Expect.scanCols("EventDate"), Expect.aboveAnchorPhysicalFields("URL", "Title"), Expect.erHasUgsi(true), - Expect.wrapperOutput("URL", "Title") + Expect.wrapperOutput("URL", "Title"), + Expect.outerProjectNames("combined") ); } @@ -300,8 +480,8 @@ public void testQtfDeclined_windowInOuterProject() { public void testQtfDeclined_expressionProjectBelowAnchor() { // SELECT URL FROM hits ORDER BY (CounterID + 1) LIMIT 10 - // The sort key (CounterID + 1) gets materialized via a derived below-Project. - // Today's algorithm declines (TODO in passthroughMap — derived-pushup not yet supported). + // Derived sort key → declines (collation references a derived below-Project slot; can't + // sort above the wrapper on it). Display-only expressions DO fire — see compositeExpression*. assertQtfDeclined("SELECT URL FROM hits ORDER BY (CounterID + 1) LIMIT 10", 2); } @@ -475,6 +655,37 @@ void check(Inspector ctx, String sql, String plan) { }; } + /** + * Outer Project (immediately above wrapper) output field names, in order. Asserts the + * SELECT contract — aliases re-applied by the Project even though the wrapper output is + * physical-named. + */ + static Expect outerProjectNames(String... expectedNamesInOrder) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + if (ctx.outerProject == null) { + fail("No outer Project above wrapper.\nSQL: " + sql + "\nPlan:\n" + plan); + return; + } + List actual = fieldNames(ctx.outerProject.getRowType().getFieldList()); + List expected = Arrays.asList(expectedNamesInOrder); + if (!expected.equals(actual)) { + fail( + "Outer Project output names mismatch.\n expected: " + + expected + + "\n actual: " + + actual + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + } + }; + } + /** Anchor Sort collation directions in order. */ static Expect collationDirections(RelFieldCollation.Direction... expected) { return new Expect() { @@ -547,12 +758,12 @@ private static List fieldNames(List fields) { } private RelNode optimize(String sql, int shardCount) { + return optimize(sql, shardCount, List.of(DATAFUSION, LUCENE)); + } + + private RelNode optimize(String sql, int shardCount, List backends) { ClusterState state = SqlPlannerTestFixture.clusterStateWith(ClickBench.INDEX, ClickBench.BASIC_FIELDS, "parquet", shardCount); - PlannerContext context = new PlannerContext( - new CapabilityRegistry(List.of(DATAFUSION, LUCENE), FieldStorageResolver::new), - state, - false - ); + PlannerContext context = new PlannerContext(new CapabilityRegistry(backends, FieldStorageResolver::new), state, false); RelNode parsed = SqlPlannerTestFixture.parseSql(sql, state); return PlannerImpl.runAllOptimizations(parsed, context); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java index 9419cf3595fd1..9a127f1ccc00c 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java @@ -76,6 +76,13 @@ public class MockDataFusionBackend extends MockBackend implements SearchBackEndP AggregateFunction.AVG ); + // STATE_EXPANDING aggregates carrying a literal config arg — exercised by the literal-arg + // Project-split plan-shape tests. Registered via the stateExpanding factory (not simple()). + private static final Set STATE_EXPANDING_AGG_FUNCTIONS = Set.of( + AggregateFunction.PERCENTILE_APPROX, + AggregateFunction.TAKE + ); + private static final Set FILTER_CAPS; static { Set caps = new HashSet<>(); @@ -91,6 +98,9 @@ public class MockDataFusionBackend extends MockBackend implements SearchBackEndP for (AggregateFunction func : AGG_FUNCTIONS) { caps.add(AggregateCapability.simple(func, SUPPORTED_TYPES, DATAFUSION_FORMATS)); } + for (AggregateFunction func : STATE_EXPANDING_AGG_FUNCTIONS) { + caps.add(AggregateCapability.stateExpanding(func, SUPPORTED_TYPES, DATAFUSION_FORMATS)); + } AGG_CAPS = caps; } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java index 2ef707e31ef68..f1a6ba45a84da 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java @@ -55,16 +55,18 @@ public class PlanShapeTests extends PlanShapeTestBase { public void testSortHeadAfterStats_outerSortPreserved() { RelNode input = topKAfterStats(/* withRedundantOuterSort */ true); RelNode result = runPlanner(input, multiShardContext()); - // SORT_PROJECT_TRANSPOSE + PROJECT_MERGE collapse the column-swap Projects; both Sorts - // remap their collation from $0 (post-swap) to $1 (Aggregate's cnt column) and survive. + // No SPT to lift the inner Project above the Sort, so the swap-Projects never become + // adjacent for PROJECT_MERGE to collapse — both survive around the inner Sort. assertPlanShape( """ OpenSearchSort(sort0=[$1], dir0=[ASC], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -77,15 +79,17 @@ public void testSortHeadAfterStats_outerSortPreserved() { public void testSortHeadAfterStats_singleSortFetchPreserved() { RelNode input = topKAfterStats(/* withRedundantOuterSort */ false); RelNode result = runPlanner(input, multiShardContext()); - // Same Project-merge / Sort-transpose as above; the swap-Project pair collapses, - // Sort collation remaps from $0 to $1 (cnt) over the Aggregate output. + // No SPT to lift the inner Project above the Sort, so the swap-Projects never become + // adjacent for PROJECT_MERGE to collapse — both survive around the inner Sort. assertPlanShape( """ - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -100,16 +104,18 @@ public void testSortHeadAfterStats_outerSortWithDifferentKeyKept() { // Inner sort by cnt ($0 below swap), outer sort by k ($0 above swap which maps to k). RelNode input = topKAfterStats(/* withRedundantOuterSort */ true, /* outerSortField */ 0); RelNode result = runPlanner(input, multiShardContext()); - // Project-merge collapses the swap pair; outer Sort sees Aggregate output (k=$0, cnt=$1) - // directly. Outer sort on k stays $0; inner sort on cnt remaps to $1. + // No SPT to lift the inner Project above the Sort, so the swap-Projects never become + // adjacent for PROJECT_MERGE to collapse — both survive around the inner Sort. assertPlanShape( """ OpenSearchSort(sort0=[$0], dir0=[ASC], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -292,17 +298,14 @@ public void testSortThenProjectThenLimit_multiShard() { ); RelNode result = runPlanner(limit, buildContext("parquet", 3, fields)); - // SORT_PROJECT_TRANSPOSE pushes the outer pure-LIMIT Sort below the identity Project, - // producing the QTF-friendly two-Sort shape Project(identity) ← Sort(fetch) ← Sort(coll) ← ER. - // The sort-pushdown rewriter then copies the collated Sort (with the outer fetch) below the ER. + // No SPT: the outer pure-LIMIT Sort stays above the identity Project (not lifted below it). assertPlanShape( """ - OpenSearchProject(name=[$0], score=[$1], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[3], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[3], viableBackends=[[mock-parquet]]) + OpenSearchProject(name=[$0], score=[$1], viableBackends=[[mock-parquet]]) OpenSearchSort(sort0=[$1], dir0=[ASC], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[3], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/RuleProfilingListenerTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/RuleProfilingListenerTests.java index e6b4a7c8bbad0..39fc0a8b3398b 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/RuleProfilingListenerTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/RuleProfilingListenerTests.java @@ -40,6 +40,7 @@ public class RuleProfilingListenerTests extends BasePlannerRulesTests { "pushdown-rules", "aggregate-decompose", "marking", + "agg-literal-arg-split", "cbo" ); @@ -94,6 +95,7 @@ public void testProfileAggregateOverFilterMultiShard() { Map.entry("OpenSearchTableScanRule", 1L), Map.entry("OpenSearchAggregateRule", 1L), Map.entry("OpenSearchAggregateSplitRule", 1L), + Map.entry("OpenSearchAggLiteralArgProjectSplitRule", 0L), Map.entry("OpenSearchDistributionDeriveRule", 3L), Map.entry("ExpandConversionRule", 5L) ) @@ -122,6 +124,8 @@ public void testProfileJoinWithAggregateMultiShard() { 1L, "OpenSearchJoinSplitRule", 1L, + "OpenSearchAggLiteralArgProjectSplitRule", + 0L, "ExpandConversionRule", 2L ) diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java index 485a5168c5312..363f30ad0a624 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java @@ -109,8 +109,8 @@ public void testJoinDag_case1_singleShardSameTable() { OpenSearchJoin(condition=[=($0, $2)], joinType=[left], viableBackends=[[mock-parquet]]) OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, dag @@ -129,8 +129,8 @@ public void testJoinDag_case2_multiShardSameTable() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[1], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON @@ -159,8 +159,8 @@ public void testJoinDag_case3_singleShardDifferentTables() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[left_idx]], viableBackends=[[mock-parquet]]) Stage 1 exchange=SINGLETON - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[right_idx]], viableBackends=[[mock-parquet]]) """, dag @@ -179,8 +179,8 @@ public void testJoinDag_case4_multiShardDifferentTables() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[1], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON @@ -201,10 +201,12 @@ public void testTopKAfterStatsDag_multiShard() { """ QueryDAG(queryId=) Stage 1 - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) + OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[SUM($1)], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) @@ -221,9 +223,11 @@ public void testTopKAfterStatsDag_singleShard() { assertDagShape(""" QueryDAG(queryId=) Stage 0 - OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[SINGLE], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[SINGLE], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, dag); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchAggLiteralArgProjectSplitRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchAggLiteralArgProjectSplitRuleTests.java new file mode 100644 index 0000000000000..cc0c720196313 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/rules/OpenSearchAggLiteralArgProjectSplitRuleTests.java @@ -0,0 +1,256 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.planner.rules; + +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgram; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.apache.calcite.util.Optionality; +import org.opensearch.analytics.planner.BasePlannerRulesTests; +import org.opensearch.analytics.planner.rel.AggregateMode; +import org.opensearch.analytics.planner.rel.OpenSearchAggregate; +import org.opensearch.analytics.planner.rel.OpenSearchProject; +import org.opensearch.analytics.planner.rel.OpenSearchTableScan; + +import java.math.BigDecimal; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Direct-call unit tests for {@link OpenSearchAggLiteralArgProjectSplitRule}. Builds marked + * {@code OpenSearchAggregate -> OpenSearchProject} trees by hand and runs the rule through a + * {@link HepPlanner}, asserting the literal-bearing Project is (or isn't) duplicated into a + * pinned-above-exchange copy over an unpinned physical-only copy. + */ +public class OpenSearchAggLiteralArgProjectSplitRuleTests extends BasePlannerRulesTests { + + private static final List BACKENDS = List.of("mock-parquet"); + + // ── positive cases ───────────────────────────────────────────────────────── + + public void testPercentileLiteralArgProjectIsSplit() { + // Project(field, $f1=50) under Aggregate(percentile_approx($0,$1)). + RelNode result = runRule(buildAggregateOverLiteralProject("PERCENTILE_APPROX", List.of(0, 1))); + assertSplit(result); + } + + public void testTakeLiteralArgProjectIsSplit() { + // Project(field, $f1=10) under Aggregate(take($0,$1)). + RelNode result = runRule(buildAggregateOverLiteralProject("TAKE", List.of(0, 1))); + assertSplit(result); + } + + // ── negative cases ────────────────────────────────────────────────────────── + + public void testNonLiteralArgAggregateNotSplit() { + // SUM is not a literal-config-arg aggregate — even over a Project with a literal column, + // the rule must not fire (SUM doesn't reference the literal as a config arg). + RelNode result = runRule(buildAggregateOverLiteralProject("SUM", List.of(0))); + assertNotSplit(result); + } + + public void testAllLiteralProjectNotSplit() { + // A percentile whose Project has NOTHING but the literal to push down: split buys nothing. + OpenSearchTableScan scan = (OpenSearchTableScan) stubMarkedScan(); + RexLiteral fifty = literal(50); + OpenSearchProject project = new OpenSearchProject( + cluster, + cluster.traitSet(), + scan, + List.of(fifty), + rowType(List.of("$f0"), List.of(SqlTypeName.INTEGER)), + BACKENDS + ); + OpenSearchAggregate agg = aggregate("PERCENTILE_APPROX", List.of(0), project); + assertNotSplit(runRule(agg)); + } + + public void testAlreadyPinnedProjectNotReSplit() { + // Upper (pinned) copy must not re-fire the rule. + OpenSearchTableScan scan = (OpenSearchTableScan) stubMarkedScan(); + RexLiteral fifty = literal(50); + OpenSearchProject pinned = new OpenSearchProject( + cluster, + cluster.traitSet(), + scan, + List.of(rexBuilder.makeInputRef(scan, 0), fifty), + rowType(List.of("field", "$f1"), List.of(SqlTypeName.INTEGER, SqlTypeName.INTEGER)), + BACKENDS, + true // pinAboveExchange + ); + OpenSearchAggregate agg = aggregate("PERCENTILE_APPROX", List.of(0, 1), pinned); + // Already pinned → the rule's guard returns early; the child stays a single Project over the scan. + RelNode result = runRule(agg); + RelNode child = result.getInput(0); + assertTrue("child stays an OpenSearchProject", child instanceof OpenSearchProject); + assertTrue("child's input stays the scan (not a duplicated Project)", child.getInput(0) instanceof OpenSearchTableScan); + } + + // ── assertions ────────────────────────────────────────────────────────────── + + /** Asserts the result is Aggregate -> pinned-upper Project -> unpinned-lower Project, with the + * literal kept only on the upper copy and the lower copy carrying just the physical column. */ + private void assertSplit(RelNode result) { + assertTrue("top should be an aggregate", result instanceof OpenSearchAggregate); + RelNode upper = result.getInput(0); + assertTrue("upper should be an OpenSearchProject", upper instanceof OpenSearchProject); + OpenSearchProject upperProject = (OpenSearchProject) upper; + assertTrue("upper Project must be pinned above the exchange", upperProject.isPinAboveExchange()); + assertTrue("upper Project keeps the literal column", upperProject.getProjects().stream().anyMatch(e -> e instanceof RexLiteral)); + + RelNode lower = upperProject.getInput(); + assertTrue("lower should be an OpenSearchProject", lower instanceof OpenSearchProject); + OpenSearchProject lowerProject = (OpenSearchProject) lower; + assertFalse("lower Project must NOT be pinned", lowerProject.isPinAboveExchange()); + assertTrue( + "lower Project must carry only physical (non-literal) columns", + lowerProject.getProjects().stream().noneMatch(e -> e instanceof RexLiteral) + ); + assertFalse("lower Project must keep at least one column to push down", lowerProject.getProjects().isEmpty()); + } + + /** Asserts the rule did NOT split: the aggregate's child is a single Project directly over the scan + * (no duplicated Project-over-Project stack, nothing pinned). */ + private void assertNotSplit(RelNode result) { + assertTrue("top should be an aggregate", result instanceof OpenSearchAggregate); + RelNode child = result.getInput(0); + assertTrue("child should be an OpenSearchProject", child instanceof OpenSearchProject); + assertFalse("child Project must not be pinned", ((OpenSearchProject) child).isPinAboveExchange()); + assertTrue("child Project's input must be the scan, not a duplicated Project", child.getInput(0) instanceof OpenSearchTableScan); + } + + // ── builders ────────────────────────────────────────────────────────────── + + /** Aggregate({@code aggName}(argList)) over {@code Project(field=$0, $f1=50)}. */ + private OpenSearchAggregate buildAggregateOverLiteralProject(String aggName, List argList) { + OpenSearchTableScan scan = (OpenSearchTableScan) stubMarkedScan(); + RexNode fieldRef = rexBuilder.makeInputRef(scan, 0); + RexLiteral literal = literal(50); + OpenSearchProject project = new OpenSearchProject( + cluster, + cluster.traitSet(), + scan, + List.of(fieldRef, literal), + rowType(List.of("field", "$f1"), List.of(SqlTypeName.INTEGER, SqlTypeName.INTEGER)), + BACKENDS + ); + return aggregate(aggName, argList, project); + } + + private OpenSearchAggregate aggregate(String aggName, List argList, OpenSearchProject input) { + AggregateCall call = "SUM".equals(aggName) + ? AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + false, + false, + List.of(), + argList, + -1, + null, + RelCollations.EMPTY, + 0, + input, + null, + aggName + ) + : AggregateCall.create( + udaf(aggName), + false, + false, + false, + List.of(), + argList, + -1, + null, + RelCollations.EMPTY, + 0, + input, + // Matches the UDAF's ReturnTypes.ARG0_FORCE_NULLABLE (arg 0 = field, INTEGER). + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + aggName + ); + return new OpenSearchAggregate( + cluster, + cluster.traitSet(), + input, + ImmutableBitSet.of(), + null, + List.of(call), + AggregateMode.SINGLE, + BACKENDS, + Map.of(), + Map.of(), + Collections.singletonList(null) + ); + } + + /** A minimal user-defined agg function with the given name (two ANY operands). */ + private static SqlAggFunction udaf(String name) { + return new SqlAggFunction( + name, + null, + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0_FORCE_NULLABLE, + null, + OperandTypes.ANY_ANY, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ) { + }; + } + + private RexLiteral literal(int value) { + return (RexLiteral) rexBuilder.makeLiteral(BigDecimal.valueOf(value), typeFactory.createSqlType(SqlTypeName.INTEGER), true); + } + + private RelDataType rowType(List names, List types) { + RelDataTypeFactory.Builder b = typeFactory.builder(); + for (int i = 0; i < names.size(); i++) { + b.add(names.get(i), typeFactory.createSqlType(types.get(i))); + } + return b.build(); + } + + /** A marked OpenSearchTableScan (single INTEGER column "field") usable as a rule leaf. */ + private RelNode stubMarkedScan() { + return new OpenSearchTableScan( + cluster, + cluster.traitSet(), + mockTable("test_index", new String[] { "field" }, new SqlTypeName[] { SqlTypeName.INTEGER }), + BACKENDS, + List.of() + ); + } + + private RelNode runRule(RelNode input) { + HepProgram program = new HepProgramBuilder().addRuleInstance(new OpenSearchAggLiteralArgProjectSplitRule()).build(); + HepPlanner planner = new HepPlanner(program); + planner.setRoot(input); + return planner.findBestExp(); + } +}