Onboard PPL dedup to the analytics-engine path (ROW_NUMBER window function) (opensearch-project#21622)

Signed-off-by: Jialiang Liang <jiallian@amazon.com>

Author: RyanL1997

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

@@ -12,10 +12,9 @@
 
 /**
  * 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. PARTITION BY is not currently supported by the planner. Ranking functions
- * (ROW_NUMBER / RANK / DENSE_RANK) are not yet reachable on this route since
- * {@code streamstats} (which lowers to them) isn't wired here.
+ * (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).
  *
  * @opensearch.internal
  */
@@ -24,7 +23,9 @@ public enum WindowFunction {
     AVG(SqlKind.AVG),
     COUNT(SqlKind.COUNT),
     MIN(SqlKind.MIN),
-    MAX(SqlKind.MAX);
+    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

@@ -388,7 +388,19 @@ public Set<JoinCapability> joinCapabilities() {
             public Set<WindowCapability> windowCapabilities() {
                 return Set.of(
                     new WindowCapability(
-                        Set.of(WindowFunction.SUM, WindowFunction.AVG, WindowFunction.COUNT, WindowFunction.MIN, WindowFunction.MAX),
+                        Set.of(
+                            WindowFunction.SUM,
+                            WindowFunction.AVG,
+                            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-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchProjectRule.java

@@ -85,9 +85,10 @@ public void onMatch(RelOptRuleCall call) {
             : childViableBackends;
 
         // Window functions (RexOver): PPL `eventstats` / `appendcol` emit window calls inline
-        // on the projection. Narrow viable backends to those whose WindowCapability declares
-        // every required function. PARTITION BY is rejected up front — no shuffle exchange
-        // exists today.
+        // on the projection; PPL `dedup` lowers to ROW_NUMBER OVER (PARTITION BY ...). Narrow
+        // viable backends to those whose WindowCapability declares every required function.
+        // PARTITION BY is rejected for aggregate-as-window functions (no shuffle exchange
+        // available yet) but allowed for ROW_NUMBER since its partition is local.
         Set<WindowFunction> requiredWindowFns = collectWindowFunctions(project.getProjects());
         if (!requiredWindowFns.isEmpty()) {
             viableBackends = narrowByWindowCapability(viableBackends, requiredWindowFns);
@@ -284,23 +285,28 @@ private boolean isOpaqueOperation(String funcName) {
     }
 
     /** Walks project expressions and collects the {@link WindowFunction}s used by any {@link RexOver}.
-     *  PARTITION BY is rejected — no shuffle exchange exists today. Unrecognized window SqlKinds
-     *  (LAG, LEAD, NTILE, etc.) also fail here. */
+     *  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. */
     private static Set<WindowFunction> collectWindowFunctions(List<? extends RexNode> exprs) {
         Set<WindowFunction> fns = new LinkedHashSet<>();
         for (RexNode expr : exprs) {
             expr.accept(new org.apache.calcite.rex.RexShuttle() {
                 @Override
                 public RexNode visitOver(RexOver over) {
-                    if (!over.getWindow().partitionKeys.isEmpty()) {
-                        throw new IllegalStateException(
-                            "Window OVER (PARTITION BY ...) is not supported — no shuffle exchange available yet"
-                        );
-                    }
                     WindowFunction fn = WindowFunction.fromSqlKind(over.getAggOperator().getKind());
                     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);
                 }