Merge branch 'main' into unified-allocator-followup

Author: gaurav-amz

sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/WindowFunction.java

@@ -12,9 +12,14 @@
 
 /**
  * Window functions a backend may support. Covers aggregate-as-window
- * (SUM/AVG/COUNT/MIN/MAX over a frame) — these are what PPL {@code eventstats}
- * lowers to — plus ranking functions (ROW_NUMBER) used by PPL {@code dedup}
- * lowering (ROW_NUMBER OVER PARTITION BY ... <= N).
+ * (SUM/AVG/COUNT/MIN/MAX over a frame, what PPL {@code eventstats} lowers to)
+ * plus ROW_NUMBER, used by PPL {@code dedup} (ROW_NUMBER OVER PARTITION BY … <= N)
+ * and emitted by PPL {@code streamstats … by …} as the helper sequence column
+ * {@code __row_number_for_streamstats__}.
+ *
+ * <p>PARTITION BY is allowed: {@code OpenSearchProject}'s cost gate forces SINGLETON
+ * input on any RexOver-bearing project, so the coordinator's {@code WindowAggExec}
+ * sees the full partition regardless of how partition keys span shards.
  *
  * @opensearch.internal
  */
@@ -24,7 +29,6 @@ public enum WindowFunction {
     COUNT(SqlKind.COUNT),
     MIN(SqlKind.MIN),
     MAX(SqlKind.MAX),
-    /** Sequence number per window partition — backs PPL dedup's row-number filter. */
     ROW_NUMBER(SqlKind.ROW_NUMBER);
 
     private final SqlKind sqlKind;

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java

@@ -430,6 +430,13 @@ public Set<JoinCapability> joinCapabilities() {
 
             @Override
             public Set<WindowCapability> windowCapabilities() {
+                // SUM/AVG/COUNT/MIN/MAX cover PPL eventstats; ROW_NUMBER covers PPL dedup
+                // (ROW_NUMBER OVER PARTITION BY … <= N) and the helper sequence column
+                // PPL streamstats … by … emits as __row_number_for_streamstats__.
+                // isthmus's RexExpressionConverter.visitOver serializes the RexOver inline as a
+                // Substrait WindowFunctionInvocation; DataFusion's substrait consumer splits it
+                // into a dedicated LogicalPlan::Window. No adapter or Rust UDF is needed —
+                // row_number is a Substrait-stdlib window function and a DataFusion built-in.
                 return Set.of(
                     new WindowCapability(
                         Set.of(
@@ -438,11 +445,6 @@ public Set<WindowCapability> windowCapabilities() {
                             WindowFunction.COUNT,
                             WindowFunction.MIN,
                             WindowFunction.MAX,
-                            // ROW_NUMBER backs PPL `dedup` lowering (ROW_NUMBER OVER PARTITION BY ... <= N).
-                            // isthmus's RexExpressionConverter.visitOver serializes the RexOver inline as a
-                            // Substrait WindowFunctionInvocation; DataFusion's substrait consumer splits it
-                            // into a dedicated LogicalPlan::Window. No adapter or Rust UDF is needed —
-                            // row_number is a Substrait-stdlib window function and a DataFusion built-in.
                             WindowFunction.ROW_NUMBER
                         ),
                         Set.copyOf(plugin.getSupportedFormats())

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java

@@ -548,6 +548,13 @@ private static Rel replaceInput(Rel wrapper, Rel newInput) {
             return Filter.builder().from(filter).input(newInput).build();
         }
         if (wrapper instanceof Project project) {
+            // Lifted-window shape: OpenSearchProject.liftNestedRexOver emits Project(Project(input))
+            // where the lower Project computes a window column the outer references. A flat swap
+            // here would drop the lower Project and break the outer's input refs.
+            if (project.getInput() instanceof Project lower && containsWindowFunction(lower)) {
+                Rel rewiredLower = replaceInput(lower, newInput);
+                return Project.builder().from(project).input(rewiredLower).build();
+            }
             return Project.builder().from(project).input(newInput).build();
         }
         if (wrapper instanceof Fetch fetch) {
@@ -561,6 +568,15 @@ private static Rel replaceInput(Rel wrapper, Rel newInput) {
         );
     }
 
+    private static boolean containsWindowFunction(Project project) {
+        for (Expression expr : project.getExpressions()) {
+            if (expr instanceof Expression.WindowFunctionInvocation) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Overrides the {@link Expression.AggregationPhase} on every {@link Aggregate.Measure}
      * inside an {@link Aggregate} wrapper. No-op for non-aggregate wrappers.

sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java

@@ -641,4 +641,61 @@ public void testOtherFunctionsNotRenamed() throws Exception {
         assertTrue("must find sum in extension declarations", foundSum);
     }
 
+    /**
+     * Regression: a lifted-window Project wrapper, shaped {@code Project_outer(Project_lower(input))}
+     * with the outer's RexInputRefs pointing into the lower's appended window column, used
+     * to lose its lower layer when re-wired. The next attach-on-top then crashed deserialising
+     * with "Field reference offset (N) must be less than number of fields in struct (N)".
+     */
+    public void testAttachFragmentOnTop_PreservesLiftedWindowProjectLayer() throws Exception {
+        DataFusionFragmentConvertor convertor = newConvertor();
+
+        RelDataType inputRowType = rowType("a");
+        RelNode innerStageScan = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, inputRowType, List.of("datafusion"));
+        byte[] innerBytes = convertor.convertFragment(innerStageScan);
+
+        RelNode placeholderInput = buildTableScan("__placeholder__", "a");
+
+        RelDataType bigintType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true);
+        RexNode rowNumberOver = rexBuilder.makeOver(
+            bigintType,
+            (org.apache.calcite.sql.SqlAggFunction) SqlStdOperatorTable.ROW_NUMBER,
+            List.of(),
+            List.of(),
+            com.google.common.collect.ImmutableList.<org.apache.calcite.rex.RexFieldCollation>of(),
+            org.apache.calcite.rex.RexWindowBounds.UNBOUNDED_PRECEDING,
+            org.apache.calcite.rex.RexWindowBounds.CURRENT_ROW,
+            true,
+            true,
+            false,
+            false,
+            false
+        );
+        RelNode lowerProject = org.apache.calcite.rel.logical.LogicalProject.create(
+            placeholderInput,
+            List.of(),
+            List.of(rexBuilder.makeInputRef(placeholderInput, 0), rowNumberOver),
+            List.of("a", "rn"),
+            java.util.Set.of()
+        );
+
+        RelNode outerProject = org.apache.calcite.rel.logical.LogicalProject.create(
+            lowerProject,
+            List.of(),
+            List.of(rexBuilder.makeInputRef(lowerProject, 0), rexBuilder.makeInputRef(lowerProject, 1)),
+            List.of("a", "rn"),
+            java.util.Set.of()
+        );
+
+        byte[] combined = convertor.attachFragmentOnTop(outerProject, innerBytes);
+
+        Plan plan = decodeSubstrait(combined);
+        Rel root = rootRel(plan);
+        assertTrue("root must be a ProjectRel (outer lift)", root.hasProject());
+        Rel innerOfOuter = root.getProject().getInput();
+        assertTrue("outer's input must remain the lower lift Project", innerOfOuter.hasProject());
+        Rel innerOfLower = innerOfOuter.getProject().getInput();
+        assertTrue("lower's input must be the rewired stage-scan", innerOfLower.hasRead());
+    }
+
 }

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java

@@ -15,9 +15,12 @@
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexFieldCollation;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
 import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.rex.RexWindow;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.analytics.planner.CapabilityRegistry;
@@ -204,6 +207,52 @@ private static RexNode adaptRex(
             if (adapted != operand) operandsChanged = true;
         }
 
+        // Window functions: adapter recursion has to descend into PARTITION BY / ORDER BY
+        // expressions too — they live on RexOver.window, not in getOperands(), and isthmus's
+        // WindowFunctionConverter walks them when emitting substrait. Without this,
+        // calls like SPAN that need a backend-specific rewrite (SPAN(field, n, NULL) →
+        // FLOOR(field/n)*n) survive into substrait emission carrying their NULL-typed
+        // operand and trip TypeConverter.
+        if (call instanceof RexOver over) {
+            RexWindow window = over.getWindow();
+            List<RexNode> adaptedPartitionKeys = new ArrayList<>(window.partitionKeys.size());
+            boolean windowChanged = false;
+            for (RexNode key : window.partitionKeys) {
+                RexNode adapted = adaptRex(key, adapters, fieldStorage, cluster);
+                adaptedPartitionKeys.add(adapted);
+                if (adapted != key) windowChanged = true;
+            }
+            List<RexFieldCollation> adaptedOrderKeys = new ArrayList<>(window.orderKeys.size());
+            for (RexFieldCollation order : window.orderKeys) {
+                RexNode adapted = adaptRex(order.left, adapters, fieldStorage, cluster);
+                if (adapted != order.left) {
+                    adaptedOrderKeys.add(new RexFieldCollation(adapted, order.right));
+                    windowChanged = true;
+                } else {
+                    adaptedOrderKeys.add(order);
+                }
+            }
+            if (operandsChanged || windowChanged) {
+                RexBuilder rexBuilder = cluster.getRexBuilder();
+                return rexBuilder.makeOver(
+                    over.getType(),
+                    over.getAggOperator(),
+                    adaptedOperands,
+                    adaptedPartitionKeys,
+                    com.google.common.collect.ImmutableList.copyOf(adaptedOrderKeys),
+                    window.getLowerBound(),
+                    window.getUpperBound(),
+                    window.getExclude(),
+                    window.isRows(),
+                    true,
+                    false,
+                    over.isDistinct(),
+                    over.ignoreNulls()
+                );
+            }
+            return over;
+        }
+
         RexCall current = operandsChanged ? call.clone(call.getType(), adaptedOperands) : call;
 
         // Look up adapter for this function

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java

@@ -285,13 +285,15 @@ private boolean isOpaqueOperation(String funcName) {
     }
 
     /** Walks project expressions and collects the {@link WindowFunction}s used by any {@link RexOver}.
-     *  PARTITION BY is rejected for aggregate-as-window functions (SUM/AVG/COUNT/MIN/MAX) since
-     *  shuffle exchange isn't wired yet — those run frame-only across all rows on a single fragment.
-     *  ROW_NUMBER is the exception: PPL {@code dedup} lowers to ROW_NUMBER OVER (PARTITION BY ...)
-     *  whose semantics are local to each partition; isthmus + DataFusion's substrait consumer
-     *  emit a Window rel that executes the partition without requiring shuffle (the partition
-     *  is the row-group boundary, not a data-redistribution operator). Unrecognized window
-     *  SqlKinds (LAG, LEAD, NTILE, etc.) also fail here. */
+     *  Unrecognized window SqlKinds (LAG, LEAD, NTILE, etc.) fail here.
+     *
+     *  <p>PARTITION BY and ORDER BY are both accepted — {@code OpenSearchProject}'s cost gate
+     *  already forces SINGLETON input on any RexOver-bearing project, so all rows in a partition
+     *  arrive on the coordinator regardless of whether partition keys span shards. The
+     *  coordinator's WindowAggExec then computes the window correctly per partition / per frame.
+     *  Covers ROW_NUMBER OVER PARTITION BY (PPL dedup), SUM/AVG/COUNT/MIN/MAX OVER PARTITION BY
+     *  (PPL eventstats by ...), and the empty-OVER aggregate-as-window forms. HASH-shuffle is a
+     *  future strict improvement, not a correctness prerequisite. */
     private static Set<WindowFunction> collectWindowFunctions(List<? extends RexNode> exprs) {
         Set<WindowFunction> fns = new LinkedHashSet<>();
         for (RexNode expr : exprs) {
@@ -302,11 +304,6 @@ public RexNode visitOver(RexOver over) {
                     if (fn == null) {
                         throw new IllegalStateException("Window function [" + over.getAggOperator().getName() + "] is not supported");
                     }
-                    if (fn != WindowFunction.ROW_NUMBER && !over.getWindow().partitionKeys.isEmpty()) {
-                        throw new IllegalStateException(
-                            "Window OVER (PARTITION BY ...) is not supported for [" + fn + "] — no shuffle exchange available yet"
-                        );
-                    }
                     fns.add(fn);
                     return super.visitOver(over);
                 }

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java

@@ -138,6 +138,9 @@ protected Set<JoinCapability> joinCapabilities() {
 
     @Override
     protected Set<WindowCapability> windowCapabilities() {
+        // Mirrors DataFusionAnalyticsBackendPlugin.windowCapabilities — SUM/AVG/COUNT/MIN/MAX
+        // for PPL eventstats; ROW_NUMBER backs PPL top/rare/dedup and the streamstats … by …
+        // helper sequence column.
         return Set.of(
             new WindowCapability(
                 Set.of(
@@ -146,8 +149,6 @@ protected Set<WindowCapability> windowCapabilities() {
                     WindowFunction.COUNT,
                     WindowFunction.MIN,
                     WindowFunction.MAX,
-                    // ROW_NUMBER backs PPL `top` / `rare` / `dedup` lowering
-                    // (ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...)).
                     WindowFunction.ROW_NUMBER
                 ),
                 Set.of(PARQUET_DATA_FORMAT)