[analytics-engine] Per-shard bucket oversampling for TopK aggregation queries

sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs

@@ -158,6 +158,7 @@ pub mod str_to_date;
 pub mod strftime;
 pub mod time_format;
 pub mod width_bucket;
+pub mod reduce_eval;
 
 // Dev note: if a freshly added UDF here fails at runtime with
 // "Unsupported function name: <X>" despite the Java side being wired, the
@@ -203,6 +204,7 @@ pub fn register_all(ctx: &SessionContext) {
     strftime::register_all(ctx);
     time_format::register_all(ctx);
     width_bucket::register_all(ctx);
+    reduce_eval::register_all(ctx);
     log::info!(
         "OpenSearch UDF register_all: convert_tz, conversion(numeric_conversion: num/auto/memk/rmcomma/rmunit/dur2sec/mstime, time_conversion: ctime/mktime), crc32, date_format, extract, from_unixtime, item, json_append, json_array_length, json_delete, json_extend, json_extract, json_extract_all, json_keys, json_set, makedate, maketime, minspan_bucket, mvappend, mvfind, mvzip, parse, range_bucket, rex_extract, rex_extract_multi, rex_offset, sha1, span_bucket, str_to_date, strftime, time_format, width_bucket registered"
     );

sandbox/plugins/analytics-backend-datafusion/rust/src/udf/reduce_eval.rs

@@ -0,0 +1,137 @@
+/*
+ * 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.
+ */
+
+//! `reduce_eval(agg_name, state)` — evaluates an opaque aggregate's partial state
+//! to produce a sortable scalar. Used by the TopK rewriter's reduce Project.
+
+use std::sync::Arc;
+
+use datafusion::arrow::array::{Array, ArrayRef, BinaryArray, UInt64Array};
+use datafusion::arrow::datatypes::{DataType, Field};
+use datafusion::common::{DataFusionError, Result, ScalarValue};
+use datafusion::execution::context::SessionContext;
+use datafusion::functions_aggregate::approx_distinct::approx_distinct_udaf;
+use datafusion::logical_expr::function::AccumulatorArgs;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility};
+use datafusion::physical_expr::expressions::Column;
+
+pub fn register_all(ctx: &SessionContext) {
+    ctx.register_udf(ScalarUDF::new_from_impl(ReduceEvalUdf));
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct ReduceEvalUdf;
+
+impl ScalarUDFImpl for ReduceEvalUdf {
+    fn as_any(&self) -> &dyn std::any::Any { self }
+    fn name(&self) -> &str { "reduce_eval" }
+
+    fn signature(&self) -> &Signature {
+        static SIG: std::sync::LazyLock<Signature> = std::sync::LazyLock::new(|| {
+            Signature::any(2, Volatility::Immutable)
+        });
+        &SIG
+    }
+
+    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+        Ok(DataType::UInt64)
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let agg_name = match &args.args[0] {
+            ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) => s.clone(),
+            _ => return Err(DataFusionError::Execution("reduce_eval: first arg must be literal agg name".into())),
+        };
+        let state_col = match &args.args[1] {
+            ColumnarValue::Array(a) => a.clone(),
+            ColumnarValue::Scalar(s) => s.to_array_of_size(args.number_rows)?,
+        };
+
+        match agg_name.as_str() {
+            "approx_distinct" => eval_approx_distinct(&state_col),
+            other => Err(DataFusionError::Execution(format!("reduce_eval: unsupported aggregate '{other}'"))),
+        }
+    }
+}
+
+fn eval_approx_distinct(state_col: &ArrayRef) -> Result<ColumnarValue> {
+    let binary = state_col.as_any().downcast_ref::<BinaryArray>()
+        .ok_or_else(|| DataFusionError::Execution("reduce_eval(approx_distinct): expected Binary state".into()))?;
+
+    let field: Arc<Field> = Arc::new(Field::new("x", DataType::Int64, true));
+    let schema = Arc::new(datafusion::arrow::datatypes::Schema::new(vec![field.as_ref().clone()]));
+    let expr: Arc<dyn datafusion::physical_plan::PhysicalExpr> = Arc::new(Column::new("x", 0));
+    let ret_field: Arc<Field> = Arc::new(Field::new("r", DataType::UInt64, true));
+
+    let mut results = Vec::with_capacity(binary.len());
+    for i in 0..binary.len() {
+        if binary.is_null(i) {
+            results.push(0u64);
+            continue;
+        }
+        let mut acc = approx_distinct_udaf().accumulator(AccumulatorArgs {
+            return_field: ret_field.clone(),
+            schema: &schema,
+            ignore_nulls: false,
+            order_bys: &[],
+            name: "x",
+            is_distinct: false,
+            exprs: &[expr.clone()],
+            expr_fields: &[field.clone()],
+            is_reversed: false,
+        })?;
+        let state_array: ArrayRef = Arc::new(BinaryArray::from(vec![binary.value(i)]));
+        acc.merge_batch(&[state_array])?;
+        match acc.evaluate()? {
+            ScalarValue::UInt64(Some(v)) => results.push(v),
+            _ => results.push(0),
+        }
+    }
+    Ok(ColumnarValue::Array(Arc::new(UInt64Array::from(results))))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::arrow::array::BinaryArray;
+    use datafusion::functions_aggregate::approx_distinct::approx_distinct_udaf;
+    use datafusion::logical_expr::Accumulator;
+
+    #[test]
+    fn test_reduce_eval_approx_distinct() {
+        // Build an HLL state by updating an accumulator
+        let field: Arc<Field> = Arc::new(Field::new("x", DataType::Int64, true));
+        let schema = Arc::new(datafusion::arrow::datatypes::Schema::new(vec![field.as_ref().clone()]));
+        let expr: Arc<dyn datafusion::physical_plan::PhysicalExpr> = Arc::new(Column::new("x", 0));
+        let ret_field: Arc<Field> = Arc::new(Field::new("r", DataType::UInt64, true));
+        let mut acc = approx_distinct_udaf().accumulator(AccumulatorArgs {
+            return_field: ret_field.clone(), schema: &schema, ignore_nulls: false,
+            order_bys: &[], name: "x", is_distinct: false,
+            exprs: &[expr.clone()], expr_fields: &[field.clone()], is_reversed: false,
+        }).unwrap();
+
+        // Feed some values
+        let values: ArrayRef = Arc::new(datafusion::arrow::array::Int64Array::from(vec![1, 2, 3, 4, 5]));
+        acc.update_batch(&[values]).unwrap();
+        let state = acc.state().unwrap();
+
+        // Extract the Binary state
+        let state_array = state[0].to_array_of_size(1).unwrap();
+        let binary = state_array.as_any().downcast_ref::<BinaryArray>().unwrap();
+
+        // Run reduce_eval
+        let result = eval_approx_distinct(&(Arc::new(binary.clone()) as ArrayRef)).unwrap();
+        match result {
+            ColumnarValue::Array(arr) => {
+                let uint_arr = arr.as_any().downcast_ref::<UInt64Array>().unwrap();
+                assert_eq!(uint_arr.value(0), 5); // 5 distinct values
+            }
+            _ => panic!("expected Array"),
+        }
+    }
+}

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

