[analytics-engine] Wire PPL eventstats / streamstats commands end-to-end (#21734)

* [analytics-engine] Allow PARTITION BY in window functions; add ROW_NUMBER

Build directly on #21668's framework: cost gate on OpenSearchProject already
forces SINGLETON input on any RexOver-bearing project. With all rows on the
coordinator, WindowAggExec computes per-partition windows correctly regardless
of whether partition keys span shards. HASH-shuffle parallel execution remains
a future strict improvement, not a correctness prerequisite.

Changes:
- OpenSearchProjectRule.collectWindowFunctions: drop the partition-by rejection
  branch. ORDER BY was already unchecked (good: streamstats running window uses
  it).
- WindowFunction enum: add ROW_NUMBER. PPL streamstats … by … emits a helper
  __row_number_for_streamstats__ column via ROW_NUMBER() OVER ROWS UNBOUNDED.
  Other ranking functions (RANK / DENSE_RANK) are not yet emitted by the
  PPL paths exercised by analytics-engine.
- DataFusionAnalyticsBackendPlugin.windowCapabilities: add ROW_NUMBER to the
  declared set so the backend is viable for streamstats … by ….
- WindowPlanShapeTests.testRexOverPartitionBy_rejected → renamed to
  testRexOverPartitionBy_2shard, asserting the cost-gate now produces the
  same SINGLETON-gather-then-project shape as the empty OVER case.

Unblocks PPL eventstats … by <field> and the simpler streamstats forms (no
reset / no global+window+by — those use a self-join lowering that doesn't
flow through RexOver and aren't covered by this PR).

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [QA] Port EventstatsCommandIT + StreamstatsCommandIT from sql repo

One-to-one ports of CalcitePPLEventstatsIT and CalciteStreamstatsCommandIT
from opensearch-project/sql. Every @test method here corresponds to a method
of the same name in the source IT.

Schema is translated to calcs (the only fixture this QA module provisions):
  age     → int0      (11 non-null + 6 nulls — gives the null-bubbling coverage
                       that the sql IT gets from explicit null partition rows)
  country → str0      (3 distinct: FURNITURE, OFFICE SUPPLIES, TECHNOLOGY)
  state   → str3      (e + null)
  name    → key       (unique row id)

Reachable cases assert exact rows. Cases that depend on functionality not yet
wired through the analytics-engine route are expectThrows negative tests
pinning the precise error fragment — failing those is how a future PR observes
that they should be promoted to passing assertions.

EventstatsCommandIT (28 cases mirroring sql IT):
- 12 reachable: testEventstats, *WithNull, *By, *ByWithNull, *ByWithNullBucket,
  *MultipleEventstats(WithNull|WithNullBucket|WithEval), *EmptyRows,
  *Variance, *VarianceWithNull, *VarianceBy, *VarianceWithNullBy
- 16 expectThrows: 7 span()-blocked cases, 1 percentile rejection, 1 dc, 1
  distinct_count, 1 dc-with-null, 1 earliest/latest, 1 variance-by-span,
  3 multiple-partition-by-span variants

StreamstatsCommandIT (47 cases mirroring sql IT):
- ~17 reachable: basic streamstats with/without by, current=false, window=N,
  bigWindow, currentAndWindow, multiple chained, streamstatsAndEventstats,
  streamstatsAndSort, multipleWithEval (1+2), variance + varianceBy
- ~30 expectThrows: span()-blocked cases, dc / distinct_count / earliest /
  latest / percentile not in WindowFunction enum, reset_before/after and
  global+window+by which use self-join lowering not flowing through RexOver,
  leftJoinWithStreamstats / whereInWithStreamstatsSubquery (composed shapes
  not exercised on the route).

Streamstats running aggregates are row-order dependent. Every reachable test
inserts | sort key as the first stage to pin output to a deterministic row
order (calcs has 17 rows; without the sort, parallel scan can return them in
any order, making running counts/avgs non-deterministic).

Both files compile clean against the analytics-engine plugin. No IT actually
runs in this commit — local cluster is unavailable due to a pre-existing
dataformat-native cargo/protoc gap. Marc / verifier should run against a
working cluster to confirm the reachable-vs-throws split matches the comment
predictions.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [analytics-engine] Plan-shape tests for ROW_NUMBER + PARTITION BY ORDER BY

Three additions to WindowPlanShapeTests pin behaviour of the relaxation in
the previous commit:

- testRowNumberOverEmpty_2shard — ROW_NUMBER OVER (), the helper sequence
  column PPL streamstats (with hasGroup=true) emits.
- testSumOverPartitionByOrderBy_2shard — SUM with both partition keys and
  order keys. Calcite's RexOver expansion yields different rendered text
  (no NULLS LAST suffix on default null ordering); the assertion pins the
  actual rendered shape.
- testRowNumberOverPartitionByOrderBy_2shard — combined PARTITION BY +
  ORDER BY for ROW_NUMBER, the streamstats … by … helper-sequence shape.

Also extends MockDataFusionBackend.windowCapabilities to declare
ROW_NUMBER, mirroring the production DataFusionAnalyticsBackendPlugin.
Without it the planner rejects ROW_NUMBER on the mock backend with
"No backend supports window functions [ROW_NUMBER]".

All 23 WindowPlanShapeTests + analytics-framework tests pass.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [analytics-engine] Preserve lifted-window Project layer in fragment rewire

OpenSearchProject.liftNestedRexOver hoists nested RexOver expressions into a
2-layer LogicalProject(LogicalProject(input)) — outer rewrites to RexInputRefs
into the lower's appended window column, lower computes the window. After
convertStandalone the wrapper is Project_outer(Project_lower(input)).

DataFusionFragmentConvertor.replaceInput previously swapped the topmost
Project's direct child wholesale, dropping the lower Project. Subsequent
attachFragmentOnTop calls then crashed substrait deserialisation with
"Field reference offset (N) must be less than number of fields in struct (N)"
because the outer's RexInputRefs pointed past the now-narrower input struct.
Surfaced on multi-shard streamstats … by where SINGLETON-gather injects an
exchange cut that recursively rewires through the lift wrapper. Single-shard
plans avoided it because no rewire happens.

Detect the lift signature (Project whose direct child is a Project containing
a WindowFunctionInvocation) and recurse so newInput replaces the lower
Project's input — both lift layers stay above newInput. Adds a unit-level
regression in DataFusionFragmentConvertorTests and upgrades
StreamstatsCommandIT.testStreamstatsBy_3shard from an error-pinning assertion
to a correctness assertion on the per-partition final running counts.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [QA] testChainedEventstats_3shard — pin lifted-window rewire across multiple stages

Two chained eventstats on a 3-shard index, each producing its own 2-layer
LogicalProject(LogicalProject(input)) lift bundle that flows through a
separate attachFragmentOnTop call on the multi-shard exchange path. Pins
that the lifted-window detection in DataFusionFragmentConvertor.replaceInput
fires correctly per attach call rather than only on the outermost lift bundle.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [QA] Drop redundant *WithNull tests; tighten *WithNullBucket on str3

testEventstatsWithNull / testEventstatsByWithNull and the streamstats peers
were exact copies of their *non*-null variants — calcs has no nulls in str0,
so they exercised nothing the base tests didn't. Removed.

testEventstatsByWithNullBucket / testStreamstatsByWithNullBucket now partition
by str3 (7 nulls + 10 'e' rows) instead of str0 (no nulls). The default and
bucket_nullable=false outputs only diverge when a null-key partition exists,
so the prior asserts proved nothing about the bucket_nullable=false path.
With str3 the null partition's row carries NULL aggregates under
bucket_nullable=false vs concrete values under default, which is what these
tests are supposed to demonstrate.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [QA] Pin span()-in-PARTITION-BY assertions to exact substrait error

The 14 eventstats/streamstats tests that exercise span(...) inside a window
function's PARTITION BY were using assertErrorAny — anything-fails — which
hides regressions. After verifying each query against the cluster, all 14
fail with the same precise error:

  java.lang.UnsupportedOperationException: Unable to convert the type NULL
    at io.substrait.isthmus.TypeConverter.toSubstrait(TypeConverter.java:192)
    ... at WindowFunctionConverter.generateBinding (...)

PPL's span(int0, 10) lowers to SPAN($10, 10, NULL:NULL) — the third operand
is a NULL-typed literal. span() works fine in scalar contexts (stats … by
span(...)) because non-window paths don't visit literal operands the same
way, but isthmus's WindowFunctionConverter walks every operand and trips on
the NULL type.

Replaced assertErrorAny with assertErrorContains pinning the substring
"Unable to convert the type NULL" in all 14 cases, plus updated the
class-level docs in both files. When isthmus or the PPL frontend drops
the NULL operand, these tests will fail loudly so a follow-up can upgrade
them to row assertions.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [analytics-engine] Recurse BackendPlanAdapter into RexOver.window keys

PPL eventstats/streamstats … by span(field, n) lowers to
RexOver(... PARTITION BY SPAN(field, n, NULL:NULL)). Substrait conversion
runs through SpanAdapter, which rewrites SPAN(field, n, NULL) into
FLOOR(field/n)*n so the NULL-typed third operand never reaches isthmus.
But BackendPlanAdapter.adaptRex only walked RexCall.getOperands() — RexOver's
window (partitionKeys / orderKeys) lives on a separate field, so the SPAN
calls inside PARTITION BY were untouched and survived into substrait emission,
where isthmus's WindowFunctionConverter visited every operand including the
NULL literal and threw 'Unable to convert the type NULL'.

Teach adaptRex to recognize RexOver and recurse into its window's
partitionKeys and orderKeys, then rebuild the RexOver via
rexBuilder.makeOver with the adapted window. ORDER BY keys are wrapped in
RexFieldCollation; rebuild the collation only when the inner expression
changed. Same SpanAdapter rewrite path now applies to window-context SPAN
exactly as it does for stats-context SPAN.

Flips all 14 span-in-PARTITION-BY IT cases (7 in EventstatsCommandIT, 7 in
StreamstatsCommandIT) from assertErrorContains('Unable to convert the type
NULL') to assertRowsEqual against the actual per-bucket aggregates.

Signed-off-by: Songkan Tang <songkant@amazon.com>

* [QA] Drop stats-collapse, assert per-row data on slim schema

The eventstats / streamstats IT cases were collapsing 17-row outputs through
| stats max(...) by <col> chains so the assertions stayed compact. With only
17 rows in calcs the collapse hides the broadcast-vs-running shape these
commands are supposed to demonstrate.

Reworked every reachable test (no error / no flaky-by-shard) to:
* drop the | stats collapse,
* select a slim schema via | fields key, <by-col?>, int0, <new eventstats / streamstats cols>,
* sort key for deterministic order,
* assert all 17 (or the post-filter row count) per-row outputs explicitly.

Single-shard streamstats becomes deterministic with `| sort key` so per-row
running aggregates are reproducible. testStreamstatsBy_3shard keeps its
`| stats max(cnt) by str0` collapse — 3-shard streaming order is shard-
dependent and per-row would be flaky.

While in there, two `WithNullBy` companions were partitioning by str0 (no nulls
in calcs) which silently negated the by-key-null coverage they claimed:
* testEventstatsVarianceWithNullBy → switched to `by str3` (7 nulls + 10 'e')
* testStreamstatsVarianceWithNullBy → switched to `by str3`

Added testStreamstatsByWithNull as the default-mode companion to
testStreamstatsByWithNullBucket (eventstats already has the equivalent pair).

Verified locally:
* 29/29 EventstatsCommandIT pass
* 47/47 StreamstatsCommandIT pass
* :sandbox:qa:analytics-engine-rest:precommit green

Signed-off-by: Songkan Tang <songkant@amazon.com>

---------

Signed-off-by: Songkan Tang <songkant@amazon.com>

Author: songkant-aws

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)