opensearch-project · mch2 · May 6, 2026
@@ -33,11 +33,26 @@ public enum AggregateFunction {
     VAR_POP(Type.STATISTICAL, SqlKind.VAR_POP),
     VAR_SAMP(Type.STATISTICAL, SqlKind.VAR_SAMP),
 
+    // Simple — first/last value semantics. PPL emits SqlAggFunction named "first" /
+    // "last"; NAME_ALIASES in NameBasedAggregateFunctionConverter rewrites those to
+    // DataFusion's "first_value"/"last_value" before substrait emission. Planner-side
+    // lookup goes via AggregateFunction.fromNameOrError("FIRST") / ...("LAST").
+    FIRST(Type.SIMPLE, SqlKind.OTHER),
+    LAST(Type.SIMPLE, SqlKind.OTHER),
+
     // State-expanding — state grows with input rows per key
     PERCENTILE_CONT(Type.STATE_EXPANDING, SqlKind.PERCENTILE_CONT),
     PERCENTILE_DISC(Type.STATE_EXPANDING, SqlKind.PERCENTILE_DISC),
     COLLECT(Type.STATE_EXPANDING, SqlKind.COLLECT),
     LISTAGG(Type.STATE_EXPANDING, SqlKind.LISTAGG),
+    TAKE(Type.STATE_EXPANDING, SqlKind.OTHER),
+    // PPL `list(field)` and `values(field)` — NAME_ALIASES in
+    // NameBasedAggregateFunctionConverter rewrites both to DataFusion's native
+    // "array_agg" on the substrait wire. Planner-side lookup goes via
+    // fromNameOrError("LIST") / ("VALUES"). VALUES additionally gets
+    // DISTINCT + ORDER BY forced by AliasConfig; LIST is a pure rename.
+    LIST(Type.STATE_EXPANDING, SqlKind.OTHER),
+    VALUES(Type.STATE_EXPANDING, SqlKind.OTHER),
 
     // Approximate — probabilistic, fixed-size state
     APPROX_COUNT_DISTINCT(Type.APPROXIMATE, SqlKind.OTHER);
@@ -76,10 +91,12 @@ public static AggregateFunction fromSqlKind(SqlKind kind) {
         return null;
     }
 
-    /** Maps an aggregate function name to an AggregateFunction. Throws if not recognized. */
+    /** Maps an aggregate function name to an AggregateFunction. Throws if not recognized.
+     *  Lookup is case-insensitive — Calcite SqlAggFunction names are lowercase
+     *  while enum constants follow Java convention (uppercase). */
     public static AggregateFunction fromNameOrError(String name) {
         try {
-            return valueOf(name);
+            return valueOf(name.toUpperCase(java.util.Locale.ROOT));
         } catch (IllegalArgumentException e) {
             throw new IllegalStateException("Unrecognized aggregate function [" + name + "]", e);
         }

@@ -80,6 +80,11 @@ task buildRustLibrary(type: Exec) {
 
     inputs.files fileTree("${rustWorkspaceDir}/common/src")
     inputs.files fileTree("${rustWorkspaceDir}/lib/src")
+    // The opensearch-datafusion crate (in analytics-backend-datafusion/rust) is a
+    // path-dependency of opensearch-native-lib. Without listing its sources here,
+    // Gradle's UP-TO-DATE check misses changes there and ships a stale dylib.
+    inputs.files fileTree("${projectDir}/../../plugins/analytics-backend-datafusion/rust/src")
+    inputs.file "${projectDir}/../../plugins/analytics-backend-datafusion/rust/Cargo.toml"
     inputs.file "${rustWorkspaceDir}/Cargo.toml"
     outputs.file nativeLibFile
 }

@@ -29,3 +29,4 @@ pub mod query_memory_pool_tracker;
 pub mod runtime_manager;
 pub mod session_context;
 pub mod statistics_cache;
+pub mod udaf;
@@ -70,9 +70,14 @@ impl LocalSession {
             .with_runtime_env(runtime_env)
             .with_default_features()
             .build();
-        Self {
-            ctx: SessionContext::new_with_state(state),
-        }
+        let ctx = SessionContext::new_with_state(state);
+        // Register OpenSearch UDAFs on the coordinator-reduce session so that
+        // aggregates with non-DF-native names (e.g. `approx_count_distinct`
+        // as emitted by isthmus for PPL distinct_count/dc, or `take` for
+        // PPL take) resolve during the substrait consumer's name lookup
+        // at reduce time too, not just on per-shard scan sessions.
+        crate::udaf::register_all(&ctx);
+        Self { ctx }
     }
 
     /// Registers a streaming input on the session under `name` and returns the

@@ -100,6 +100,7 @@ pub async fn execute_query(
         .build();
 
     let ctx = SessionContext::new_with_state(state);
+    crate::udaf::register_all(&ctx);
 
     // Register table via ListingTable — all IO goes through object store
     let file_format = ParquetFormat::new();

@@ -99,6 +99,7 @@ pub async unsafe fn create_session_context(
         .build();
 
     let ctx = SessionContext::new_with_state(state);
+    crate::udaf::register_all(&ctx);
 
     // Register default ListingTable for parquet scans
     let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new()))

@@ -0,0 +1,229 @@
+/*
+ * 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.
+ */
+
+//! `approx_count_distinct` — DataFusion-substrait-facing alias for DF's native
+//! `approx_distinct` UDAF.
+//!
+//! Context: PPL's `distinct_count` / `dc` aliases map to Calcite's
+//! `SqlStdOperatorTable.APPROX_COUNT_DISTINCT`. Isthmus's built-in
+//! `AGGREGATE_SIGS` emits substrait with function name `"approx_count_distinct"`
+//! (matching the core `functions_aggregate_approx.yaml` declaration). But
+//! DataFusion's UDAF registry keys by primary name only, and DF registers its
+//! HyperLogLog-backed impl under `"approx_distinct"` — no alias entry for
+//! `approx_count_distinct`. The substrait consumer's name lookup
+//! (`FunctionRegistry::udaf(name)`) therefore misses.
+//!
+//! This module wraps `datafusion::functions_aggregate::approx_distinct::ApproxDistinct`
+//! and overrides only `name()` to return `"approx_count_distinct"`. Everything
+//! else (signature, accumulator, state fields, aliases) delegates to the inner
+//! impl via a fresh `Arc<AggregateUDF>` we construct at call time. Registered
+//! alongside DF's existing `approx_distinct` UDAF — both keys resolve to the
+//! same HLL implementation, differing only in the registry key.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use datafusion::arrow::datatypes::{DataType, FieldRef};
+use datafusion::common::{Result, ScalarValue};
+use datafusion::logical_expr::function::{
+    AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs,
+};
+use datafusion::logical_expr::utils::AggregateOrderSensitivity;
+use datafusion::logical_expr::{
+    Accumulator, AggregateUDF, AggregateUDFImpl, GroupsAccumulator, ReversedUDAF, Signature,
+    StatisticsArgs,
+};
+
+/// Wrapper around DataFusion's native `ApproxDistinct` UDAF that reports its
+/// `name()` as `"approx_count_distinct"`. Every other trait method delegates to
+/// the inner `approx_distinct` impl.
+///
+/// `AggregateUDFImpl` requires `DynEq + DynHash`; `Arc<AggregateUDF>` implements
+/// `PartialEq + Eq + Hash`, so the derives propagate to the single field.
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ApproxCountDistinctAlias {
+    inner: Arc<AggregateUDF>,
+}
+
+impl ApproxCountDistinctAlias {
+    /// Wraps DataFusion's native `approx_distinct` UDAF (from the
+    /// `datafusion-functions-aggregate` crate, re-exported via
+    /// `datafusion::functions_aggregate`).
+    pub fn new() -> Self {
+        Self {
+            inner: datafusion::functions_aggregate::approx_distinct::approx_distinct_udaf(),
+        }
+    }
+}
+
+impl Default for ApproxCountDistinctAlias {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl AggregateUDFImpl for ApproxCountDistinctAlias {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    /// Only override — the reason this wrapper exists.
+    fn name(&self) -> &str {
+        "approx_count_distinct"
+    }
+
+    fn signature(&self) -> &Signature {
+        self.inner.signature()
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        self.inner.return_type(arg_types)
+    }
+
+    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
+        self.inner.accumulator(acc_args)
+    }
+
+    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
+        self.inner.state_fields(args)
+    }
+
+    fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
+        self.inner.groups_accumulator_supported(args)
+    }
+
+    fn create_groups_accumulator(
+        &self,
+        args: AccumulatorArgs,
+    ) -> Result<Box<dyn GroupsAccumulator>> {
+        self.inner.create_groups_accumulator(args)
+    }
+
+    fn aliases(&self) -> &[String] {
+        self.inner.aliases()
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        self.inner.coerce_types(arg_types)
+    }
+
+    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
+        self.inner.order_sensitivity()
+    }
+
+    fn reverse_expr(&self) -> ReversedUDAF {
+        // AggregateUDF exposes `reverse_udf()` (not `reverse_expr()` of the
+        // underlying impl). It returns `ReversedUDAF`; delegate.
+        self.inner.reverse_udf()
+    }
+
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        self.inner.simplify()
+    }
+
+    fn is_nullable(&self) -> bool {
+        self.inner.is_nullable()
+    }
+
+    fn is_descending(&self) -> Option<bool> {
+        self.inner.is_descending()
+    }
+
+    fn value_from_stats(&self, statistics_args: &StatisticsArgs) -> Option<ScalarValue> {
+        self.inner.value_from_stats(statistics_args)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::arrow::array::{Int32Array, StringArray};
+    use datafusion::arrow::record_batch::RecordBatch;
+    use datafusion::common::Result;
+    use datafusion::execution::context::SessionContext;
+    use datafusion::logical_expr::AggregateUDF;
+    use std::sync::Arc;
+
+    /// Register the alias and verify the substrait-consumer-facing lookup
+    /// `ctx.udaf("approx_count_distinct")` resolves. Guards against DataFusion
+    /// ever adding the alias itself (in which case `with_default_features()`
+    /// would already wire it up and this wrapper would become redundant).
+    #[tokio::test]
+    async fn alias_resolves_by_name() -> Result<()> {
+        let ctx = SessionContext::new();
+
+        // Pre-condition: DF's default features DO register `approx_distinct`,
+        // but NOT `approx_count_distinct`.
+        assert!(
+            ctx.state().aggregate_functions().contains_key("approx_distinct"),
+            "approx_distinct must be registered by default"
+        );
+        assert!(
+            !ctx.state().aggregate_functions().contains_key("approx_count_distinct"),
+            "approx_count_distinct must NOT be registered by default — if this flips, the alias wrapper is redundant"
+        );
+
+        // Register the alias.
+        ctx.register_udaf(AggregateUDF::from(ApproxCountDistinctAlias::new()));
+
+        assert!(
+            ctx.state().aggregate_functions().contains_key("approx_count_distinct"),
+            "approx_count_distinct must be registered after alias registration"
+        );
+        Ok(())
+    }
+
+    /// End-to-end evaluation via SQL: `SELECT approx_count_distinct(col) FROM t`
+    /// must return the expected distinct count on a small enough input that HLL
+    /// is exact. Exercises the full accumulator + groups_accumulator path via
+    /// delegation.
+    #[tokio::test]
+    async fn alias_executes_via_sql() -> Result<()> {
+        use datafusion::arrow::datatypes::{Field, Schema};
+
+        let ctx = SessionContext::new();
+        ctx.register_udaf(AggregateUDF::from(ApproxCountDistinctAlias::new()));
+
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("tag", DataType::Utf8, false),
+        ]));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(Int32Array::from(vec![1, 2, 3, 1, 2, 3, 4])),
+                Arc::new(StringArray::from(vec!["a", "b", "c", "a", "b", "c", "d"])),
+            ],
+        )?;
+        ctx.register_batch("t", batch)?;
+
+        let rows = ctx
+            .sql("SELECT approx_count_distinct(id) AS n_ids, approx_count_distinct(tag) AS n_tags FROM t")
+            .await?
+            .collect()
+            .await?;
+        assert_eq!(rows.len(), 1);
+        let batch = &rows[0];
+        // HLL is exact for <= 100-ish distinct values in DF's 52.x impl.
+        let n_ids = batch
+            .column(0)
+            .as_any()
+            .downcast_ref::<datafusion::arrow::array::UInt64Array>()
+            .expect("n_ids column")
+            .value(0);
+        let n_tags = batch
+            .column(1)
+            .as_any()
+            .downcast_ref::<datafusion::arrow::array::UInt64Array>()
+            .expect("n_tags column")
+            .value(0);
+        assert_eq!(n_ids, 4, "expected 4 distinct id values");
+        assert_eq!(n_tags, 4, "expected 4 distinct tag values");
+        Ok(())
+    }
+}
@@ -0,0 +1,48 @@
+/*
+ * 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.
+ */
+
+//! OpenSearch-specific user-defined aggregate functions registered on every
+//! DataFusion `SessionContext` used by this plugin (per-shard scan + coordinator
+//! reduce). The substrait consumer resolves aggregate references by name against
+//! the session's registry, so it's enough to register here and ship matching
+//! YAML extension entries (see `extensions/opensearch_aggregate.yaml`) on the
+//! Java side.
+
+use std::sync::Arc;
+
+use datafusion::execution::context::SessionContext;
+use datafusion::logical_expr::AggregateUDF;
+
+pub mod approx_count_distinct_alias;
+pub mod take;
+
+/// Register every OpenSearch UDAF on `ctx`. Call once at session construction.
+pub fn register_all(ctx: &SessionContext) {
+    ctx.register_udaf(AggregateUDF::from(take::TakeUdaf::new()));
+    // Alias DataFusion's native `approx_distinct` UDAF under the name
+    // `approx_count_distinct` so the substrait consumer's name-lookup resolves
+    // the core-substrait-YAML signature emitted by isthmus for PPL's
+    // distinct_count / dc. See approx_count_distinct_alias.rs for rationale.
+    ctx.register_udaf(AggregateUDF::from(
+        approx_count_distinct_alias::ApproxCountDistinctAlias::new(),
+    ));
+    log::info!(
+        "OpenSearch UDAF register_all: take, approx_count_distinct (alias for approx_distinct) registered"
+    );
+}
+
+/// Same as [`register_all`] but builds an `Arc<AggregateUDF>` for callers that
+/// only have a `SessionStateBuilder`.
+pub fn all_udafs() -> Vec<Arc<AggregateUDF>> {
+    vec![
+        Arc::new(AggregateUDF::from(take::TakeUdaf::new())),
+        Arc::new(AggregateUDF::from(
+            approx_count_distinct_alias::ApproxCountDistinctAlias::new(),
+        )),
+    ]
+}