@@ -106,8 +106,19 @@ public class DataFusionFragmentConvertor implements FragmentConvertor {
         SqlFunctionCategory.USER_DEFINED_FUNCTION
     );
 
+    /** TopK reduce expression: evaluates opaque aggregate state to a sortable scalar. */
+    static final SqlOperator REDUCE_EVAL_OP = new SqlFunction(
+        "reduce_eval",
+        SqlKind.OTHER_FUNCTION,
+        ReturnTypes.BIGINT_NULLABLE,
+        null,
+        OperandTypes.ANY_ANY,
+        SqlFunctionCategory.USER_DEFINED_FUNCTION
+    );
+
     private static final List<FunctionMappings.Sig> ADDITIONAL_SCALAR_SIGS = List.of(
         FunctionMappings.s(DelegatedPredicateFunction.FUNCTION, DelegatedPredicateFunction.NAME),
+        FunctionMappings.s(REDUCE_EVAL_OP, "reduce_eval"),
         FunctionMappings.s(DelegationPossibleFunction.FUNCTION, DelegationPossibleFunction.NAME),
         FunctionMappings.s(SqlStdOperatorTable.ASCII, "ascii"),
         FunctionMappings.s(SqlStdOperatorTable.CHAR_LENGTH, "length"),

sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_scalar_functions.yaml

@@ -1190,3 +1190,12 @@ scalar_functions:
           - { name: format, value: "varchar<L2>" }
         nullability: DECLARED_OUTPUT
         return: fp64
+  - name: "reduce_eval"
+    description: "Evaluates opaque aggregate partial state to a sortable scalar for TopK."
+    impls:
+      - args:
+          - name: agg_name
+            value: string
+          - name: state
+            value: any1
+        return: i64

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

@@ -171,7 +171,11 @@ public Collection<Module> createGuiceModules() {
 
     @Override
     public List<Setting<?>> getSettings() {
-        return List.of(COORDINATOR_BUFFER_LIMIT, ReaderContextStore.READER_CONTEXT_KEEP_ALIVE);
+        List<Setting<?>> settings = new java.util.ArrayList<>();
+        settings.add(COORDINATOR_BUFFER_LIMIT);
+        settings.add(ReaderContextStore.READER_CONTEXT_KEEP_ALIVE);
+        settings.addAll(org.opensearch.analytics.settings.AnalyticsApproximationSettings.all());
+        return List.copyOf(settings);
     }
 
     @Override

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

@@ -41,6 +41,7 @@
 import org.opensearch.analytics.planner.rules.OpenSearchSortRule;
 import org.opensearch.analytics.planner.rules.OpenSearchSortSplitRule;
 import org.opensearch.analytics.planner.rules.OpenSearchTableScanRule;
+import org.opensearch.analytics.planner.rules.OpenSearchTopKRewriter;
 import org.opensearch.analytics.planner.rules.OpenSearchUnionRule;
 import org.opensearch.analytics.planner.rules.OpenSearchUnionSplitRule;
 import org.opensearch.analytics.planner.rules.OpenSearchValuesRule;
@@ -108,6 +109,11 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con
             modifiedRelNode = lateMat.get();
             LOGGER.info("After late-materialization:\n{}", RelOptUtil.toString(modifiedRelNode));
         }
+        Optional<RelNode> topK = OpenSearchTopKRewriter.rewrite(modifiedRelNode, context);
+        if (topK.isPresent()) {
+            modifiedRelNode = topK.get();
+            LOGGER.info("After TopK rewrite:\n{}", RelOptUtil.toString(modifiedRelNode));
+        }
 
         if (listener != null) {
             RuleProfilingListener.PlannerProfile profile = listener.snapshot();
@@ -288,6 +294,7 @@ private static RelNode mark(RelNode input, PlannerContext context, RuleProfiling
     }
 
     /** Phase 2: VolcanoPlanner for trait propagation + exchange insertion. */
+
     private static RelNode cbo(RelNode marked, RelNode rawRelNode, PlannerContext context, RuleProfilingListener listener) {
         VolcanoPlanner volcanoPlanner = new VolcanoPlanner();
         volcanoPlanner.addRelTraitDef(ConventionTraitDef.INSTANCE);

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

@@ -34,6 +34,7 @@
 public class OpenSearchSort extends Sort implements OpenSearchRelNode {
 
     private final List<String> viableBackends;
+    private final boolean perPartition;
 
     public OpenSearchSort(
         RelOptCluster cluster,
@@ -43,9 +44,28 @@ public OpenSearchSort(
         RexNode offset,
         RexNode fetch,
         List<String> viableBackends
+    ) {
+        this(cluster, traitSet, input, collation, offset, fetch, viableBackends, false);
+    }
+
+    public OpenSearchSort(
+        RelOptCluster cluster,
+        RelTraitSet traitSet,
+        RelNode input,
+        RelCollation collation,
+        RexNode offset,
+        RexNode fetch,
+        List<String> viableBackends,
+        boolean perPartition
     ) {
         super(cluster, traitSet, input, collation, offset, fetch);
         this.viableBackends = viableBackends;
+        this.perPartition = perPartition;
+    }
+
+    /** True when this Sort runs per-shard (shard-bucket oversampling). */
+    public boolean isPerPartition() {
+        return perPartition;
     }
 
     @Override
@@ -65,7 +85,7 @@ public List<FieldStorageInfo> getOutputFieldStorage() {
 
     @Override
     public Sort copy(RelTraitSet traitSet, RelNode input, RelCollation collation, RexNode offset, RexNode fetch) {
-        return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, viableBackends);
+        return new OpenSearchSort(getCluster(), traitSet, input, collation, offset, fetch, viableBackends, perPartition);
     }
 
     /**