Onboard PPL dedup to the analytics-engine path (ROW_NUMBER window function)

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

Author: RyanL1997

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

@@ -47,6 +47,14 @@ default Set<ProjectCapability> projectCapabilities() {
         return Set.of();
     }
 
+    /** Window functions this backend can evaluate when a {@link org.apache.calcite.rex.RexOver}
+     *  appears inside a project. Reported separately from {@link #projectCapabilities()} so
+     *  the planner can pick a window-capable backend without conflating window evaluation with
+     *  ordinary scalar projection. */
+    default Set<WindowCapability> windowCapabilities() {
+        return Set.of();
+    }
+
     /**
      * Delegation types this backend can initiate — it has a custom physical operator
      * that calls Analytics Core's delegation API to offload work to another backend.

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

@@ -0,0 +1,22 @@
+/*
+ * 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.spi;
+
+import java.util.Set;
+
+/**
+ * Declares that a backend can evaluate a specific {@link WindowFunction} as part
+ * of a {@link org.apache.calcite.rex.RexOver} inside a project. Window functions
+ * are independent of {@link FieldType} because Substrait serializes them inline
+ * with their argument expressions — the backend dispatches on the substrait
+ * function reference at execution time, not on the Calcite operator's typed inputs.
+ *
+ * @opensearch.internal
+ */
+public record WindowCapability(WindowFunction function, Set<String> formats) {}

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

@@ -0,0 +1,51 @@
+/*
+ * 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.spi;
+
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Window functions a backend may support — distinct from {@link AggregateFunction}
+ * because {@link org.apache.calcite.rex.RexOver} occurrences are evaluated by the
+ * backend's window operator (Substrait inline {@code WindowFunctionInvocation}),
+ * not as scalar projections or rolled-up aggregates.
+ *
+ * @opensearch.internal
+ */
+public enum WindowFunction {
+    /** Sequence number per window partition — backs PPL dedup's row-number filter. */
+    ROW_NUMBER(SqlStdOperatorTable.ROW_NUMBER);
+
+    private static final Map<SqlOperator, WindowFunction> OPERATOR_INDEX = new HashMap<>();
+
+    static {
+        for (WindowFunction fn : values()) {
+            OPERATOR_INDEX.put(fn.operator, fn);
+        }
+    }
+
+    private final SqlOperator operator;
+
+    WindowFunction(SqlOperator operator) {
+        this.operator = operator;
+    }
+
+    public SqlOperator operator() {
+        return operator;
+    }
+
+    /** Resolves a Calcite {@link SqlOperator} to a {@link WindowFunction}, or {@code null} if unsupported. */
+    public static WindowFunction fromOperator(SqlOperator operator) {
+        return OPERATOR_INDEX.get(operator);
+    }
+}

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

@@ -31,6 +31,8 @@
 import org.opensearch.analytics.spi.ScanCapability;
 import org.opensearch.analytics.spi.SearchExecEngineProvider;
 import org.opensearch.analytics.spi.StdOperatorRewriteAdapter;
+import org.opensearch.analytics.spi.WindowCapability;
+import org.opensearch.analytics.spi.WindowFunction;
 import org.opensearch.be.datafusion.indexfilter.FilterTreeCallbacks;
 import org.opensearch.index.engine.dataformat.DataFormatRegistry;
 
@@ -361,6 +363,17 @@ public Set<ProjectCapability> projectCapabilities() {
                 return Set.copyOf(caps);
             }
 
+            @Override
+            public Set<WindowCapability> windowCapabilities() {
+                Set<String> formats = Set.copyOf(plugin.getSupportedFormats());
+                // ROW_NUMBER inside a project (e.g. PPL dedup's row-number filter) goes through
+                // isthmus's RexExpressionConverter.visitOver → Substrait WindowFunctionInvocation
+                // → DataFusion's Substrait consumer splits it into 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(WindowFunction.ROW_NUMBER, formats));
+            }
+
             @Override
             public Set<AggregateCapability> aggregateCapabilities() {
                 Set<String> formats = Set.copyOf(plugin.getSupportedFormats());

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

@@ -20,6 +20,8 @@
 import org.opensearch.analytics.spi.ProjectCapability;
 import org.opensearch.analytics.spi.ScalarFunction;
 import org.opensearch.analytics.spi.ScanCapability;
+import org.opensearch.analytics.spi.WindowCapability;
+import org.opensearch.analytics.spi.WindowFunction;
 import org.opensearch.cluster.metadata.IndexMetadata;
 
 import java.util.ArrayList;
@@ -62,6 +64,7 @@ public class CapabilityRegistry {
     private final Map<ScalarKey, Map<String, List<String>>> filterIndex = new HashMap<>();
     private final Map<AggregateKey, Map<String, List<String>>> aggregateIndex = new HashMap<>();
     private final Map<ScalarKey, Map<String, List<String>>> scalarIndex = new HashMap<>();
+    private final Map<WindowFunction, Map<String, List<String>>> windowIndex = new HashMap<>();
     // Backends that declared supportsLiteralEvaluation=true for a (function, fieldType)
     private final Map<ScalarKey, List<String>> literalScalarIndex = new HashMap<>();
     // Opaque operations keyed by name (e.g. "painless") rather than a typed key
@@ -80,6 +83,7 @@ public class CapabilityRegistry {
     private final Set<String> filterCapableBackends = new HashSet<>();
     private final Set<String> aggregateCapableBackends = new HashSet<>();
     private final Set<String> projectCapableBackends = new HashSet<>();
+    private final Set<String> windowCapableBackends = new HashSet<>();
 
     public CapabilityRegistry(
         List<AnalyticsSearchBackendPlugin> backends,
@@ -169,6 +173,13 @@ public CapabilityRegistry(
                 }
                 projectCapableBackends.add(name);
             }
+            for (WindowCapability cap : caps.windowCapabilities()) {
+                Map<String, List<String>> formatMap = windowIndex.computeIfAbsent(cap.function(), k -> new HashMap<>());
+                for (String format : cap.formats()) {
+                    formatMap.computeIfAbsent(format, k -> new ArrayList<>()).add(name);
+                }
+                windowCapableBackends.add(name);
+            }
         }
     }
 
@@ -204,6 +215,15 @@ public Set<String> projectCapableBackends() {
         return projectCapableBackends;
     }
 
+    public Set<String> windowCapableBackends() {
+        return windowCapableBackends;
+    }
+
+    /** All backends declaring window-function support for {@code function} ignoring storage formats. */
+    public List<String> windowBackendsAnyFormat(WindowFunction function) {
+        return allBackends(windowIndex.getOrDefault(function, Map.of()));
+    }
+
     // ---- Scan lookups ----
 
     public List<String> scanBackends(Class<? extends ScanCapability> kind, FieldType fieldType, String format) {

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

@@ -15,6 +15,7 @@
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexOver;
 import org.apache.calcite.sql.SqlFunction;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
@@ -27,6 +28,7 @@
 import org.opensearch.analytics.spi.DelegationType;
 import org.opensearch.analytics.spi.FieldType;
 import org.opensearch.analytics.spi.ScalarFunction;
+import org.opensearch.analytics.spi.WindowFunction;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -167,6 +169,17 @@ private RexNode annotateExpr(RexNode expr, List<String> childViableBackends) {
             return expr;
         }
 
+        // Window functions (RexOver) — Calcite emits these from rules like PPL's dedup
+        // (ROW_NUMBER() OVER (PARTITION BY ...)). isthmus's RexExpressionConverter.visitOver
+        // serializes them inline as Substrait WindowFunctionInvocation expressions, and the
+        // consumer (DataFusion) splits them into a dedicated Window plan node. The planner's
+        // job here is only to confine the project to a backend that declared support for the
+        // specific WindowFunction; no scalar-style annotation walk applies because RexOver's
+        // operands are window arguments, not subordinate scalar expressions to resolve.
+        if (rexCall instanceof RexOver rexOver) {
+            return annotateWindow(rexOver, childViableBackends);
+        }
+
         // Baseline operators — arithmetic, CAST, null-handling, conditional — are assumed
         // supported by every backend and are not subject to capability-registry enforcement.
         // Recurse into operands so a non-baseline function nested inside (e.g.
@@ -219,6 +232,27 @@ private RexNode annotateExpr(RexNode expr, List<String> childViableBackends) {
         return new AnnotatedProjectExpression(target.getType(), target, scalarViable, context.nextAnnotationId());
     }
 
+    private RexNode annotateWindow(RexOver rexOver, List<String> childViableBackends) {
+        SqlOperator op = rexOver.getOperator();
+        WindowFunction windowFn = WindowFunction.fromOperator(op);
+        if (windowFn == null) {
+            throw new IllegalStateException("Unsupported window function [" + op.getName() + "]");
+        }
+        List<String> windowCapable = context.getCapabilityRegistry().windowBackendsAnyFormat(windowFn);
+        List<String> viable = new ArrayList<>();
+        for (String candidate : childViableBackends) {
+            if (windowCapable.contains(candidate)) {
+                viable.add(candidate);
+            }
+        }
+        if (viable.isEmpty()) {
+            throw new IllegalStateException(
+                "No backend supports window function [" + windowFn + "] among " + childViableBackends
+            );
+        }
+        return new AnnotatedProjectExpression(rexOver.getType(), rexOver, viable, context.nextAnnotationId());
+    }
+
     private List<String> resolveOpaqueViableBackends(String funcName, List<String> childViableBackends) {
         CapabilityRegistry registry = context.getCapabilityRegistry();
         List<String> viable = registry.opaqueBackendsAnyFormat(funcName);