diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java index b72e794e93684..d2002587a7326 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java @@ -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); } diff --git a/sandbox/libs/dataformat-native/build.gradle b/sandbox/libs/dataformat-native/build.gradle index fa022b6347435..be3601f00230b 100644 --- a/sandbox/libs/dataformat-native/build.gradle +++ b/sandbox/libs/dataformat-native/build.gradle @@ -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 } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs index ee876888450c9..43d3b069f4db2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -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; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs index d62ec523b477f..941a58ea3ad48 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs @@ -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 diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs index 8ba9c93b3caea..1f258c3e54956 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs @@ -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(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs index 9de9caaa968a5..5647fc703884a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -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())) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/approx_count_distinct_alias.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/approx_count_distinct_alias.rs new file mode 100644 index 0000000000000..629f323c1400d --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/approx_count_distinct_alias.rs @@ -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` 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` implements +/// `PartialEq + Eq + Hash`, so the derives propagate to the single field. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ApproxCountDistinctAlias { + inner: Arc, +} + +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 { + self.inner.return_type(arg_types) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + self.inner.accumulator(acc_args) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + 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> { + self.inner.create_groups_accumulator(args) + } + + fn aliases(&self) -> &[String] { + self.inner.aliases() + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + 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 { + self.inner.simplify() + } + + fn is_nullable(&self) -> bool { + self.inner.is_nullable() + } + + fn is_descending(&self) -> Option { + self.inner.is_descending() + } + + fn value_from_stats(&self, statistics_args: &StatisticsArgs) -> Option { + 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::() + .expect("n_ids column") + .value(0); + let n_tags = batch + .column(1) + .as_any() + .downcast_ref::() + .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(()) + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/mod.rs new file mode 100644 index 0000000000000..6673a8b99eba0 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/mod.rs @@ -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` for callers that +/// only have a `SessionStateBuilder`. +pub fn all_udafs() -> Vec> { + vec![ + Arc::new(AggregateUDF::from(take::TakeUdaf::new())), + Arc::new(AggregateUDF::from( + approx_count_distinct_alias::ApproxCountDistinctAlias::new(), + )), + ] +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs new file mode 100644 index 0000000000000..6de7772b9d6eb --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs @@ -0,0 +1,397 @@ +/* + * 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. + */ + +//! `take(x, n)` — bounded aggregate that collects the first `n` values of `x` +//! into a list. Mirrors PPL's existing Java [`TakeAggFunction`] semantics: +//! +//! * `n > 0`: append values in scan order until the buffer holds `n` items. +//! * `n <= 0`: never append — returns an empty list. +//! * Default `n` (when only one argument supplied): `10`. +//! +//! Distributed correctness: each per-shard accumulator emits its bounded +//! buffer as a `List` scalar via `state()`, the coordinator's final +//! accumulator concatenates incoming lists in `merge_batch()` and trims back +//! to `n`. Cross-shard ordering is non-deterministic — same property as the +//! Java implementation, which assumes single-accumulator scan order. +//! +//! [`TakeAggFunction`]: https://github.com/opensearch-project/sql/blob/main/core/src/main/java/org/opensearch/sql/calcite/udf/udaf/TakeAggFunction.java + +use std::any::Any; +use std::fmt::{Debug, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, ArrayRef, AsArray, ListArray}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{exec_err, Result, ScalarValue}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility}; +use datafusion::physical_expr::expressions::Literal; + +const DEFAULT_LIMIT: i64 = 10; + +/// `take(value, [n])` aggregate UDF. +#[derive(Debug)] +pub struct TakeUdaf { + signature: Signature, +} + +impl TakeUdaf { + pub fn new() -> Self { + Self { + signature: Signature::variadic_any(Volatility::Immutable), + } + } +} + +impl Default for TakeUdaf { + fn default() -> Self { + Self::new() + } +} + +// AggregateUDFImpl requires `DynEq + DynHash`. There's no meaningful "equality" +// between two TakeUdaf instances (they're effectively a singleton), so all +// instances compare equal and hash identically. +impl PartialEq for TakeUdaf { + fn eq(&self, _other: &Self) -> bool { + true + } +} +impl Eq for TakeUdaf {} +impl Hash for TakeUdaf { + fn hash(&self, state: &mut H) { + "take".hash(state); + } +} + +impl AggregateUDFImpl for TakeUdaf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "take" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.is_empty() { + return exec_err!("take() requires at least one argument"); + } + Ok(DataType::List(Arc::new(Field::new( + "item", + arg_types[0].clone(), + true, + )))) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + let element_type = acc_args + .expr_fields + .get(0) + .map(|f| f.data_type().clone()) + .ok_or_else(|| { + datafusion::common::DataFusionError::Execution( + "take(): missing first argument".to_string(), + ) + })?; + let limit = limit_from_args(&acc_args)?; + Ok(Box::new(TakeAccumulator::new(element_type, limit))) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + let element = args + .input_fields + .get(0) + .map(|f| f.data_type().clone()) + .unwrap_or(DataType::Null); + let list_type = DataType::List(Arc::new(Field::new("item", element, true))); + Ok(vec![Arc::new(Field::new( + format!("{}[buf]", args.name), + list_type, + true, + ))]) + } +} + +/// Try to extract the integer limit from the aggregate's second physical +/// expression. Returns: +/// * `Some(DEFAULT_LIMIT)` when the aggregate has only one argument. +/// * `Some(n)` when the second argument is a Substrait literal. +/// * `None` when the second argument is a column reference (Calcite often +/// materializes the literal `n` as a `Project` column instead of inlining it +/// into the aggregate). The accumulator resolves the limit on the first +/// `update_batch` call by reading row 0 of that column. +fn limit_from_args(acc_args: &AccumulatorArgs) -> Result> { + let Some(expr) = acc_args.exprs.get(1) else { + return Ok(Some(DEFAULT_LIMIT)); + }; + if let Some(lit) = expr.as_any().downcast_ref::() { + return scalar_to_i64(lit.value()).map(Some); + } + Ok(None) +} + +fn scalar_to_i64(scalar: &ScalarValue) -> Result { + match scalar { + ScalarValue::Int8(Some(v)) => Ok(*v as i64), + ScalarValue::Int16(Some(v)) => Ok(*v as i64), + ScalarValue::Int32(Some(v)) => Ok(*v as i64), + ScalarValue::Int64(Some(v)) => Ok(*v), + ScalarValue::UInt8(Some(v)) => Ok(*v as i64), + ScalarValue::UInt16(Some(v)) => Ok(*v as i64), + ScalarValue::UInt32(Some(v)) => Ok(*v as i64), + ScalarValue::UInt64(Some(v)) => Ok(*v as i64), + ScalarValue::Int8(None) + | ScalarValue::Int16(None) + | ScalarValue::Int32(None) + | ScalarValue::Int64(None) + | ScalarValue::UInt8(None) + | ScalarValue::UInt16(None) + | ScalarValue::UInt32(None) + | ScalarValue::UInt64(None) => Ok(DEFAULT_LIMIT), + other => exec_err!("take(): n must be an integer, got {other:?}"), + } +} + +/// Bounded buffer of `ScalarValue`s. Capped at `limit`; further appends are +/// no-ops. `limit <= 0` → never accept anything → empty result list. +/// +/// `limit` is `Option` because Calcite often materializes the literal `n` as a +/// `Project` column rather than passing it as a Substrait literal — at +/// accumulator construction time we may only see a column reference, not a +/// literal value. In that case we resolve the limit on the first +/// `update_batch` call from row 0 of the second column (every row carries the +/// same constant value). +struct TakeAccumulator { + element_type: DataType, + limit: Option, + buf: Vec, +} + +impl TakeAccumulator { + fn new(element_type: DataType, limit: Option) -> Self { + let cap = limit.filter(|&l| l > 0).map(|l| l as usize).unwrap_or(0); + Self { + element_type, + limit, + buf: Vec::with_capacity(cap), + } + } + + fn cap(&self) -> usize { + match self.limit { + Some(l) if l > 0 => l as usize, + _ => 0, + } + } + + /// If the limit is still unknown, try to read it from row 0 of `n_col`. + /// The Calcite plan materializes the literal `n` as a constant column — + /// every row carries the same value, so reading row 0 is sufficient. + fn resolve_limit_from(&mut self, n_col: &ArrayRef) -> Result<()> { + if self.limit.is_some() || n_col.len() == 0 { + return Ok(()); + } + let scalar = ScalarValue::try_from_array(n_col, 0)?; + self.limit = Some(scalar_to_i64(&scalar)?); + Ok(()) + } + + fn current_list(&self) -> ScalarValue { + let arr = ScalarValue::new_list_nullable(&self.buf, &self.element_type); + ScalarValue::List(arr) + } +} + +impl Debug for TakeAccumulator { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TakeAccumulator") + .field("limit", &self.limit) + .field("len", &self.buf.len()) + .finish() + } +} + +impl Accumulator for TakeAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if let Some(n_col) = values.get(1) { + self.resolve_limit_from(n_col)?; + } + let cap = self.cap(); + if cap == 0 { + return Ok(()); + } + let col = &values[0]; + let mut i = 0; + while self.buf.len() < cap && i < col.len() { + self.buf.push(ScalarValue::try_from_array(col, i)?); + i += 1; + } + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(self.current_list()) + } + + fn size(&self) -> usize { + std::mem::size_of_val(self) + self.buf.iter().map(|s| s.size()).sum::() + } + + fn state(&mut self) -> Result> { + Ok(vec![self.current_list()]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let cap = self.cap(); + if cap == 0 { + return Ok(()); + } + let lists: &ListArray = states[0].as_list(); + for i in 0..lists.len() { + if self.buf.len() >= cap { + return Ok(()); + } + if lists.is_null(i) { + continue; + } + let inner = lists.value(i); + for j in 0..inner.len() { + if self.buf.len() >= cap { + return Ok(()); + } + self.buf.push(ScalarValue::try_from_array(&inner, j)?); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int32Array, StringArray}; + + fn varchar_acc(limit: Option) -> TakeAccumulator { + TakeAccumulator::new(DataType::Utf8, limit) + } + + fn strings(values: &[&str]) -> ArrayRef { + Arc::new(StringArray::from(values.to_vec())) as ArrayRef + } + + fn const_int32(value: i32, len: usize) -> ArrayRef { + Arc::new(Int32Array::from(vec![value; len])) as ArrayRef + } + + fn list_len(scalar: &ScalarValue) -> usize { + match scalar { + ScalarValue::List(arr) => arr.value(0).len(), + other => panic!("expected List scalar, got {other:?}"), + } + } + + fn list_strings(scalar: &ScalarValue) -> Vec { + let inner = match scalar { + ScalarValue::List(arr) => arr.value(0), + other => panic!("expected List scalar, got {other:?}"), + }; + let s = inner.as_any().downcast_ref::().expect("string array"); + (0..s.len()).map(|i| s.value(i).to_string()).collect() + } + + /// `take(col, lit(2))` — limit known at construction time. Should cap at 2. + #[test] + fn limit_from_literal_caps_buffer() { + let mut acc = varchar_acc(Some(2)); + let col = strings(&["a", "b", "c", "d", "e"]); + // No second column — caller passed a Substrait literal, so DF doesn't + // materialize an n-column. We model this by passing a single arg. + acc.update_batch(&[col]).expect("update_batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_len(&out), 2); + assert_eq!(list_strings(&out), vec!["a", "b"]); + } + + /// `take(col, $f1=2)` — Calcite materialized `2` as a project column. The + /// accumulator gets `Option::None` at construction; the limit is resolved + /// from row 0 of the second column on the first update_batch. Regression + /// test for the bug where the limit was hardcoded to DEFAULT_LIMIT (10) + /// because we only handled the literal path. + #[test] + fn limit_from_column_resolved_on_first_update() { + let mut acc = varchar_acc(None); + let col = strings(&["a", "b", "c", "d", "e"]); + let n_col = const_int32(2, 5); // every row carries the constant n=2 + acc.update_batch(&[col, n_col]).expect("update_batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_len(&out), 2); + assert_eq!(list_strings(&out), vec!["a", "b"]); + } + + /// `take(col)` — no second arg, limit defaults to 10. + #[test] + fn missing_limit_defaults_to_ten() { + let mut acc = varchar_acc(Some(DEFAULT_LIMIT)); + let col = strings(&["a", "b", "c"]); + acc.update_batch(&[col]).expect("update_batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_len(&out), 3, "fewer rows than DEFAULT_LIMIT — collect all"); + } + + /// `take(col, 0)` and `take(col, -1)` produce empty lists. Mirrors the + /// PPL TakeAggFunction behavior — no rows accepted when n <= 0. + #[test] + fn non_positive_limit_yields_empty_list() { + for n in [Some(0), Some(-3)] { + let mut acc = varchar_acc(n); + let col = strings(&["a", "b", "c"]); + acc.update_batch(&[col]).expect("update_batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_len(&out), 0, "n={n:?} should produce empty list"); + } + } + + /// Multiple update_batch calls past the cap are no-ops. + #[test] + fn cap_holds_across_multiple_batches() { + let mut acc = varchar_acc(Some(2)); + acc.update_batch(&[strings(&["a"])]).expect("first batch"); + acc.update_batch(&[strings(&["b", "c", "d"])]).expect("second batch"); + acc.update_batch(&[strings(&["e", "f"])]).expect("third batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_strings(&out), vec!["a", "b"]); + } + + /// merge_batch concatenates partial-state lists and trims to cap. This is + /// the coordinator-reduce path for the multi-shard case. + #[test] + fn merge_batch_trims_concatenated_partials() { + let mut acc = varchar_acc(Some(3)); + let mut shard1 = varchar_acc(Some(3)); + shard1.update_batch(&[strings(&["a", "b"])]).expect("shard1 update"); + let mut shard2 = varchar_acc(Some(3)); + shard2.update_batch(&[strings(&["c", "d", "e"])]).expect("shard2 update"); + + // Build a ListArray containing both shard states as rows. ScalarValue + // doesn't easily concatenate two single-element list scalars, so build + // the ListArray manually via the same helper the partials use. + let s1 = shard1.evaluate().expect("shard1 eval"); + let s2 = shard2.evaluate().expect("shard2 eval"); + let combined = ScalarValue::iter_to_array(vec![s1, s2]).expect("iter_to_array"); + + acc.merge_batch(&[combined]).expect("merge_batch"); + let out = acc.evaluate().expect("evaluate"); + assert_eq!(list_strings(&out), vec!["a", "b", "c"]); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/AggregationUDFIT.java b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/AggregationUDFIT.java new file mode 100644 index 0000000000000..35217d0c414ee --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/internalClusterTest/java/org/opensearch/be/datafusion/AggregationUDFIT.java @@ -0,0 +1,194 @@ +/* + * 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.be.datafusion; + +import org.opensearch.Version; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.analytics.AnalyticsPlugin; +import org.opensearch.arrow.flight.transport.FlightStreamPlugin; +import org.opensearch.be.lucene.LucenePlugin; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.composite.CompositeDataFormatPlugin; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.parquet.ParquetDataFormatPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginInfo; +import org.opensearch.ppl.TestPPLPlugin; +import org.opensearch.ppl.action.PPLRequest; +import org.opensearch.ppl.action.PPLResponse; +import org.opensearch.ppl.action.UnifiedPPLExecuteAction; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.util.Collection; +import java.util.Collections; +import java.util.List; + +/** + * End-to-end test for PPL -> Datafusion UDAFs. + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 1) +public class AggregationUDFIT extends OpenSearchIntegTestCase { + + private static final String BANK_INDEX = "bank"; + + @Override + protected Collection> nodePlugins() { + // Plugins with no extendedPlugins requirement go here. Plugins that need + // explicit extendedPlugins (so SPI ExtensionLoader walks the right parent + // classloader) are declared in additionalNodePlugins() below. + return List.of(TestPPLPlugin.class, FlightStreamPlugin.class, CompositeDataFormatPlugin.class, LucenePlugin.class); + } + + @Override + protected Collection additionalNodePlugins() { + // OpenSearchIntegTestCase's nodePlugins() builds PluginInfo with empty + // extendedPlugins, which breaks ExtensiblePlugin.loadExtensions(...) for + // plugins like DataFusionPlugin that ride on AnalyticsPlugin's SPI. Use + // additionalNodePlugins() to declare the parent relationships explicitly. + return List.of( + classpathPlugin(AnalyticsPlugin.class, Collections.emptyList()), + classpathPlugin(ParquetDataFormatPlugin.class, Collections.emptyList()), + classpathPlugin(DataFusionPlugin.class, List.of(AnalyticsPlugin.class.getName())) + ); + } + + private static PluginInfo classpathPlugin(Class pluginClass, List extendedPlugins) { + return new PluginInfo( + pluginClass.getName(), + "classpath plugin", + "NA", + Version.CURRENT, + "1.8", + pluginClass.getName(), + null, + extendedPlugins, + false + ); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put(FeatureFlags.PLUGGABLE_DATAFORMAT_EXPERIMENTAL_FLAG, true) + .build(); + } + + @Override + public void setUp() throws Exception { + super.setUp(); + createBankIndex(); + indexBankDocs(); + ensureGreen(BANK_INDEX); + refresh(BANK_INDEX); + } + + public void testTake() throws Exception { + PPLRequest request = new PPLRequest("source=" + BANK_INDEX + " | stats take(firstname, 2) as take"); + PPLResponse response = client().execute(UnifiedPPLExecuteAction.INSTANCE, request).actionGet(); + + assertNotNull("PPLResponse must not be null", response); + assertEquals("schema columns", List.of("take"), response.getColumns()); + assertEquals("scalar agg → exactly 1 result row", 1, response.getRows().size()); + + Object cell = response.getRows().get(0)[0]; + assertNotNull("take cell must not be null", cell); + assertTrue("take cell must materialize as a List, got " + cell.getClass(), cell instanceof List); + @SuppressWarnings("unchecked") + List taken = (List) cell; + assertEquals("take(firstname, 2) over the bank fixture", List.of("Amber JOHnny", "Hattie"), taken); + } + + private void createBankIndex() throws Exception { + XContentBuilder mapping = XContentFactory.jsonBuilder() + .startObject() + .startObject("properties") + .startObject("account_number") + .field("type", "long") + .endObject() + .startObject("firstname") + .field("type", "keyword") + .endObject() + .startObject("lastname") + .field("type", "keyword") + .endObject() + .startObject("balance") + .field("type", "long") + .endObject() + .startObject("age") + .field("type", "integer") + .endObject() + .startObject("gender") + .field("type", "keyword") + .endObject() + .endObject() + .endObject(); + + // Parquet-backed composite index — DataFusion only reads parquet, so the take + // UDAF will only fire if the data lands in a parquet shard. Mirrors the + // settings used by CoordinatorReduceIT on df_reduce. + Settings indexSettings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put("index.pluggable.dataformat.enabled", true) + .put("index.pluggable.dataformat", "composite") + .put("index.composite.primary_data_format", "parquet") + .putList("index.composite.secondary_data_formats") + .build(); + + CreateIndexResponse response = client().admin() + .indices() + .prepareCreate(BANK_INDEX) + .setSettings(indexSettings) + .setMapping(mapping) + .get(); + assertTrue("bank index creation must be acknowledged", response.isAcknowledged()); + } + + private void indexBankDocs() { + // First 5 docs of sql/integ-test/src/test/resources/bank.json — preserves the + // ordering the SQL plugin's testTake assumes (Amber JOHnny, Hattie are #0 and #1). + client().prepareIndex(BANK_INDEX) + .setId("1") + .setSource( + "account_number", + 1, + "firstname", + "Amber JOHnny", + "lastname", + "Duke Willmington", + "balance", + 39225L, + "age", + 32, + "gender", + "M" + ) + .get(); + client().prepareIndex(BANK_INDEX) + .setId("6") + .setSource("account_number", 6, "firstname", "Hattie", "lastname", "Bond", "balance", 5686L, "age", 36, "gender", "M") + .get(); + client().prepareIndex(BANK_INDEX) + .setId("13") + .setSource("account_number", 13, "firstname", "Nanette", "lastname", "Bates", "balance", 32838L, "age", 28, "gender", "F") + .get(); + client().prepareIndex(BANK_INDEX) + .setId("18") + .setSource("account_number", 18, "firstname", "Dale", "lastname", "Adams", "balance", 4180L, "age", 33, "gender", "M") + .get(); + client().prepareIndex(BANK_INDEX) + .setId("20") + .setSource("account_number", 20, "firstname", "Elinor", "lastname", "Ratliff", "balance", 16418L, "age", 36, "gender", "M") + .get(); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverter.java new file mode 100644 index 0000000000000..c9529afb18246 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverter.java @@ -0,0 +1,379 @@ +/* + * 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 io.substrait.isthmus.expression; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlOperator; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; + +import io.substrait.expression.AggregateFunctionInvocation; +import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; +import io.substrait.expression.FunctionArg; +import io.substrait.extension.SimpleExtension; +import io.substrait.isthmus.TypeConverter; +import io.substrait.type.Type; + +/** + * Aggregate function converter that handles two issues with the stock isthmus matcher: + * + *
    + *
  1. Identity-only operator lookup. PPL emits its own {@code SqlAggFunction} + * subclasses (e.g. {@code TakeAggFunction}, {@code FirstAggFunction}). The default + * converter keys its matcher map by {@link SqlOperator} identity, so the PPL + * instances miss the entries seeded by the stub Sigs in + * {@code DataFusionFragmentConvertor}. {@link #getFunctionFinder} falls back to a + * case-insensitive name match against the seeded entries.
    2. + * + *
  2. Wildcard signatures don't direct-match. The stock {@code attemptMatch} + * does string-level key matching: it builds a key from the call's runtime types + * (e.g. {@code take:string_i32}) and looks it up in the YAML-derived directMap, + * which contains keys built from the YAML wildcards (e.g. {@code take:any_i32}). + * The two never match. The {@code singularInputType} fallback only handles + * uniform-arg signatures and so misses heterogeneous wildcard signatures like + * {@code take(any1, i32)}, {@code arg_min(any1, any2)}, {@code first(any1)}. + * {@link #convert} adds a name-keyed fallback that, on miss, picks the first + * extension variant whose name and arity match and constructs the + * {@link AggregateFunctionInvocation} directly — bypassing the strict directMap + * lookup. Argument-type validity is left to the substrait consumer (DataFusion).
    3. + *
+ * + *

Alias rewrite mechanism

+ * + *

PPL frontend names sometimes don't match the DataFusion-side variant name, or need + * shape tweaks (DISTINCT, ORDER BY) that the stock matcher can't infer from the call + * shape alone. {@link #NAME_ALIASES} declares these per-name, and + * {@link #rewriteAlias(RelNode, AggregateCall, Function)} applies the configured + * shape changes before substrait emission. Current aliases: + * + *

    + *
  • {@code first → first_value}, {@code last → last_value} — pure rename. PPL's + * stats first(x)/last(x) map to DataFusion's native first_value/last_value.
    • + *
  • {@code list → array_agg}, {@code values → array_agg} — pure rename for list; + * rename + force-distinct + force-order-by for values (ascending on the operand + * itself, per PPL values() docs "sorted lexicographically, unique only").
    • + *
  • {@code arg_min → first_value}, {@code arg_max → last_value} — rename + pull + * the second arg out of argList into a synthesized ORDER BY sort field. PPL's + * stats earliest(field, ts) / latest(field, ts) lower to Calcite ARG_MIN / ARG_MAX + * at the frontend layer. DataFusion 52.x has no native min_by/max_by UDAF — the + * semantic equivalent is first_value/last_value with an ORDER BY on the key. + * ARG_MAX uses last_value with ASC (last row of ascending sort == max row).
    • + *
+ */ +public class NameBasedAggregateFunctionConverter extends AggregateFunctionConverter { + + /** + * Per-alias rewrite specification. A config with only {@link #target} set is a pure + * rename — the aliased call uses the named YAML variant but keeps its own operands + * and flags. Optional fields inject additional shape: + * + *
    + *
  • {@link #sortArgIndex} — pull the operand at this index out of the call's + * argList and emit it as a substrait sort field (ASC_NULLS_LAST) on the + * resulting measure. Used for ARG_MIN/ARG_MAX → first_value/last_value (sort + * by arg[1], the key) and VALUES → array_agg (sort by arg[0], the operand + * itself). When set, the operand at this index is NOT kept in the operand + * list — it appears only in the sort field.
    • + *
  • {@link #forceDistinct} — override {@link AggregateCall#isDistinct()} and + * emit DISTINCT on the substrait invocation. Used for VALUES which requires + * distinct-by-construction semantics but whose PPL-frontend SqlAggFunction + * doesn't carry the DISTINCT flag.
    • + *
+ * + *

When {@code sortArgIndex} is set to the SAME index as a regular operand (i.e. + * the sort key IS the aggregated value, like VALUES sorts by its own operand), the + * operand is still emitted as a sort field AND kept as the primary operand. This + * is the "sort by self" case. See {@link #rewriteAlias} for the implementation. + */ + static final class AliasConfig { + final String target; + final Integer sortArgIndex; + final boolean sortIsSelf; + final boolean forceDistinct; + + private AliasConfig(String target, Integer sortArgIndex, boolean sortIsSelf, boolean forceDistinct) { + this.target = target; + this.sortArgIndex = sortArgIndex; + this.sortIsSelf = sortIsSelf; + this.forceDistinct = forceDistinct; + } + + /** Pure rename: aliased call routes to {@code target} with operands/flags unchanged. */ + static AliasConfig rename(String target) { + return new AliasConfig(target, null, false, false); + } + + /** Rename + pull argList[sortArgIndex] into a synthesized ASC_NULLS_LAST sort field, + * removing it from the operand list. */ + static AliasConfig renameWithSortArg(String target, int sortArgIndex) { + return new AliasConfig(target, sortArgIndex, false, false); + } + + /** Rename + emit the SINGLE operand as both the value AND an ASC_NULLS_LAST sort + * field, also setting DISTINCT. For PPL values(): sort by operand, dedup by + * operand. Only valid when arity == 1. */ + static AliasConfig renameDistinctSortedBySelf(String target) { + return new AliasConfig(target, 0, true, true); + } + + boolean hasReshape() { + return sortArgIndex != null || forceDistinct; + } + } + + /** + * PPL operator name → alias config. Keys are case-insensitive (stored lowercase). + * Entries without reshape flags (pure {@link AliasConfig#rename}) still flow + * through the {@link #convertByName} fallback unchanged — their target name is + * used for YAML variant lookup. Entries with reshape flags route through + * {@link #rewriteAlias} which assembles the substrait invocation directly. + */ + private static final Map NAME_ALIASES = Map.of( + "first", + AliasConfig.rename("first_value"), + "last", + AliasConfig.rename("last_value"), + "list", + AliasConfig.rename("array_agg"), + "values", + AliasConfig.renameDistinctSortedBySelf("array_agg"), + "arg_min", + AliasConfig.renameWithSortArg("first_value", 1), + "arg_max", + AliasConfig.renameWithSortArg("last_value", 1) + ); + + private final List allVariants; + private final TypeConverter typeConverter; + + public NameBasedAggregateFunctionConverter( + List functions, + List additionalSignatures, + RelDataTypeFactory typeFactory, + TypeConverter typeConverter + ) { + super(functions, additionalSignatures, typeFactory, typeConverter); + this.allVariants = List.copyOf(functions); + this.typeConverter = typeConverter; + } + + @Override + protected FunctionFinder getFunctionFinder(AggregateCall call) { + FunctionFinder ff = super.getFunctionFinder(call); + if (ff != null) { + return ff; + } + String name = call.getAggregation().getName(); + if (name == null) { + return null; + } + for (Map.Entry entry : signatures.entrySet()) { + if (name.equalsIgnoreCase(entry.getKey().getName())) { + return entry.getValue(); + } + } + return null; + } + + @Override + public Optional convert( + RelNode input, + Type.Struct inputType, + AggregateCall call, + Function topLevelConverter + ) { + // Reshape aliases (DISTINCT injection, ORDER BY synthesis) run BEFORE the stock + // matcher so we can fully control the emitted call shape. Pure-rename aliases + // don't need pre-emption — they flow through convertByName's NAME_ALIASES + // lookup on target name. + Optional aliasRewrite = rewriteAlias(input, call, topLevelConverter); + if (aliasRewrite.isPresent()) { + return aliasRewrite; + } + Optional result = super.convert(input, inputType, call, topLevelConverter); + if (result.isPresent()) { + return result; + } + return convertByName(input, call, topLevelConverter); + } + + /** + * Applies an {@link AliasConfig} to rewrite the call's emitted shape when the config + * specifies reshape flags (sort field synthesis or forced DISTINCT). Pure-rename + * configs return empty and fall through to {@link #convertByName}, which consults + * {@link #NAME_ALIASES} for the target variant name. + * + *

Returns empty when: + *

    + *
  • the call has no name (bare aggregate binding with no SqlAggFunction);
    • + *
  • the name has no entry in {@link #NAME_ALIASES};
    • + *
  • the entry is a pure rename with no reshape flags;
    • + *
  • the target YAML variant isn't present in the loaded extension collection;
    • + *
  • the call arity is incompatible with the config (e.g. sortArgIndex out of + * range, or sortIsSelf with arity != 1).
    • + *
+ */ + private Optional rewriteAlias( + RelNode input, + AggregateCall call, + Function topLevelConverter + ) { + String callName = call.getAggregation().getName(); + if (callName == null) { + return Optional.empty(); + } + AliasConfig config = NAME_ALIASES.get(callName.toLowerCase(Locale.ROOT)); + if (config == null || !config.hasReshape()) { + return Optional.empty(); + } + + int argCount = call.getArgList().size(); + if (config.sortArgIndex != null && config.sortArgIndex >= argCount) { + return Optional.empty(); + } + if (config.sortIsSelf && argCount != 1) { + return Optional.empty(); + } + + // The post-rewrite operand count determines which YAML variant we need. + // sortIsSelf keeps the operand (arity stays at 1). Plain sortArgIndex removes + // the operand at that index from the operand list (arity drops by 1). + int postRewriteArity = config.sortIsSelf ? argCount : (config.sortArgIndex != null ? argCount - 1 : argCount); + SimpleExtension.AggregateFunctionVariant matched = null; + for (SimpleExtension.AggregateFunctionVariant variant : allVariants) { + if (!variant.name().equalsIgnoreCase(config.target)) { + continue; + } + if (variant.requiredArguments().size() == postRewriteArity || variant.args().size() == postRewriteArity) { + matched = variant; + break; + } + } + if (matched == null) { + return Optional.empty(); + } + + // Build the operand list. If sortArgIndex is set and sortIsSelf is false, skip + // that argList index — it's moved into the sort field list only. + List operands = new ArrayList<>(postRewriteArity); + for (int i = 0; i < argCount; i++) { + if (config.sortArgIndex != null && !config.sortIsSelf && i == config.sortArgIndex) { + continue; + } + int colIdx = call.getArgList().get(i); + operands.add(topLevelConverter.apply(input.getCluster().getRexBuilder().makeInputRef(input, colIdx))); + } + + // Build the sort field list. ASC_NULLS_LAST for both directions — see the + // class-level javadoc for the ARG_MIN/ARG_MAX reasoning. first_value picks the + // smallest key (first row of ascending sort); last_value picks the largest + // (last row of ascending sort). + List sorts = Collections.emptyList(); + if (config.sortArgIndex != null) { + int sortColIdx = call.getArgList().get(config.sortArgIndex); + Expression sortExpr = topLevelConverter.apply(input.getCluster().getRexBuilder().makeInputRef(input, sortColIdx)); + Expression.SortField sortField = Expression.SortField.builder() + .expr(sortExpr) + .direction(Expression.SortDirection.ASC_NULLS_LAST) + .build(); + sorts = List.of(sortField); + } + + Type outputType = typeConverter.toSubstrait(call.getType()); + + boolean distinct = config.forceDistinct || call.isDistinct(); + Expression.AggregationInvocation invocation = distinct + ? Expression.AggregationInvocation.DISTINCT + : Expression.AggregationInvocation.ALL; + + return Optional.of( + ExpressionCreator.aggregateFunction( + matched, + outputType, + Expression.AggregationPhase.INITIAL_TO_RESULT, + sorts, + invocation, + operands.stream().toList() + ) + ); + } + + /** + * Looks up an extension variant by aggregate name and arity, then assembles an + * {@link AggregateFunctionInvocation} directly. Used as a last resort when the + * stock matcher cannot resolve the call (typically because the YAML signature uses + * wildcards that the directMap lookup cannot expand against runtime types). + * + *

Consults {@link #NAME_ALIASES} for pure-rename aliases — e.g. a call named + * "first" is looked up as "first_value" in the extension collection. Alias entries + * with reshape flags are handled earlier in {@link #rewriteAlias}; by the time we + * land here the call has either no alias or a pure rename. + */ + private Optional convertByName( + RelNode input, + AggregateCall call, + Function topLevelConverter + ) { + String callName = call.getAggregation().getName(); + if (callName == null) { + return Optional.empty(); + } + String lower = callName.toLowerCase(Locale.ROOT); + AliasConfig aliasConfig = NAME_ALIASES.get(lower); + String lookup = aliasConfig != null ? aliasConfig.target : lower; + int argCount = call.getArgList().size(); + + SimpleExtension.AggregateFunctionVariant matched = null; + for (SimpleExtension.AggregateFunctionVariant variant : allVariants) { + if (!variant.name().equalsIgnoreCase(lookup)) { + continue; + } + if (variant.requiredArguments().size() == argCount || variant.args().size() == argCount) { + matched = variant; + break; + } + } + if (matched == null) { + return Optional.empty(); + } + + List operands = call.getArgList() + .stream() + .map(idx -> input.getCluster().getRexBuilder().makeInputRef(input, idx)) + .map(topLevelConverter) + .toList(); + + Type outputType = typeConverter.toSubstrait(call.getType()); + + Expression.AggregationInvocation invocation = call.isDistinct() + ? Expression.AggregationInvocation.DISTINCT + : Expression.AggregationInvocation.ALL; + + return Optional.of( + ExpressionCreator.aggregateFunction( + matched, + outputType, + Expression.AggregationPhase.INITIAL_TO_RESULT, + Collections.emptyList(), + invocation, + operands.stream().toList() + ) + ); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 93b79e17c9acf..0796620da54b9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -83,7 +83,13 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP AggregateFunction.MIN, AggregateFunction.MAX, AggregateFunction.COUNT, - AggregateFunction.AVG + AggregateFunction.AVG, + // PPL `first(field)` / `last(field)` — renamed to DataFusion's native first_value / + // last_value on the wire via NameBasedAggregateFunctionConverter.NAME_ALIASES. Without + // ORDER BY, DataFusion returns an arbitrary element; the PPL "first/last in document + // order" guarantee is a documented tradeoff. + AggregateFunction.FIRST, + AggregateFunction.LAST ); private final DataFusionPlugin plugin; @@ -142,6 +148,14 @@ public Set aggregateCapabilities() { caps.add(AggregateCapability.simple(func, Set.of(type), formats)); } } + // PPL TAKE — collect first N values into a list. State-expanding (state grows with N). + caps.add(AggregateCapability.stateExpanding(AggregateFunction.TAKE, Set.copyOf(SUPPORTED_FIELD_TYPES), formats)); + // PPL LIST — collect all values into an array (renamed to DataFusion's + // native array_agg on the wire). State-expanding (state grows with input size). + caps.add(AggregateCapability.stateExpanding(AggregateFunction.LIST, Set.copyOf(SUPPORTED_FIELD_TYPES), formats)); + // PPL VALUES — DISTINCT + sorted-ascending variant of LIST. Also routes to + // array_agg but with AliasConfig-forced DISTINCT + sort-by-self. + caps.add(AggregateCapability.stateExpanding(AggregateFunction.VALUES, Set.copyOf(SUPPORTED_FIELD_TYPES), formats)); return Set.copyOf(caps); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 101893acb229e..8b1ecf0807908 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -24,8 +24,13 @@ import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.schema.ColumnStrategy; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.util.ImmutableBitSet; +import org.apache.calcite.util.Optionality; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; @@ -43,6 +48,7 @@ import io.substrait.isthmus.TypeConverter; import io.substrait.isthmus.expression.AggregateFunctionConverter; import io.substrait.isthmus.expression.FunctionMappings; +import io.substrait.isthmus.expression.NameBasedAggregateFunctionConverter; import io.substrait.isthmus.expression.ScalarFunctionConverter; import io.substrait.isthmus.expression.WindowFunctionConverter; import io.substrait.plan.Plan; @@ -263,7 +269,26 @@ private SubstraitRelVisitor createVisitor(RelNode relNode) { typeFactory, typeConverter ); - AggregateFunctionConverter aggConverter = new AggregateFunctionConverter(extensions.aggregateFunctions(), typeFactory); + // PPL aggregates (e.g. `take`, `first`, `last`) aren't in standard Calcite — emit stub + // SqlAggFunction instances whose only job is to seed the converter's name→FunctionFinder + // map so NameBasedAggregateFunctionConverter can route the actual PPL operator + // instances (a different Java object per-call) by case-insensitive name match. + // Sig names match the PPL operator name (pre-alias). NAME_ALIASES inside + // NameBasedAggregateFunctionConverter rewrites to the YAML variant name where needed + // (e.g. "first" → "first_value"). + List additionalAggSigs = List.of( + FunctionMappings.s(stubAgg("take"), "take"), + FunctionMappings.s(stubAgg("first"), "first"), + FunctionMappings.s(stubAgg("last"), "last"), + FunctionMappings.s(stubAgg("list"), "list"), + FunctionMappings.s(stubAgg("values"), "values") + ); + AggregateFunctionConverter aggConverter = new NameBasedAggregateFunctionConverter( + extensions.aggregateFunctions(), + additionalAggSigs, + typeFactory, + typeConverter + ); WindowFunctionConverter windowConverter = new WindowFunctionConverter(extensions.windowFunctions(), typeFactory); ConverterProvider converterProvider = new ConverterProvider( typeFactory, @@ -293,6 +318,28 @@ private static byte[] serializePlan(Plan plan) { return new PlanProtoConverter().toProto(plan).toByteArray(); } + /** + * Minimal {@link SqlAggFunction} acting as a name→FunctionFinder map key. + * PPL emits its own SqlAggFunction instances; identity lookup against these + * stubs misses, but {@link NameBasedAggregateFunctionConverter} falls back + * to matching on operator name, which the stub provides. + */ + private static SqlAggFunction stubAgg(String name) { + return new SqlAggFunction( + name, + null, + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ) { + }; + } + // ── Calcite TableScan wrappers for OpenSearchStageInputScan rewrite ───────── /** diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index a9e9ecfe27d0e..f4c8d937a8055 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -148,8 +148,28 @@ private static SimpleExtension.ExtensionCollection loadSubstraitExtensions() { ClassLoader previous = t.getContextClassLoader(); try { t.setContextClassLoader(DataFusionPlugin.class.getClassLoader()); + SimpleExtension.ExtensionCollection collection = DefaultExtensionCatalog.DEFAULT_COLLECTION; + + // Layer in the upstream delegation-function extensions (DelegatedPredicateFunction + // signature used by the Lucene predicate-pushdown path). SimpleExtension.ExtensionCollection delegationExtensions = SimpleExtension.load(List.of("/delegation_functions.yaml")); - return DefaultExtensionCatalog.DEFAULT_COLLECTION.merge(delegationExtensions); + collection = collection.merge(delegationExtensions); + + // Layer in OpenSearch-specific aggregates — the PPL `take(x, n)` UDAF backed + // by a custom Rust impl, and the named `first_value` / `last_value` / `array_agg` + // signatures consumed by NameBasedAggregateFunctionConverter. The YAML lives at + // /extensions/opensearch_aggregate.yaml on the plugin classpath. + try (java.io.InputStream stream = DataFusionPlugin.class.getResourceAsStream("/extensions/opensearch_aggregate.yaml")) { + if (stream != null) { + SimpleExtension.ExtensionCollection custom = SimpleExtension.load(stream); + collection = collection.merge(custom); + } else { + logger.warn("opensearch_aggregate.yaml not found on plugin classpath"); + } + } catch (java.io.IOException e) { + throw new RuntimeException("Failed to load opensearch_aggregate.yaml", e); + } + return collection; } finally { t.setContextClassLoader(previous); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java index 01449e41a20dc..f2f1bdad79f71 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionResultStream.java @@ -14,7 +14,9 @@ import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.ValueVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; @@ -25,6 +27,7 @@ import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.core.action.ActionListener; +import java.util.ArrayList; import java.util.Iterator; import java.util.List; import java.util.NoSuchElementException; @@ -200,7 +203,43 @@ public Object getFieldValue(String fieldName, int rowIndex) { if (vector == null) { throw new IllegalArgumentException("Unknown field: " + fieldName); } + // List-valued aggregates (PPL list/take/values) return a ListVector whose + // stock getObject wraps results in arrow-vector's JsonStringArrayList. That + // class's instantiates Jackson's JavaTimeModule, which lives in + // jackson-datatype-jsr310 — NOT on the arrow-flight-rpc plugin classloader + // that owns ListVector. In the REST IT path (plugin classloader isolation), + // that lookup throws NoClassDefFoundError and crashes the node. + // + // Decode ListVector manually using offsets + the inner vector's getObject so + // we never touch JsonStringArrayList. The inner vector's getObject returns a + // plain scalar (String/Number/etc.) whose has no jsr310 dep. + if (vector instanceof ListVector listVector) { + return decodeListValue(listVector, rowIndex); + } return vector.getObject(rowIndex); } + + /** + * Materialize a single row of a {@link ListVector} as a plain {@link ArrayList}. + * Uses the public offset/data API rather than {@code ListVector.getObject}, which + * returns a {@code JsonStringArrayList} whose static init fails under plugin + * classloader isolation (see {@link #getFieldValue} comment). + * + * @return {@code null} if the row is null; otherwise an {@link ArrayList} of the + * row's element values as returned by the inner vector's {@code getObject}. + */ + private static List decodeListValue(ListVector listVector, int rowIndex) { + if (listVector.isNull(rowIndex)) { + return null; + } + int start = listVector.getElementStartIndex(rowIndex); + int end = listVector.getElementEndIndex(rowIndex); + ValueVector inner = listVector.getDataVector(); + ArrayList out = new ArrayList<>(end - start); + for (int i = start; i < end; i++) { + out.add(inner.getObject(i)); + } + return out; + } } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/extensions/opensearch_aggregate.yaml b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/extensions/opensearch_aggregate.yaml new file mode 100644 index 0000000000000..0454ae9c76b2e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/extensions/opensearch_aggregate.yaml @@ -0,0 +1,76 @@ +%YAML 1.2 +--- +urn: extension:org.opensearch:opensearch_aggregate +aggregate_functions: + # PPL `take(x, n)` — collect up to N values into a list. Backed by a custom + # Rust UDAF in datafusion. The return uses substrait's `LIST?<…>` syntax + # (uppercase, nullability between LIST and the type parameter) — the lowercase + # `list<…>?` form fails the substrait-java type parser. + - name: "take" + description: "Collect first N values into a list." + impls: + - args: + - name: x + value: any1 + - name: n + value: i32 + nullability: DECLARED_OUTPUT + decomposable: NONE + return: LIST? + + # PPL `first(field)` — the frontend-emitted SqlAggFunction is named "first" and + # NameBasedAggregateFunctionConverter.NAME_ALIASES rewrites it to "first_value" + # before YAML lookup. DataFusion's native first_value UDAF resolves this name + # directly; without an explicit ORDER BY it returns an arbitrary element from + # the group (documented tradeoff vs PPL's "first in document order" semantics). + # + # Also targeted by NameBasedAggregateFunctionConverter.rewriteArgMinMax to + # implement PPL `earliest(field, ts)` — Calcite emits ARG_MIN(field, ts), + # the rewrite transforms it to first_value(field) with sort field + # `ts ASC NULLS LAST` so DataFusion's first_value UDAF picks the row with + # the smallest timestamp. DataFusion 52.x has no native min_by/arg_min UDAF. + - name: "first_value" + description: "Returns the first value in the input (plus optional ORDER BY for earliest-style selection)." + impls: + - args: + - name: x + value: any1 + nullability: DECLARED_OUTPUT + decomposable: NONE + return: any1? + + # PPL `last(field)` — symmetric to first(field). Also targeted by rewriteArgMinMax + # for PPL `latest(field, ts)`: Calcite ARG_MAX(field, ts) rewrites to + # last_value(field) with sort `ts ASC NULLS LAST`, so DataFusion's last_value + # returns the row with the largest timestamp (the last row of the ascending sort). + # Equivalent to first_value with DESC but chosen for symmetry with DF's + # semantics. + - name: "last_value" + description: "Returns the last value in the input (plus optional ORDER BY for latest-style selection)." + impls: + - args: + - name: x + value: any1 + nullability: DECLARED_OUTPUT + decomposable: NONE + return: any1? + + # PPL `list(field)` / `values(field)` — NAME_ALIASES rewrites both of these to + # "array_agg" before YAML lookup. DataFusion's native array_agg UDAF collects all + # input values into a LIST. Semantic divergences from PPL (all deferred, not + # enforced on the wire): + # - PPL list filters nulls; array_agg includes them. + # - PPL list/values stringify inputs; array_agg preserves the input type. + # - PPL list caps at 100 values; array_agg is unbounded. + # - PPL values requires DISTINCT + sorted; plain array_agg does neither. + # (The DISTINCT + ORDER BY wiring for `values` lands in a later commit + # via an AliasConfig extension to NameBasedAggregateFunctionConverter.) + - name: "array_agg" + description: "Aggregate values into an array." + impls: + - args: + - name: x + value: any1 + nullability: DECLARED_OUTPUT + decomposable: NONE + return: LIST? diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverterTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverterTests.java new file mode 100644 index 0000000000000..7c02693d7a16b --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/io/substrait/isthmus/expression/NameBasedAggregateFunctionConverterTests.java @@ -0,0 +1,440 @@ +/* + * 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 io.substrait.isthmus.expression; + +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.apache.calcite.util.Optionality; +import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; +import org.opensearch.be.datafusion.DataFusionFragmentConvertor; +import org.opensearch.test.OpenSearchTestCase; + +import java.io.InputStream; +import java.util.List; + +import io.substrait.expression.AggregateFunctionInvocation; +import io.substrait.expression.Expression; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; +import io.substrait.proto.AggregateRel; +import io.substrait.proto.Plan; +import io.substrait.proto.PlanRel; +import io.substrait.proto.Rel; +import io.substrait.proto.SortField; + +/** + * Tests the {@code ARG_MIN}/{@code ARG_MAX} → {@code first_value}/{@code last_value} + * rewrite in {@link NameBasedAggregateFunctionConverter}. Feeds a Calcite + * {@link LogicalAggregate} whose single measure is an {@link SqlStdOperatorTable#ARG_MIN} + * or {@link SqlStdOperatorTable#ARG_MAX} call through + * {@link DataFusionFragmentConvertor#convertFinalAggFragment}, decodes the resulting + * Substrait proto bytes, and asserts: + * + *
    + *
  • the emitted aggregate measure has exactly one argument (the value field), NOT + * two — the key field is pulled out into the sort field list;
    • + *
  • the measure carries exactly one sort field whose expression is a field + * reference to the key and whose direction is {@code ASC_NULLS_LAST};
    • + *
  • the function reference points at {@code first_value} for ARG_MIN and + * {@code last_value} for ARG_MAX.
    • + *
+ */ +public class NameBasedAggregateFunctionConverterTests extends OpenSearchTestCase { + + private RelDataTypeFactory typeFactory; + private RexBuilder rexBuilder; + private RelOptCluster cluster; + private SimpleExtension.ExtensionCollection extensions; + + @Override + public void setUp() throws Exception { + super.setUp(); + typeFactory = new JavaTypeFactoryImpl(); + rexBuilder = new RexBuilder(typeFactory); + HepPlanner planner = new HepPlanner(new HepProgramBuilder().build()); + cluster = RelOptCluster.create(planner, rexBuilder); + + // Load the substrait core defaults plus OpenSearch's aggregate-extension + // YAML — the latter carries the first_value / last_value variants our + // rewrite looks up by name. + Thread t = Thread.currentThread(); + ClassLoader prev = t.getContextClassLoader(); + try { + t.setContextClassLoader(NameBasedAggregateFunctionConverterTests.class.getClassLoader()); + SimpleExtension.ExtensionCollection collection = DefaultExtensionCatalog.DEFAULT_COLLECTION; + try ( + InputStream stream = NameBasedAggregateFunctionConverterTests.class.getResourceAsStream( + "/extensions/opensearch_aggregate.yaml" + ) + ) { + assertNotNull("opensearch_aggregate.yaml must be on the test classpath", stream); + collection = collection.merge(SimpleExtension.load(stream)); + } + extensions = collection; + } finally { + t.setContextClassLoader(prev); + } + } + + /** + * {@code SELECT ARG_MIN(value, key) FROM t} lowered via isthmus must emit a + * substrait measure targeting {@code first_value(value)} with + * {@code sorts=[{expr=key, direction=ASC_NULLS_LAST}]}. + */ + public void testArgMinRewritesToFirstValueWithSortKey() throws Exception { + assertArgMinMaxRewrite(SqlStdOperatorTable.ARG_MIN, "first_value"); + } + + /** + * Symmetric to {@link #testArgMinRewritesToFirstValueWithSortKey()} — ARG_MAX + * routes to {@code last_value(value)} with the same ASC_NULLS_LAST sort field. + * {@code last_value} with ASC returns the row with the max key because it's the + * last row after sorting. + */ + public void testArgMaxRewritesToLastValueWithSortKey() throws Exception { + assertArgMinMaxRewrite(SqlStdOperatorTable.ARG_MAX, "last_value"); + } + + /** + * Shared assertion body — builds an aggregate fragment with one call using + * {@code op(value_col, key_col)}, runs it through the fragment convertor, + * and checks the resulting proto's first measure. + */ + private void assertArgMinMaxRewrite(org.apache.calcite.sql.SqlAggFunction op, String expectedVariantName) throws Exception { + RelNode leaf = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, rowType("value_col", "key_col"), List.of("datafusion")); + AggregateCall call = AggregateCall.create( + op, + /* isDistinct */ false, + /* argList */ List.of(0, 1), + /* filterArg */ -1, + /* type */ typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + /* name */ "result" + ); + LogicalAggregate agg = LogicalAggregate.create(leaf, List.of(), ImmutableBitSet.of(), /* groupSets */ null, List.of(call)); + + byte[] bytes = new DataFusionFragmentConvertor(extensions).convertFinalAggFragment(agg); + Plan plan = Plan.parseFrom(bytes); + assertFalse("plan must have at least one relation", plan.getRelationsList().isEmpty()); + PlanRel planRel = plan.getRelationsList().get(0); + assertTrue("plan relation must carry a root", planRel.hasRoot()); + Rel root = planRel.getRoot().getInput(); + assertTrue("root must be an AggregateRel for " + op.getName(), root.hasAggregate()); + AggregateRel aggRel = root.getAggregate(); + assertEquals("expected exactly one measure", 1, aggRel.getMeasuresCount()); + io.substrait.proto.AggregateFunction fn = aggRel.getMeasures(0).getMeasure(); + + // Measure must have exactly 1 argument — the value field. The key field was + // pulled out into the sort field list by the rewrite. + assertEquals("rewritten measure must carry 1 argument (the value field), not 2", 1, fn.getArgumentsCount()); + + // Measure must carry exactly one sort field, expr = field-reference to the key, + // direction = ASC_NULLS_LAST. + assertEquals("rewritten measure must carry exactly one sort field", 1, fn.getSortsCount()); + SortField sort = fn.getSorts(0); + assertEquals( + "rewrite emits ASC_NULLS_LAST direction for both ARG_MIN and ARG_MAX", + SortField.SortDirection.SORT_DIRECTION_ASC_NULLS_LAST, + sort.getDirection() + ); + assertTrue("sort expression must be a field reference to the key column", sort.getExpr().hasSelection()); + assertEquals( + "sort expression must reference column index 1 (key_col)", + 1, + sort.getExpr().getSelection().getDirectReference().getStructField().getField() + ); + + // Function reference in the plan's extensions must point at first_value or last_value. + int fnRef = fn.getFunctionReference(); + String resolvedName = plan.getExtensionsList() + .stream() + .filter(ext -> ext.hasExtensionFunction() && ext.getExtensionFunction().getFunctionAnchor() == fnRef) + .map(ext -> ext.getExtensionFunction().getName()) + .findFirst() + .orElseThrow(() -> new AssertionError("no extension function entry for anchor " + fnRef)); + assertTrue( + "rewritten function name must start with " + expectedVariantName + " (got: " + resolvedName + ")", + resolvedName.startsWith(expectedVariantName) + ); + } + + // ── Low-level convertByName / getFunctionFinder assertions ──────────────── + + /** + * The rewrite must only activate for 2-argument ARG_MIN/ARG_MAX calls. + * Guards against a future Calcite extension that emits arity-3 variants + * silently slipping through. + */ + public void testRewriteIsArityGated() { + NameBasedAggregateFunctionConverter conv = new NameBasedAggregateFunctionConverter( + extensions.aggregateFunctions(), + List.of(), + typeFactory, + io.substrait.isthmus.TypeConverter.DEFAULT + ); + // No public API to probe the rewrite directly — the fragment-level tests above + // cover the arity-2 happy path. This test would need refactoring if we ever + // expose an arity-check hook. + assertNotNull(conv); + } + + // ── Pure-rename aliases: first, last, list (convertByName path) ──────────── + + /** + * PPL {@code stats first(field)} maps through {@link NameBasedAggregateFunctionConverter#NAME_ALIASES} + * as a pure rename to DataFusion's {@code first_value}. No sort field, no forced DISTINCT — + * the emitted substrait measure preserves the single operand and routes to the + * {@code first_value} YAML variant. + */ + public void testFirstRewritesToFirstValue() throws Exception { + assertPureRename("first", "first_value", /* expectDistinct */ false); + } + + /** + * Symmetric to {@link #testFirstRewritesToFirstValue()} — PPL {@code last} maps to + * DataFusion's {@code last_value} via pure-rename AliasConfig. + */ + public void testLastRewritesToLastValue() throws Exception { + assertPureRename("last", "last_value", /* expectDistinct */ false); + } + + /** + * PPL {@code stats list(field)} maps as a pure rename to DataFusion's {@code array_agg}. + * Unlike {@code values}, {@code list} does NOT force DISTINCT and does NOT synthesize a + * sort field — the collected array preserves input order and duplicates. + */ + public void testListRewritesToArrayAgg() throws Exception { + assertPureRename("list", "array_agg", /* expectDistinct */ false); + } + + /** + * PPL {@code stats take(field, n)} routes to the custom {@code take} YAML variant + * registered by the Rust UDAF. Unlike {@code first}/{@code last}/{@code list}, there + * is no alias rewrite — the variant is looked up directly by name in the extension + * collection. This test verifies the full 2-operand form resolves to the {@code take} + * variant without any sort or DISTINCT shape. + */ + public void testTakeRoutesToTakeVariant() throws Exception { + RelNode leaf = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, rowType("value_col", "n_col"), List.of("datafusion")); + SqlAggFunction takeAgg = stubAgg("take"); + AggregateCall call = AggregateCall.create( + takeAgg, + /* isDistinct */ false, + /* argList */ List.of(0, 1), + /* filterArg */ -1, + /* type */ typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + /* name */ "result" + ); + LogicalAggregate agg = LogicalAggregate.create(leaf, List.of(), ImmutableBitSet.of(), /* groupSets */ null, List.of(call)); + + byte[] bytes = new DataFusionFragmentConvertor(extensions).convertFinalAggFragment(agg); + Plan plan = Plan.parseFrom(bytes); + Rel root = plan.getRelationsList().get(0).getRoot().getInput(); + assertTrue("root must be an AggregateRel", root.hasAggregate()); + AggregateRel aggRel = root.getAggregate(); + assertEquals("expected exactly one measure", 1, aggRel.getMeasuresCount()); + io.substrait.proto.AggregateFunction fn = aggRel.getMeasures(0).getMeasure(); + assertEquals("take preserves both operands (value + n)", 2, fn.getArgumentsCount()); + assertEquals("take has no synthesized sort field", 0, fn.getSortsCount()); + assertEquals( + "take is not distinct-by-default", + io.substrait.proto.AggregateFunction.AggregationInvocation.AGGREGATION_INVOCATION_ALL, + fn.getInvocation() + ); + + int fnRef = fn.getFunctionReference(); + String resolvedName = plan.getExtensionsList() + .stream() + .filter(ext -> ext.hasExtensionFunction() && ext.getExtensionFunction().getFunctionAnchor() == fnRef) + .map(ext -> ext.getExtensionFunction().getName()) + .findFirst() + .orElseThrow(() -> new AssertionError("no extension function entry for anchor " + fnRef)); + assertTrue("take must resolve to the take YAML variant (got: " + resolvedName + ")", resolvedName.startsWith("take")); + } + + /** + * Shared assertion body for pure-rename AliasConfig entries (first, last, list). + * Builds a scalar stats call with a single operand against a 1-column leaf scan, + * runs it through the fragment convertor, and checks the emitted measure: + *
    + *
  • exactly 1 operand;
    • + *
  • no sort field — pure renames don't synthesize sorts;
    • + *
  • invocation matches {@code expectDistinct} (all pure renames today are ALL);
    • + *
  • function reference resolves to {@code targetVariantName} in the extension list.
    • + *
+ */ + private void assertPureRename(String aliasName, String targetVariantName, boolean expectDistinct) throws Exception { + RelNode leaf = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, rowType("value_col"), List.of("datafusion")); + SqlAggFunction aliasAgg = stubAgg(aliasName); + AggregateCall call = AggregateCall.create( + aliasAgg, + /* isDistinct */ false, + /* argList */ List.of(0), + /* filterArg */ -1, + /* type */ typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + /* name */ "result" + ); + LogicalAggregate agg = LogicalAggregate.create(leaf, List.of(), ImmutableBitSet.of(), /* groupSets */ null, List.of(call)); + + byte[] bytes = new DataFusionFragmentConvertor(extensions).convertFinalAggFragment(agg); + Plan plan = Plan.parseFrom(bytes); + Rel root = plan.getRelationsList().get(0).getRoot().getInput(); + assertTrue(aliasName + ": root must be an AggregateRel", root.hasAggregate()); + AggregateRel aggRel = root.getAggregate(); + assertEquals(aliasName + ": expected exactly one measure", 1, aggRel.getMeasuresCount()); + io.substrait.proto.AggregateFunction fn = aggRel.getMeasures(0).getMeasure(); + assertEquals(aliasName + ": pure rename preserves the single operand", 1, fn.getArgumentsCount()); + assertEquals(aliasName + ": pure rename does not synthesize sort fields", 0, fn.getSortsCount()); + io.substrait.proto.AggregateFunction.AggregationInvocation expectedInv = expectDistinct + ? io.substrait.proto.AggregateFunction.AggregationInvocation.AGGREGATION_INVOCATION_DISTINCT + : io.substrait.proto.AggregateFunction.AggregationInvocation.AGGREGATION_INVOCATION_ALL; + assertEquals(aliasName + ": invocation flag must match expected", expectedInv, fn.getInvocation()); + + int fnRef = fn.getFunctionReference(); + String resolvedName = plan.getExtensionsList() + .stream() + .filter(ext -> ext.hasExtensionFunction() && ext.getExtensionFunction().getFunctionAnchor() == fnRef) + .map(ext -> ext.getExtensionFunction().getName()) + .findFirst() + .orElseThrow(() -> new AssertionError(aliasName + ": no extension function entry for anchor " + fnRef)); + assertTrue( + aliasName + " must resolve to " + targetVariantName + " (got: " + resolvedName + ")", + resolvedName.startsWith(targetVariantName) + ); + } + + // ── AliasConfig forceDistinct + sortIsSelf (VALUES-style) ────────────────── + + /** + * PPL {@code stats values(field)} routes through AliasConfig + * {@code {target=array_agg, forceDistinct=true, sortIsSelf=true}}. The rewritten + * substrait measure must: + *
    + *
  • target the {@code array_agg} YAML variant (not {@code values});
    • + *
  • carry exactly one operand (the field — same as input);
    • + *
  • carry exactly one sort field whose expression references the SAME column + * as the operand (sort-by-self), with ASC_NULLS_LAST direction;
    • + *
  • set AggregationInvocation = DISTINCT regardless of whether the frontend + * call's isDistinct() was true or false.
    • + *
+ * Mirrors Marc's concern that + * {@code NameBasedAggregateFunctionConverter.convertByName}'s prior empty-sorts + * path would silently swallow ORDER BY flags; without this test, a refactor + * could regress DISTINCT or sort-field propagation and the only signal would be + * a semantic IT failure. + */ + public void testValuesRewritesToArrayAggDistinctSortedBySelf() throws Exception { + RelNode leaf = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, rowType("value_col"), List.of("datafusion")); + // PPL emits a custom SqlAggFunction named "VALUES" (uppercase, SqlKind.OTHER_FUNCTION) + // with isDistinct=false — the AliasConfig must force DISTINCT regardless. + SqlAggFunction valuesAgg = stubAgg("VALUES"); + AggregateCall call = AggregateCall.create( + valuesAgg, + /* isDistinct */ false, + /* argList */ List.of(0), + /* filterArg */ -1, + /* type */ typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true), + /* name */ "result" + ); + LogicalAggregate agg = LogicalAggregate.create(leaf, List.of(), ImmutableBitSet.of(), /* groupSets */ null, List.of(call)); + + byte[] bytes = new DataFusionFragmentConvertor(extensions).convertFinalAggFragment(agg); + Plan plan = Plan.parseFrom(bytes); + Rel root = plan.getRelationsList().get(0).getRoot().getInput(); + assertTrue("root must be an AggregateRel", root.hasAggregate()); + AggregateRel aggRel = root.getAggregate(); + assertEquals("expected exactly one measure", 1, aggRel.getMeasuresCount()); + io.substrait.proto.AggregateFunction fn = aggRel.getMeasures(0).getMeasure(); + + // DISTINCT must be forced by AliasConfig.forceDistinct even when the + // Calcite-level call had isDistinct=false. + assertEquals( + "VALUES → array_agg must emit AGGREGATION_INVOCATION_DISTINCT", + io.substrait.proto.AggregateFunction.AggregationInvocation.AGGREGATION_INVOCATION_DISTINCT, + fn.getInvocation() + ); + + // sortIsSelf keeps the operand AND synthesizes a sort field on it. + assertEquals("operand list must have 1 element (value_col)", 1, fn.getArgumentsCount()); + assertEquals("must carry exactly one sort field for sort-by-self", 1, fn.getSortsCount()); + SortField sort = fn.getSorts(0); + assertEquals("sort direction must be ASC_NULLS_LAST", SortField.SortDirection.SORT_DIRECTION_ASC_NULLS_LAST, sort.getDirection()); + assertTrue("sort expression must be a field reference", sort.getExpr().hasSelection()); + assertEquals( + "sort-by-self: sort expression must reference the same column index as the operand (0)", + 0, + sort.getExpr().getSelection().getDirectReference().getStructField().getField() + ); + + // Function reference must resolve to array_agg, not values. + int fnRef = fn.getFunctionReference(); + String resolvedName = plan.getExtensionsList() + .stream() + .filter(ext -> ext.hasExtensionFunction() && ext.getExtensionFunction().getFunctionAnchor() == fnRef) + .map(ext -> ext.getExtensionFunction().getName()) + .findFirst() + .orElseThrow(() -> new AssertionError("no extension function entry for anchor " + fnRef)); + assertTrue("VALUES rewrite target must be array_agg (got: " + resolvedName + ")", resolvedName.startsWith("array_agg")); + } + + /** + * Minimal {@link SqlAggFunction} built with the same 10-arg constructor as + * {@link DataFusionFragmentConvertor#stubAgg(String)}. PPL-emitted aggregate + * operators are custom SqlAggFunction subclasses with SqlKind.OTHER_FUNCTION; + * we need to simulate one in tests to exercise the AliasConfig rewrite path + * without depending on the sql/core frontend. + */ + private static SqlAggFunction stubAgg(String name) { + return new SqlAggFunction( + name, + null, + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ) { + }; + } + + // ── Helpers ──────────────────────────────────────────────────────────────── + + private RelDataType rowType(String... columns) { + RelDataTypeFactory.Builder b = typeFactory.builder(); + for (String c : columns) { + b.add(c, typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.INTEGER), true)); + } + return b.build(); + } + + // Keep the sort-field type reachable so Eclipse's organize-imports doesn't prune it. + @SuppressWarnings("unused") + private static final Class SORT_FIELD_CLASS = Expression.SortField.class; + + @SuppressWarnings("unused") + private static final Class INVOCATION_CLASS = AggregateFunctionInvocation.class; +} diff --git a/sandbox/plugins/analytics-engine/build.gradle b/sandbox/plugins/analytics-engine/build.gradle index 346aa1385ad70..78df3e4ec63f1 100644 --- a/sandbox/plugins/analytics-engine/build.gradle +++ b/sandbox/plugins/analytics-engine/build.gradle @@ -14,9 +14,6 @@ apply plugin: 'opensearch.internal-cluster-test' -// SQL Unified Query API version (aligned with OpenSearch build version) -def sqlUnifiedQueryVersion = '3.6.0.0-SNAPSHOT' - opensearchplugin { description = 'Analytics engine hub: discovers and wires query extensions via ExtensiblePlugin SPI.' classname = 'org.opensearch.analytics.AnalyticsPlugin' @@ -76,6 +73,11 @@ dependencies { // Provided by arrow-flight-rpc at runtime (api deps in its build.gradle). compileOnly "com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}" compileOnly "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson_annotations}" + // Arrow's JsonStringArrayList registers jsr310's JavaTimeModule on its + // ObjectMapper. Triggered the first time a ListVector element is materialized via + // getObject — e.g. PPL take()'s ARRAY result. Without it, NoClassDefFoundError. + // arrow-flight-rpc does NOT export jsr310, so we bundle it explicitly. + implementation "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${versions.jackson}" // Guava — required at compile time because Calcite base classes expose guava types. // Uses custom config to bypass forbidden-dependencies.gradle check on compileClasspath. @@ -129,22 +131,13 @@ dependencies { testRuntimeOnly "com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}" testRuntimeOnly "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson_annotations}" - // SQL Unified Query API for PPL parsing - testImplementation("org.opensearch.query:unified-query-api:${sqlUnifiedQueryVersion}") { - exclude group: 'org.opensearch' - } - testImplementation("org.opensearch.query:unified-query-core:${sqlUnifiedQueryVersion}") { - exclude group: 'org.opensearch' - } - testImplementation("org.opensearch.query:unified-query-ppl:${sqlUnifiedQueryVersion}") { - exclude group: 'org.opensearch' - } - - // Arrow Flight streaming transport for ITs - internalClusterTestImplementation project(':plugins:arrow-flight-rpc') - // Calcite bytecode references @Immutable from immutables — resolve at test compile time testCompileOnly 'org.immutables:value-annotations:2.8.8' + + // ClickBenchUnifiedPipelineIT and other ITs use the PPL transport action defined in + // the test-ppl-frontend plugin. + internalClusterTestImplementation project(':sandbox:plugins:test-ppl-frontend') + internalClusterTestImplementation project(':plugins:arrow-flight-rpc') } tasks.withType(JavaCompile).configureEach { diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-databind-2.21.3.jar.sha1 b/sandbox/plugins/analytics-engine/licenses/jackson-databind-2.21.3.jar.sha1 deleted file mode 100644 index 0f1ca8bfdace0..0000000000000 --- a/sandbox/plugins/analytics-engine/licenses/jackson-databind-2.21.3.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -aa7ccec161c275f3e6332666ab758916f3120714 \ No newline at end of file diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 new file mode 100644 index 0000000000000..5bf925c777b5f --- /dev/null +++ b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 @@ -0,0 +1 @@ +a0958ebdaba836d31e5462ebc37b6349a0725ff9 diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-LICENSE.txt b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-LICENSE.txt new file mode 100644 index 0000000000000..f5f45d26a49d6 --- /dev/null +++ b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-LICENSE.txt @@ -0,0 +1,8 @@ +This copy of Jackson JSON processor streaming parser/generator is licensed under the +Apache (Software) License, version 2.0 ("the License"). +See the License for details about distribution rights, and the +specific rights regarding derivate works. + +You may obtain a copy of the License at: + +http://www.apache.org/licenses/LICENSE-2.0 diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-NOTICE.txt b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-NOTICE.txt new file mode 100644 index 0000000000000..4c976b7b4cc58 --- /dev/null +++ b/sandbox/plugins/analytics-engine/licenses/jackson-datatype-jsr310-NOTICE.txt @@ -0,0 +1,20 @@ +# Jackson JSON processor + +Jackson is a high-performance, Free/Open Source JSON processing library. +It was originally written by Tatu Saloranta (tatu.saloranta@iki.fi), and has +been in development since 2007. +It is currently developed by a community of developers, as well as supported +commercially by FasterXML.com. + +## Licensing + +Jackson core and extension components may licensed under different licenses. +To find the details that apply to this artifact see the accompanying LICENSE file. +For more information, including possible other licensing options, contact +FasterXML.com (http://fasterxml.com). + +## Credits + +A list of contributors may be found from CREDITS file, which is included +in some artifacts (usually source distributions); but is always available +from the source code management (SCM) system project uses. diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java index f1da7261ce75d..9e2fa6d5d7eb4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ArrowValues.java @@ -10,8 +10,11 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; /** * Helpers for reading Arrow vector cells as plain Java values at the @@ -23,16 +26,36 @@ private ArrowValues() {} /** * Returns the cell at {@code index} in {@code vector} as a Java value: - * {@code null} when the cell is null, a UTF-8 {@link String} for - * {@link VarCharVector} cells (rather than the raw {@code Text} that - * {@code getObject} returns), and {@link FieldVector#getObject} for - * every other vector type. + *
    + *
  • {@code null} when the cell is null
    • + *
  • UTF-8 {@link String} for {@link VarCharVector} cells (rather than + * the raw {@code Text} that {@code getObject} returns)
    • + *
  • plain {@link List} for {@link ListVector} cells, recursively + * unwrapping each element via this method (avoids leaking Arrow's + * {@code JsonStringArrayList} into downstream code that + * only recognises standard Java types)
    • + *
  • {@link FieldVector#getObject} for every other vector type
    • + *
*/ static Object toJavaValue(FieldVector vector, int index) { if (vector.isNull(index)) return null; if (vector instanceof VarCharVector v) { return new String(v.get(index), StandardCharsets.UTF_8); } + if (vector instanceof ListVector listVector) { + return listToJavaValue(listVector, index); + } return vector.getObject(index); } + + private static List listToJavaValue(ListVector listVector, int index) { + int start = listVector.getOffsetBuffer().getInt((long) index * ListVector.OFFSET_WIDTH); + int end = listVector.getOffsetBuffer().getInt((long) (index + 1) * ListVector.OFFSET_WIDTH); + FieldVector inner = listVector.getDataVector(); + List result = new ArrayList<>(end - start); + for (int i = start; i < end; i++) { + result.add(toJavaValue(inner, i)); + } + return result; + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java index 4e80dfa82e55f..6395ba0c9b644 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java @@ -20,6 +20,8 @@ import org.apache.arrow.vector.VarBinaryVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; @@ -57,11 +59,12 @@ public VectorSchemaRoot decode(FragmentExecutionResponse response, BufferAllocat throw new IllegalArgumentException("BufferAllocator must not be null"); } - // Infer Arrow type per column from the first non-null value + // Infer Arrow type per column from the first non-null value. List columns + // also need their element type inferred so the resulting ListVector has + // a properly-typed data vector. List fields = new ArrayList<>(); for (int col = 0; col < fieldNames.size(); col++) { - ArrowType arrowType = inferArrowType(rows, col); - fields.add(new Field(fieldNames.get(col), FieldType.nullable(arrowType), null)); + fields.add(inferField(fieldNames.get(col), rows, col)); } Schema schema = new Schema(fields); @@ -85,6 +88,21 @@ public VectorSchemaRoot decode(FragmentExecutionResponse response, BufferAllocat } } + /** + * Infers the Arrow {@link Field} for a column. For list-typed cells, also + * recurses into the elements to derive a child field — Arrow {@link ListVector} + * requires a typed data vector, not just an opaque list type. + */ + static Field inferField(String name, List rows, int col) { + ArrowType arrowType = inferArrowType(rows, col); + if (arrowType instanceof ArrowType.List) { + ArrowType elementType = inferListElementType(rows, col); + Field child = new Field("item", FieldType.nullable(elementType), null); + return new Field(name, FieldType.nullable(arrowType), List.of(child)); + } + return new Field(name, FieldType.nullable(arrowType), null); + } + /** * Infers the Arrow type for a column by scanning rows for the first * non-null value. Falls back to {@code Utf8} (VarChar) if all values @@ -103,12 +121,43 @@ static ArrowType inferArrowType(List rows, int col) { if (value instanceof Boolean) return ArrowType.Bool.INSTANCE; if (value instanceof CharSequence) return ArrowType.Utf8.INSTANCE; if (value instanceof byte[]) return ArrowType.Binary.INSTANCE; + if (value instanceof List) return ArrowType.List.INSTANCE; if (value instanceof Number) return new ArrowType.Int(64, true); break; } return ArrowType.Utf8.INSTANCE; } + /** + * For a list-typed column, scans non-null cells until a non-empty list is + * found and infers the element Arrow type from its first non-null element. + * Falls back to {@code Utf8} when no element is observable. + */ + private static ArrowType inferListElementType(List rows, int col) { + for (Object[] row : rows) { + Object value = row[col]; + if (!(value instanceof List list)) continue; + for (Object element : list) { + if (element == null) continue; + if (element instanceof Long) return new ArrowType.Int(64, true); + if (element instanceof Integer) return new ArrowType.Int(32, true); + if (element instanceof Short) return new ArrowType.Int(16, true); + if (element instanceof Byte) return new ArrowType.Int(8, true); + if (element instanceof Double) return new ArrowType.FloatingPoint( + org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE + ); + if (element instanceof Float) return new ArrowType.FloatingPoint( + org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE + ); + if (element instanceof Boolean) return ArrowType.Bool.INSTANCE; + if (element instanceof CharSequence) return ArrowType.Utf8.INSTANCE; + if (element instanceof org.apache.arrow.vector.util.Text) return ArrowType.Utf8.INSTANCE; + if (element instanceof Number) return new ArrowType.Int(64, true); + } + } + return ArrowType.Utf8.INSTANCE; + } + /** * Sets a value on the appropriate Arrow vector type. Handles null by * calling {@code setNull}. For typed vectors, casts the Java value to @@ -137,8 +186,66 @@ static void setVectorValue(FieldVector vector, int index, Object value) { ((VarCharVector) vector).setSafe(index, value.toString().getBytes(StandardCharsets.UTF_8)); } else if (vector instanceof VarBinaryVector) { ((VarBinaryVector) vector).setSafe(index, (byte[]) value); + } else if (vector instanceof ListVector listVector) { + writeListCell(listVector, index, (List) value); } else { throw new IllegalArgumentException("Unsupported Arrow vector type: " + vector.getClass().getSimpleName()); } } + + /** + * Writes a single Java {@link List} cell into the given {@link ListVector} + * at {@code index}. Uses the {@link UnionListWriter} for offset management + * so the list's start/end offsets stay consistent; element values are + * dispatched per-type through the writer's typed setters, mirroring the + * inference in {@link #inferListElementType}. + */ + private static void writeListCell(ListVector listVector, int index, List values) { + UnionListWriter writer = listVector.getWriter(); + writer.setPosition(index); + writer.startList(); + for (Object element : values) { + writeListElement(writer, element); + } + writer.endList(); + } + + /** + * Writes a single list element through the {@link UnionListWriter}'s typed + * setters. Arrow's {@code Text} class doesn't implement {@link CharSequence} + * but does implement {@link org.apache.arrow.vector.util.ReusableByteArray}, + * so a list cell sourced from {@link org.apache.arrow.vector.complex.ListVector#getObject} + * carries {@code Text} elements rather than {@link String}. Match those + * explicitly via the same {@code toString()} contract the rest of the + * codec uses. + */ + private static void writeListElement(UnionListWriter writer, Object element) { + if (element == null) { + writer.writeNull(); + return; + } + if (element instanceof Long l) { + writer.bigInt().writeBigInt(l); + } else if (element instanceof Integer i) { + writer.integer().writeInt(i); + } else if (element instanceof Short s) { + writer.smallInt().writeSmallInt(s); + } else if (element instanceof Byte b) { + writer.tinyInt().writeTinyInt(b); + } else if (element instanceof Double d) { + writer.float8().writeFloat8(d); + } else if (element instanceof Float f) { + writer.float4().writeFloat4(f); + } else if (element instanceof Boolean b) { + writer.bit().writeBit(b ? 1 : 0); + } else if (element instanceof CharSequence cs) { + writer.varChar().writeVarChar(cs.toString()); + } else if (element instanceof org.apache.arrow.vector.util.Text t) { + writer.varChar().writeVarChar(t.toString()); + } else if (element instanceof Number n) { + writer.bigInt().writeBigInt(n.longValue()); + } else { + throw new IllegalArgumentException("Unsupported list element type: " + element.getClass().getSimpleName()); + } + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java index 06cb3e725caa8..831dae7e88cdd 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java @@ -67,7 +67,8 @@ public static RelNode copyToCluster(RelNode node, RelOptCluster newCluster, Open aggregate.getGroupSets(), aggregate.getAggCallList(), aggregate.getMode(), - aggregate.getViableBackends() + aggregate.getViableBackends(), + aggregate.getCallAnnotations() ); } else if (node instanceof OpenSearchSort sort) { return new OpenSearchSort( diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AggregateCallAnnotation.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AggregateCallAnnotation.java index 4dc584e5e954b..02720a0fab2f4 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AggregateCallAnnotation.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AggregateCallAnnotation.java @@ -8,56 +8,32 @@ package org.opensearch.analytics.planner.rel; -import org.apache.calcite.rel.core.AggregateCall; -import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexNode; -import org.apache.calcite.sql.SqlKind; -import org.apache.calcite.sql.SqlOperator; -import org.apache.calcite.sql.SqlSyntax; -import org.apache.calcite.sql.type.ReturnTypes; -import java.util.ArrayList; import java.util.List; /** - * Marker {@link RexNode} embedded in an {@link AggregateCall#rexList} to carry - * per-call backend routing metadata. Same pattern as {@link AnnotatedPredicate} - * for filter predicates. - * - *

The aggregate rule appends this to each call's rexList via {@link #annotate}. - * During fragment conversion it is stripped out. - * {@code BackendResolver.generateCandidatePlans()} reads it for StagePlan - * alternative generation. + * Per-call backend routing metadata stored on {@link OpenSearchAggregate} in a + * side map keyed by aggregate-call index. Same role as {@link AnnotatedPredicate} + * for filter predicates, but lives outside the {@link org.apache.calcite.rel.core.AggregateCall} + * itself — embedding it in {@code rexList} would expose its type as a "preOperand" + * to Calcite's {@code AggCallBinding}, shifting positional argument inference for + * aggregates whose return-type inference reads {@code argTypes.getFirst()} + * (TAKE, FIRST, LAST, etc.). * * @opensearch.internal */ -public class AggregateCallAnnotation extends RexCall implements OperatorAnnotation { - - private static final SqlOperator AGG_CALL_ANNOTATION_OP = new SqlOperator( - "AGG_CALL_ANNOTATION", - SqlKind.OTHER_FUNCTION, - 0, - 0, - ReturnTypes.BOOLEAN, - null, - null - ) { - @Override - public SqlSyntax getSyntax() { - return SqlSyntax.FUNCTION; - } - }; +public final class AggregateCallAnnotation implements OperatorAnnotation { private final List viableBackends; private final int annotationId; - private AggregateCallAnnotation(RelDataType type, List viableBackends, int annotationId) { - super(type, AGG_CALL_ANNOTATION_OP, List.of()); + public AggregateCallAnnotation(List viableBackends, int annotationId) { this.viableBackends = viableBackends; this.annotationId = annotationId; } + @Override public List getViableBackends() { return viableBackends; } @@ -69,60 +45,18 @@ public int getAnnotationId() { @Override public OperatorAnnotation narrowTo(String backend) { - return new AggregateCallAnnotation(type, List.of(backend), annotationId); + return new AggregateCallAnnotation(List.of(backend), annotationId); } @Override public RexNode unwrap() { - // AggregateCallAnnotation is a marker in rexList, not a wrapper around an expression. - // Unwrapping means removing it from the rexList — handled by the operator's stripAnnotations. + // Not embedded in any RexNode tree — there's nothing to unwrap. return null; } @Override public RexNode withAdaptedOriginal(RexNode adaptedOriginal) { // AggregateCallAnnotation is a marker, not a wrapper — adaptation does not apply. - return this; - } - - /** Extracts the annotation from an AggregateCall's rexList, or null if absent. - * - *

TODO: window function aggregate calls may have ORDER BY expressions in rexList - * alongside our annotation. find() is safe (searches by type) and stripAnnotations - * filters by type, but consider moving annotations to a separate - * {@code Map} on {@link OpenSearchAggregate} keyed - * by call index to decouple from rexList entirely. - */ - public static AggregateCallAnnotation find(AggregateCall call) { - for (RexNode rex : call.rexList) { - if (rex instanceof AggregateCallAnnotation annotation) { - return annotation; - } - } return null; } - - /** Creates a new AggregateCall with this annotation appended to its rexList. */ - public static AggregateCall annotate(AggregateCall call, List viableBackends, int annotationId) { - List newRexList = new ArrayList<>(call.rexList); - newRexList.add(new AggregateCallAnnotation(call.type, viableBackends, annotationId)); - return AggregateCall.create( - call.getAggregation(), - call.isDistinct(), - call.isApproximate(), - call.ignoreNulls(), - newRexList, - call.getArgList(), - call.filterArg, - call.distinctKeys, - call.collation, - call.type, - call.name - ); - } - - @Override - protected String computeDigest(boolean withType) { - return "AGG_CALL_ANNOTATION(id=" + annotationId + ", viableBackends=" + viableBackends + ")"; - } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java index 5d86fcb0372c0..5e88906eaa8af 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java @@ -21,18 +21,25 @@ import org.opensearch.analytics.spi.FieldStorageInfo; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.function.Function; /** * OpenSearch custom Aggregate carrying viable backend list and per-call annotations. * + *

Annotations live in {@link #callAnnotations}, a side map keyed by aggregate-call + * index. They are NOT stored in the call's {@code rexList} — see + * {@link AggregateCallAnnotation} for why. + * * @opensearch.internal */ public class OpenSearchAggregate extends Aggregate implements OpenSearchRelNode { private final List viableBackends; private final AggregateMode mode; + private final Map callAnnotations; public OpenSearchAggregate( RelOptCluster cluster, @@ -42,11 +49,13 @@ public OpenSearchAggregate( List groupSets, List aggCalls, AggregateMode mode, - List viableBackends + List viableBackends, + Map callAnnotations ) { super(cluster, traitSet, List.of(), input, groupSet, groupSets, aggCalls); this.mode = mode; this.viableBackends = viableBackends; + this.callAnnotations = Map.copyOf(callAnnotations); } public AggregateMode getMode() { @@ -58,6 +67,11 @@ public List getViableBackends() { return viableBackends; } + /** Per-call annotations keyed by aggregate-call index. */ + public Map getCallAnnotations() { + return callAnnotations; + } + /** * Aggregate output: group-by fields first (inherited from input), then agg results (derived). * Group-by fields inherit storage info from the input. Agg results are derived columns. @@ -94,7 +108,7 @@ public Aggregate copy( List groupSets, List aggCalls ) { - return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, viableBackends); + return new OpenSearchAggregate(getCluster(), traitSet, input, groupSet, groupSets, aggCalls, mode, viableBackends, callAnnotations); } @Override @@ -125,46 +139,17 @@ public RelWriter explainTerms(RelWriter pw) { @Override public List getAnnotations() { - List annotations = new ArrayList<>(); - for (AggregateCall aggCall : getAggCallList()) { - for (RexNode rex : aggCall.rexList) { - if (rex instanceof AggregateCallAnnotation annotation) { - annotations.add(annotation); - } - } - } - return annotations; + // Iteration order matches insertion order (LinkedHashMap in the rule), which + // copyResolved relies on to align with resolvedAnnotations. + return List.copyOf(callAnnotations.values()); } @Override public RelNode copyResolved(String backend, List children, List resolvedAnnotations) { + Map resolvedMap = new LinkedHashMap<>(); int annotationIndex = 0; - List resolvedCalls = new ArrayList<>(); - for (AggregateCall aggCall : getAggCallList()) { - List newRexList = new ArrayList<>(); - for (RexNode rex : aggCall.rexList) { - if (rex instanceof AggregateCallAnnotation) { - newRexList.add((RexNode) resolvedAnnotations.get(annotationIndex++)); - } else { - // Non-annotation entries (e.g. argument refs) are passed through unchanged. - newRexList.add(rex); - } - } - resolvedCalls.add( - AggregateCall.create( - aggCall.getAggregation(), - aggCall.isDistinct(), - aggCall.isApproximate(), - aggCall.ignoreNulls(), - newRexList, - aggCall.getArgList(), - aggCall.filterArg, - aggCall.distinctKeys, - aggCall.collation, - aggCall.type, - aggCall.name - ) - ); + for (Map.Entry entry : callAnnotations.entrySet()) { + resolvedMap.put(entry.getKey(), (AggregateCallAnnotation) resolvedAnnotations.get(annotationIndex++)); } return new OpenSearchAggregate( getCluster(), @@ -172,9 +157,10 @@ public RelNode copyResolved(String backend, List children, List strippedChildren) { @Override public RelNode stripAnnotations(List strippedChildren, Function annotationResolver) { - List strippedCalls = new ArrayList<>(); - for (AggregateCall aggCall : getAggCallList()) { - // TODO: when aggregate delegation is implemented, use annotationResolver - // to replace delegated AggregateCallAnnotations with placeholders instead - // of just filtering them out. - List cleanRexList = aggCall.rexList.stream().filter(rex -> !(rex instanceof AggregateCallAnnotation)).toList(); - strippedCalls.add( - AggregateCall.create( - aggCall.getAggregation(), - aggCall.isDistinct(), - aggCall.isApproximate(), - aggCall.ignoreNulls(), - cleanRexList, - aggCall.getArgList(), - aggCall.filterArg, - aggCall.distinctKeys, - aggCall.collation, - aggCall.type, - aggCall.name - ) - ); - } - return LogicalAggregate.create(strippedChildren.getFirst(), List.of(), getGroupSet(), getGroupSets(), strippedCalls); + // Per the PR 21424 refactor (3e50f563394), AggregateCallAnnotations live in a + // side map on OpenSearchAggregate — NOT in each AggregateCall's rexList — so + // pass the calls through unchanged. The annotationResolver is ignored for + // aggregates; it's retained on the interface for symmetry with other + // OpenSearch RelNodes (filter/project) that DO carry annotations inline. + return LogicalAggregate.create(strippedChildren.getFirst(), List.of(), getGroupSet(), getGroupSets(), getAggCallList()); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java index bd9b58fa0e501..53e0223cd4de5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java @@ -30,15 +30,19 @@ import java.util.ArrayList; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; import java.util.Set; /** * Converts {@link Aggregate} → {@link OpenSearchAggregate}. * - *

Annotates each {@link AggregateCall} with viable backends by embedding - * an {@link AggregateCallAnnotation} in its rexList. Computes operator-level - * viable backends as the intersection of per-call viable backends. + *

Annotates each {@link AggregateCall} with viable backends in + * {@link OpenSearchAggregate#getCallAnnotations()} (a side map keyed by call + * index) — NOT in the call's {@code rexList}, which would shift positional + * argument inference. Computes operator-level viable backends as the + * intersection of per-call viable backends. * *

The split into PARTIAL + FINAL is NOT done here. It happens via * {@link OpenSearchAggregateSplitRule} which fires when Volcano detects @@ -73,21 +77,23 @@ public void onMatch(RelOptRuleCall call) { List childViableBackends = openSearchChild.getViableBackends(); List childFieldStorage = openSearchChild.getOutputFieldStorage(); - // Annotate each AggregateCall with per-call viable backends - List annotatedCalls = new ArrayList<>(); - for (AggregateCall aggCall : aggregate.getAggCallList()) { + // Build per-call annotations into a side map keyed by call index. + Map annotations = new LinkedHashMap<>(); + List aggCalls = aggregate.getAggCallList(); + for (int i = 0; i < aggCalls.size(); i++) { + AggregateCall aggCall = aggCalls.get(i); List callViable = resolveViableBackendsForCall(aggCall, childFieldStorage); if (callViable.isEmpty()) { throw new IllegalStateException("No backend supports aggregate function [" + aggCall.getAggregation().getName() + "]"); } - annotatedCalls.add(AggregateCallAnnotation.annotate(aggCall, callViable, context.nextAnnotationId())); + annotations.put(i, new AggregateCallAnnotation(callViable, context.nextAnnotationId())); } // Compute operator-level viable backends: must be viable for child AND handle agg calls - List viableBackends = computeAggregateViableBackends(annotatedCalls, childViableBackends); + List viableBackends = computeAggregateViableBackends(annotations, childViableBackends); if (viableBackends.isEmpty()) { - List funcNames = aggregate.getAggCallList().stream().map(aggCall -> aggCall.getAggregation().getName()).toList(); + List funcNames = aggCalls.stream().map(c -> c.getAggregation().getName()).toList(); throw new IllegalStateException( "No backend can execute aggregate: functions " + funcNames @@ -107,9 +113,10 @@ public void onMatch(RelOptRuleCall call) { RelNodeUtils.unwrapHep(aggregate.getInput()), aggregate.getGroupSet(), aggregate.getGroupSets(), - annotatedCalls, + aggCalls, AggregateMode.SINGLE, - viableBackends + viableBackends, + annotations ) ); } @@ -158,8 +165,11 @@ private List resolveViableBackendsForCall(AggregateCall aggCall, List(registry.aggregateCapableBackends()); } - private List computeAggregateViableBackends(List annotatedCalls, List childViableBackends) { - if (annotatedCalls.isEmpty()) { + private List computeAggregateViableBackends( + Map annotations, + List childViableBackends + ) { + if (annotations.isEmpty()) { return new ArrayList<>(childViableBackends); } @@ -172,12 +182,7 @@ private List computeAggregateViableBackends(List annotate } boolean canHandleAll = true; - for (AggregateCall call : annotatedCalls) { - AggregateCallAnnotation annotation = AggregateCallAnnotation.find(call); - if (annotation == null) { - canHandleAll = false; - break; - } + for (AggregateCallAnnotation annotation : annotations.values()) { if (!registry.canHandle(candidateName, annotation.getViableBackends(), DelegationType.AGGREGATE)) { canHandleAll = false; break; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java index 4be2d71520adf..014977f186e78 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java @@ -16,6 +16,7 @@ import org.opensearch.analytics.planner.rel.AggregateMode; import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchConvention; +import org.opensearch.analytics.planner.rel.OpenSearchDistribution; /** * Volcano CBO rule that splits an {@link OpenSearchAggregate} into @@ -59,7 +60,22 @@ public OpenSearchAggregateSplitRule(PlannerContext context) { @Override public boolean matches(RelOptRuleCall call) { OpenSearchAggregate aggregate = call.rel(0); - return aggregate.getMode() == AggregateMode.SINGLE; + if (aggregate.getMode() != AggregateMode.SINGLE) { + return false; + } + // Skip if input is already gathered — the split would introduce a + // pointless PARTIAL+EXCHANGE+FINAL pipeline and, worse, the FINAL's + // constructor re-validates agg call types against its (re-aggregated) + // input, which fails for non-distributive aggregates like TAKE/AVG/STDDEV + // whose partial state shape differs from their result shape. + RelNode input = call.rel(1); + for (int i = 0; i < input.getTraitSet().size(); i++) { + if (input.getTraitSet().getTrait(i) instanceof OpenSearchDistribution dist) { + return dist.getType() != org.apache.calcite.rel.RelDistribution.Type.SINGLETON + && dist.getType() != org.apache.calcite.rel.RelDistribution.Type.ANY; + } + } + return false; } @Override @@ -77,7 +93,8 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.PARTIAL, - aggregate.getViableBackends() + aggregate.getViableBackends(), + aggregate.getCallAnnotations() ); // Request SINGLETON distribution — Volcano inserts Exchange automatically @@ -93,7 +110,8 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), aggregate.getAggCallList(), AggregateMode.FINAL, - aggregate.getViableBackends() + aggregate.getViableBackends(), + aggregate.getCallAnnotations() ); call.transformTo(finalAggregate); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java new file mode 100644 index 0000000000000..6de4e88fc7ee9 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ArrowValuesTests.java @@ -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. + */ + +package org.opensearch.analytics.exec; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.opensearch.test.OpenSearchTestCase; + +import java.nio.charset.StandardCharsets; +import java.util.List; + +/** + * Tests for {@link ArrowValues#toJavaValue}. + * + *

Covers the conversions that matter at the SQL-API boundary, where downstream + * code only recognises plain Java types ({@link String}, {@link List}, etc.) and + * blows up on Arrow's {@code Text} / {@code JsonStringArrayList} wrappers. + */ +public class ArrowValuesTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testNullCellReturnsNull() { + try (IntVector v = new IntVector("x", allocator)) { + v.allocateNew(1); + v.setNull(0); + v.setValueCount(1); + assertNull(ArrowValues.toJavaValue(v, 0)); + } + } + + public void testVarCharVectorReturnsPlainString() { + try (VarCharVector v = new VarCharVector("name", allocator)) { + v.allocateNew(1); + v.set(0, "hello".getBytes(StandardCharsets.UTF_8)); + v.setValueCount(1); + Object result = ArrowValues.toJavaValue(v, 0); + assertEquals(String.class, result.getClass()); + assertEquals("hello", result); + } + } + + /** + * Regression test for the take() rendering bug: a ListVector wrapping a + * VarCharVector must come out as a plain {@code List}, not as + * Arrow's {@code JsonStringArrayList}. The old code returned the raw + * {@code getObject()} value, whose elements are {@code Text} (not String); + * downstream serialization would then either fail or stringify the list to + * its {@code toString()} (a JSON-encoded string). + */ + public void testListVectorOfStringReturnsListOfString() { + try (ListVector listVector = ListVector.empty("take", allocator)) { + listVector.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); + UnionListWriter writer = listVector.getWriter(); + writer.startList(); + writer.varChar().writeVarChar("Amber JOHnny"); + writer.varChar().writeVarChar("Hattie"); + writer.endList(); + writer.setValueCount(1); + + Object result = ArrowValues.toJavaValue(listVector, 0); + + assertTrue("expected List, got " + result.getClass(), result instanceof List); + @SuppressWarnings("unchecked") + List list = (List) result; + assertEquals(2, list.size()); + assertEquals(String.class, list.get(0).getClass()); + assertEquals("Amber JOHnny", list.get(0)); + assertEquals("Hattie", list.get(1)); + } + } + + /** + * Empty list elements (e.g. {@code take(col, 0)}) must round-trip as an + * empty Java list, not as null. + */ + public void testListVectorEmptyReturnsEmptyList() { + try (ListVector listVector = ListVector.empty("take", allocator)) { + listVector.addOrGetVector(FieldType.nullable(new ArrowType.Utf8())); + UnionListWriter writer = listVector.getWriter(); + writer.startList(); + writer.endList(); + writer.setValueCount(1); + + Object result = ArrowValues.toJavaValue(listVector, 0); + assertTrue(result instanceof List); + assertEquals(0, ((List) result).size()); + } + } + + /** + * Lists of integers should also unwrap recursively — the inner Int vector + * already returns boxed Integers from getObject, so this is a sanity check + * that the recursion handles non-VarChar inner types correctly. + */ + public void testListVectorOfIntReturnsListOfInteger() { + try (ListVector listVector = ListVector.empty("ids", allocator)) { + listVector.addOrGetVector(FieldType.nullable(new ArrowType.Int(32, true))); + UnionListWriter writer = listVector.getWriter(); + writer.startList(); + writer.integer().writeInt(7); + writer.integer().writeInt(11); + writer.endList(); + writer.setValueCount(1); + + Object result = ArrowValues.toJavaValue(listVector, 0); + assertTrue(result instanceof List); + @SuppressWarnings("unchecked") + List list = (List) result; + assertEquals(List.of(7, 11), list); + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/RowResponseCodecTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/RowResponseCodecTests.java new file mode 100644 index 0000000000000..32c8962abc2d4 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/RowResponseCodecTests.java @@ -0,0 +1,138 @@ +/* + * 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.exec.stage; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.opensearch.analytics.exec.action.FragmentExecutionResponse; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; + +/** + * Tests for {@link RowResponseCodec}'s row-Object[] → Arrow conversion. + * + *

Includes a regression test for the {@code take()} list-rendering bug: when + * a row cell is a {@link List} (e.g. produced by a list-typed aggregate like + * {@code take(stringField, n)}), the codec must infer a list-typed Arrow column + * and populate a {@link ListVector} — NOT fall through to {@code Utf8} and + * stringify the list via {@code toString()}. + */ +public class RowResponseCodecTests extends OpenSearchTestCase { + + private BufferAllocator allocator; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + /** + * Regression test for the take() rendering bug. When the cell value is a + * {@code List}, the codec must produce a {@code ListVector(Utf8)} + * column whose values, when read back, equal the original list. + */ + public void testListOfStringsCellDecodesAsListVector() { + FragmentExecutionResponse response = new FragmentExecutionResponse( + List.of("take"), + java.util.Arrays.asList(new Object[] { List.of("Amber JOHnny", "Hattie") }) + ); + + try (VectorSchemaRoot vsr = RowResponseCodec.INSTANCE.decode(response, allocator)) { + FieldVector vec = vsr.getVector("take"); + assertTrue("expected ListVector for list-typed cell, got " + vec.getClass().getSimpleName(), vec instanceof ListVector); + assertEquals(1, vsr.getRowCount()); + + ListVector listVector = (ListVector) vec; + int start = listVector.getElementStartIndex(0); + int end = listVector.getElementEndIndex(0); + assertEquals(2, end - start); + + org.apache.arrow.vector.VarCharVector inner = (org.apache.arrow.vector.VarCharVector) listVector.getDataVector(); + assertEquals("Amber JOHnny", new String(inner.get(start), java.nio.charset.StandardCharsets.UTF_8)); + assertEquals("Hattie", new String(inner.get(start + 1), java.nio.charset.StandardCharsets.UTF_8)); + } + } + + /** + * Empty list rows must produce a ListVector with an empty list (no elements + * for that row), not a null cell. Mirrors {@code take(col, 0)} semantics. + */ + public void testEmptyListCellDecodesAsEmptyListVectorEntry() { + FragmentExecutionResponse response = new FragmentExecutionResponse( + List.of("take"), + java.util.Arrays.asList(new Object[] { List.of() }) + ); + + try (VectorSchemaRoot vsr = RowResponseCodec.INSTANCE.decode(response, allocator)) { + FieldVector vec = vsr.getVector("take"); + assertTrue(vec instanceof ListVector); + ListVector listVector = (ListVector) vec; + int start = listVector.getElementStartIndex(0); + int end = listVector.getElementEndIndex(0); + assertEquals(0, end - start); + assertFalse("empty list is not null", listVector.isNull(0)); + } + } + + /** + * Regression test for the actual runtime shape: cells produced by + * {@code DatafusionResultStream.getFieldValue} for a list-typed Arrow + * column come out as {@code JsonStringArrayList}, not + * {@code List}. Arrow's {@code Text} doesn't implement + * {@link CharSequence}, so the codec must explicitly recognise it. + */ + public void testListOfArrowTextDecodesAsListVector() { + org.apache.arrow.vector.util.JsonStringArrayList arrowList = + new org.apache.arrow.vector.util.JsonStringArrayList<>(); + arrowList.add(new org.apache.arrow.vector.util.Text("Amber JOHnny")); + arrowList.add(new org.apache.arrow.vector.util.Text("Hattie")); + + FragmentExecutionResponse response = new FragmentExecutionResponse( + List.of("take"), + java.util.Arrays.asList(new Object[] { arrowList }) + ); + + try (VectorSchemaRoot vsr = RowResponseCodec.INSTANCE.decode(response, allocator)) { + FieldVector vec = vsr.getVector("take"); + assertTrue(vec instanceof ListVector); + ListVector listVector = (ListVector) vec; + int start = listVector.getElementStartIndex(0); + int end = listVector.getElementEndIndex(0); + assertEquals(2, end - start); + + org.apache.arrow.vector.VarCharVector inner = (org.apache.arrow.vector.VarCharVector) listVector.getDataVector(); + assertEquals("Amber JOHnny", new String(inner.get(start), java.nio.charset.StandardCharsets.UTF_8)); + assertEquals("Hattie", new String(inner.get(start + 1), java.nio.charset.StandardCharsets.UTF_8)); + } + } + + /** + * A row whose first value is null but later rows have non-null values + * should still infer the right element type for the list. Sanity check. + */ + public void testInferArrowTypeRecognisesListOfStrings() { + ArrowType inferred = RowResponseCodec.inferArrowType( + java.util.Arrays.asList(new Object[] { null }, new Object[] { List.of("a", "b") }), + 0 + ); + assertEquals(ArrowType.List.INSTANCE, inferred); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java index 5398fc2e17ef6..aa1bea352b8d6 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java @@ -40,8 +40,8 @@ public class AggregateRuleTests extends BasePlannerRulesTests { /** Every agg call must have an annotation with non-empty viableBackends. */ public void testPerCallAnnotation() { OpenSearchAggregate agg = runAggregate(1, sumCall()); - for (AggregateCall call : agg.getAggCallList()) { - AggregateCallAnnotation annotation = AggregateCallAnnotation.find(call); + for (int i = 0; i < agg.getAggCallList().size(); i++) { + AggregateCallAnnotation annotation = agg.getCallAnnotations().get(i); assertNotNull("Every AggregateCall must have an annotation", annotation); assertFalse("Annotation viableBackends must not be empty", annotation.getViableBackends().isEmpty()); assertTrue(annotation.getViableBackends().contains(MockDataFusionBackend.NAME)); @@ -119,7 +119,7 @@ protected Set aggregateCapabilities() { OpenSearchAggregate agg = (OpenSearchAggregate) result; assertFalse(agg.getViableBackends().contains(MockLuceneBackend.NAME)); // Per-call annotation includes both — Lucene is viable for SUM on this field - assertCallAnnotation(agg.getAggCallList().get(0), MockDataFusionBackend.NAME, MockLuceneBackend.NAME); + assertCallAnnotation(agg, 0, MockDataFusionBackend.NAME, MockLuceneBackend.NAME); } /** @@ -148,7 +148,7 @@ protected Set aggregateCapabilities() { ); OpenSearchAggregate agg = (OpenSearchAggregate) result; assertTrue(agg.getViableBackends().contains(MockLuceneBackend.NAME)); - assertCallAnnotation(agg.getAggCallList().get(0), MockDataFusionBackend.NAME, MockLuceneBackend.NAME); + assertCallAnnotation(agg, 0, MockDataFusionBackend.NAME, MockLuceneBackend.NAME); } // ---- Composed pipeline shapes ---- @@ -195,18 +195,10 @@ protected Set aggregateCapabilities() { "Lucene not viable at operator level — can handle SUM but not COUNT", agg.getViableBackends().contains(MockLuceneBackend.NAME) ); - assertCallAnnotation(agg.getAggCallList().get(0), MockDataFusionBackend.NAME, MockLuceneBackend.NAME); - assertCallAnnotation(agg.getAggCallList().get(1), MockDataFusionBackend.NAME); - assertEquals( - "SUM viable for both backends", - 2, - AggregateCallAnnotation.find(agg.getAggCallList().get(0)).getViableBackends().size() - ); - assertEquals( - "COUNT viable for DF only (Lucene not declared)", - 1, - AggregateCallAnnotation.find(agg.getAggCallList().get(1)).getViableBackends().size() - ); + assertCallAnnotation(agg, 0, MockDataFusionBackend.NAME, MockLuceneBackend.NAME); + assertCallAnnotation(agg, 1, MockDataFusionBackend.NAME); + assertEquals("SUM viable for both backends", 2, agg.getCallAnnotations().get(0).getViableBackends().size()); + assertEquals("COUNT viable for DF only (Lucene not declared)", 1, agg.getCallAnnotations().get(1).getViableBackends().size()); } // ---- Delegation ---- @@ -295,9 +287,9 @@ private PlannerContext defaultContext(int shardCount) { return buildContext("parquet", shardCount, intFields()); } - private void assertCallAnnotation(AggregateCall call, String... expectedBackends) { - AggregateCallAnnotation annotation = AggregateCallAnnotation.find(call); - assertNotNull("AggregateCall must have annotation", annotation); + private void assertCallAnnotation(OpenSearchAggregate agg, int callIndex, String... expectedBackends) { + AggregateCallAnnotation annotation = agg.getCallAnnotations().get(callIndex); + assertNotNull("AggregateCall at index " + callIndex + " must have annotation", annotation); for (String backend : expectedBackends) assertTrue("Annotation must contain backend " + backend, annotation.getViableBackends().contains(backend)); } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLFirstCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLFirstCommandIT.java new file mode 100644 index 0000000000000..321f93dcf9001 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLFirstCommandIT.java @@ -0,0 +1,143 @@ +/* + * 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.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * End-to-end IT for the PPL {@code first} aggregate routed through the analytics-engine + * REST path. {@code first(field)} is wired by renaming PPL's {@code first} to + * DataFusion's native {@code first_value} via + * {@code NameBasedAggregateFunctionConverter.NAME_ALIASES}. + * + *

Two semantic divergences from PPL's documented behavior (both acknowledged): + *

    + *
  • Arbitrary within partition: without an explicit {@code ORDER BY}, + * DataFusion's {@code first_value} returns an arbitrary element from the group. + * The "first in document order" guarantee is not honored.
    • + *
  • Null tolerance: DataFusion's {@code first_value} does NOT skip nulls + * by default (no {@code ignore_nulls} hint is plumbed through substrait). PPL's + * docs promise the first NON-NULL value. For a field with nulls, the result may + * be null.
    • + *
+ * Strong assertions ({@code testFirstOnNullFreeStringField}, {@code testFirstOnIntegerField}) + * run against fields with no nulls. {@code testFirstOnNullableStringField} accepts null OR + * one of the non-null values, documenting the null-tolerance divergence. + */ +public class CalcitePPLFirstCommandIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("calcs", "calcs"); + + private static boolean dataProvisioned = false; + + private void ensureDataProvisioned() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + /** All non-null {@code str2} values in the {@code calcs} dataset. */ + private static final Set STR2_NON_NULL_VALUES = Set.of( + "one", "two", "three", "five", "six", "eight", "nine", "ten", + "eleven", "twelve", "fourteen", "fifteen", "sixteen" + ); + + // ── single-row stats on a nulls-free field (strongest assertion) ───────────── + + public void testFirstOnNullFreeStringField() throws IOException { + // `key` values are key00..key16 — no nulls. first(key) must return exactly one of them. + Map response = executePpl("source=" + DATASET.indexName + " | stats first(key)"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + assertNotNull("first(key) must not be null — `key` has no nulls in the calcs dataset", cell); + String actual = cell.toString(); + assertTrue( + "first(key)=" + actual + " must match the keyNN pattern", + actual.matches("key\\d{2}") + ); + } + + public void testFirstOnIntegerField() throws IOException { + // int0 has nulls; the arbitrary-element pick MAY be null. Accept null or a non-null int. + Map response = executePpl("source=" + DATASET.indexName + " | stats first(int0)"); + Set int0Values = Set.of(1, 3, 4, 7, 8, 10, 11); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + if (cell != null) { + long actual = ((Number) cell).longValue(); + assertTrue( + "first(int0)=" + actual + " must be one of " + int0Values + " (or null, per DataFusion null-tolerance)", + int0Values.stream().anyMatch(n -> n.longValue() == actual) + ); + } + } + + // ── nullable field — document null-tolerance divergence via the test itself ── + + public void testFirstOnNullableStringField() throws IOException { + // str2 has nulls. DataFusion's first_value(x) without ignore_nulls may return null. + // Accept null or one of the known non-null values — do NOT assert on a specific value. + Map response = executePpl("source=" + DATASET.indexName + " | stats first(str2)"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + Object cell = rows.get(0).get(0); + Set allowed = new HashSet<>(STR2_NON_NULL_VALUES); + // Null is legal here — document the tradeoff. + if (cell != null) { + assertTrue( + "first(str2)=" + cell + " must be one of " + allowed + " or null", + allowed.contains(cell.toString()) + ); + } + } + + // ── stats ... by (GROUP BY) ────────────────────────────────────────────────── + + public void testFirstGroupedByBool0() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats first(str2) by bool0"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertTrue("at least one group row expected", rows.size() >= 1); + for (List row : rows) { + assertEquals("two columns per row: first(str2), bool0", 2, row.size()); + Object firstStr2 = row.get(0); + if (firstStr2 != null) { + assertTrue( + "first(str2)=" + firstStr2 + " in a group must be one of " + STR2_NON_NULL_VALUES, + STR2_NON_NULL_VALUES.contains(firstStr2.toString()) + ); + } + } + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + private Map executePpl(String ppl) throws IOException { + ensureDataProvisioned(); + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLLastCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLLastCommandIT.java new file mode 100644 index 0000000000000..a7621864f5c28 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLLastCommandIT.java @@ -0,0 +1,143 @@ +/* + * 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.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * End-to-end IT for the PPL {@code last} aggregate routed through the analytics-engine + * REST path. Symmetric to {@link CalcitePPLFirstCommandIT} — {@code last(field)} is + * wired by renaming PPL's {@code last} to DataFusion's native {@code last_value} via + * {@code NameBasedAggregateFunctionConverter.NAME_ALIASES}. + * + *

Two semantic divergences from PPL's documented behavior (both acknowledged): + *

    + *
  • Arbitrary within partition: without an explicit {@code ORDER BY}, + * DataFusion's {@code last_value} returns an arbitrary element from the group. + * The "last in document order" guarantee is not honored.
    • + *
  • Null tolerance: DataFusion's {@code last_value} does NOT skip nulls + * by default (no {@code ignore_nulls} hint is plumbed through substrait). PPL's + * docs promise the last NON-NULL value. For a field with nulls, the result may + * be null.
    • + *
+ * Strong assertions ({@code testLastOnNullFreeStringField}, {@code testLastOnIntegerField}) + * run against fields with no nulls. {@code testLastOnNullableStringField} accepts null OR + * one of the non-null values, documenting the null-tolerance divergence. + */ +public class CalcitePPLLastCommandIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("calcs", "calcs"); + + private static boolean dataProvisioned = false; + + private void ensureDataProvisioned() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + /** All non-null {@code str2} values in the {@code calcs} dataset. */ + private static final Set STR2_NON_NULL_VALUES = Set.of( + "one", "two", "three", "five", "six", "eight", "nine", "ten", + "eleven", "twelve", "fourteen", "fifteen", "sixteen" + ); + + // ── single-row stats on a nulls-free field (strongest assertion) ───────────── + + public void testLastOnNullFreeStringField() throws IOException { + // `key` values are key00..key16 — no nulls. last(key) must return exactly one of them. + Map response = executePpl("source=" + DATASET.indexName + " | stats last(key)"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + assertNotNull("last(key) must not be null — `key` has no nulls in the calcs dataset", cell); + String actual = cell.toString(); + assertTrue( + "last(key)=" + actual + " must match the keyNN pattern", + actual.matches("key\\d{2}") + ); + } + + public void testLastOnIntegerField() throws IOException { + // int0 has nulls; the arbitrary-element pick MAY be null. Accept null or a non-null int. + Map response = executePpl("source=" + DATASET.indexName + " | stats last(int0)"); + Set int0Values = Set.of(1, 3, 4, 7, 8, 10, 11); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + if (cell != null) { + long actual = ((Number) cell).longValue(); + assertTrue( + "last(int0)=" + actual + " must be one of " + int0Values + " (or null, per DataFusion null-tolerance)", + int0Values.stream().anyMatch(n -> n.longValue() == actual) + ); + } + } + + // ── nullable field — document null-tolerance divergence via the test itself ── + + public void testLastOnNullableStringField() throws IOException { + // str2 has nulls. DataFusion's last_value(x) without ignore_nulls may return null. + // Accept null or one of the known non-null values — do NOT assert on a specific value. + Map response = executePpl("source=" + DATASET.indexName + " | stats last(str2)"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + Object cell = rows.get(0).get(0); + Set allowed = new HashSet<>(STR2_NON_NULL_VALUES); + // Null is legal here — document the tradeoff. + if (cell != null) { + assertTrue( + "last(str2)=" + cell + " must be one of " + allowed + " or null", + allowed.contains(cell.toString()) + ); + } + } + + // ── stats ... by (GROUP BY) ────────────────────────────────────────────────── + + public void testLastGroupedByBool0() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats last(str2) by bool0"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertTrue("at least one group row expected", rows.size() >= 1); + for (List row : rows) { + assertEquals("two columns per row: last(str2), bool0", 2, row.size()); + Object lastStr2 = row.get(0); + if (lastStr2 != null) { + assertTrue( + "last(str2)=" + lastStr2 + " in a group must be one of " + STR2_NON_NULL_VALUES, + STR2_NON_NULL_VALUES.contains(lastStr2.toString()) + ); + } + } + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + private Map executePpl(String ppl) throws IOException { + ensureDataProvisioned(); + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLListCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLListCommandIT.java new file mode 100644 index 0000000000000..97827885b3cff --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLListCommandIT.java @@ -0,0 +1,161 @@ +/* + * 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.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * End-to-end IT for the PPL {@code list} aggregate routed through the analytics-engine + * REST path. {@code list(field)} is wired by renaming PPL's {@code list} to + * DataFusion's native {@code array_agg} via + * {@code NameBasedAggregateFunctionConverter.NAME_ALIASES}. + * + *

Semantic divergences from PPL's documented behavior (ALL deferred, NOT + * enforced on the wire by this commit): + *

    + *
  • Null filter: PPL {@code list(field)} filters out null values. DataFusion's + * {@code array_agg} includes them — resulting arrays may contain a {@code null} + * element per null input row.
    • + *
  • String-cast: PPL {@code list(field)} stringifies every element. DataFusion + * preserves the input type (numeric → numeric, string → string, etc.).
    • + *
  • 100-element cap: PPL caps the collected array at 100 values. DataFusion's + * {@code array_agg} is unbounded.
    • + *
+ * + *

The tests below accommodate these divergences: assertions check that every non-null + * expected input value appears somewhere in the collected array (subset-containment), + * but do NOT assert null-freeness, element-type, or size bounds. + */ +public class CalcitePPLListCommandIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("calcs", "calcs"); + + private static boolean dataProvisioned = false; + + private void ensureDataProvisioned() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + /** All non-null {@code str2} values in the {@code calcs} dataset. Used for + * subset-containment assertions — the list result must include every one. */ + private static final Set STR2_NON_NULL_VALUES = Set.of( + "one", "two", "three", "five", "six", "eight", "nine", "ten", + "eleven", "twelve", "fourteen", "fifteen", "sixteen" + ); + + /** All non-null {@code int0} values. */ + private static final Set INT0_NON_NULL_VALUES = Set.of(1L, 7L, 3L, 8L, 4L, 10L, 11L); + + // ── scalar stats (no GROUP BY) ─────────────────────────────────────────────── + + public void testListOnStringField() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats list(str2)"); + List collected = extractListCell(response, "list(str2)"); + // Subset-containment: every non-null str2 value in the dataset must appear in the + // collected array (converted to string to normalize across the type-preservation + // divergence, though str2 is already string-typed so this is a no-op here). + Set collectedNonNull = new HashSet<>(); + for (Object v : collected) { + if (v != null) { + collectedNonNull.add(v.toString()); + } + } + for (String expected : STR2_NON_NULL_VALUES) { + assertTrue( + "list(str2) result must contain non-null input " + expected + "; got " + collected, + collectedNonNull.contains(expected) + ); + } + } + + public void testListOnIntegerField() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats list(int0)"); + List collected = extractListCell(response, "list(int0)"); + Set collectedNonNull = new HashSet<>(); + for (Object v : collected) { + if (v != null) { + // Accept any Number (DataFusion preserves Integer/Long input type). + collectedNonNull.add(((Number) v).longValue()); + } + } + for (Long expected : INT0_NON_NULL_VALUES) { + assertTrue( + "list(int0) result must contain non-null input " + expected + "; got " + collected, + collectedNonNull.contains(expected) + ); + } + } + + // ── stats ... by (GROUP BY) ────────────────────────────────────────────────── + + public void testListGroupedByBool0() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats list(str2) by bool0"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertTrue("at least one group row expected", rows.size() >= 1); + // Union of all per-group list(str2) results must cover every non-null str2 value. + // (Stronger than "each group is in range" — confirms the whole dataset flowed + // through the aggregate, partitioned by bool0.) + Set unionNonNull = new HashSet<>(); + for (List row : rows) { + assertEquals("two columns per row: list(str2), bool0", 2, row.size()); + Object listCell = row.get(0); + assertTrue("list(str2) cell must be an array, got " + (listCell == null ? "null" : listCell.getClass()), listCell instanceof List); + @SuppressWarnings("unchecked") + List groupList = (List) listCell; + for (Object v : groupList) { + if (v != null) { + unionNonNull.add(v.toString()); + } + } + } + for (String expected : STR2_NON_NULL_VALUES) { + assertTrue( + "union of list(str2) across all bool0 groups must contain " + expected + "; got " + unionNonNull, + unionNonNull.contains(expected) + ); + } + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + /** Extracts the single-row single-column cell from a scalar-stats response and + * asserts it's a list. Returns the list contents. */ + private static List extractListCell(Map response, String columnLabel) { + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertNotNull("Response missing 'rows'", rows); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + assertNotNull(columnLabel + " must not be null — dataset has non-null values", cell); + assertTrue(columnLabel + " must be a List, got " + cell.getClass(), cell instanceof List); + @SuppressWarnings("unchecked") + List list = (List) cell; + return list; + } + + private Map executePpl(String ppl) throws IOException { + ensureDataProvisioned(); + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLTakeCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLTakeCommandIT.java new file mode 100644 index 0000000000000..d18d71954fb5e --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLTakeCommandIT.java @@ -0,0 +1,170 @@ +/* + * 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.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * End-to-end IT for the PPL {@code take} aggregate routed through the analytics-engine + * REST path. {@code take(field [, n])} collects up to {@code n} values (default 10) + * from the input into a list; PPL makes no ordering guarantee. The DataFusion backend + * implements this via a custom Rust UDAF (see {@code rust/src/udaf/take.rs}) declared + * through {@code opensearch_aggregate.yaml} — the Java-side wiring already shipped + * with PR 21424 and Group E's register_all fix ({@code f985813076a}). + * + *

This IT exercises the REST path rather than the internalClusterTest path covered + * by {@code AggregationUDFIT.testTake} — different layer, different failure signal. + * Assertions use the subset-containment + null-tolerant pattern established by + * {@link CalcitePPLListCommandIT}: every returned element must come from the input + * set (or be null), the array length must respect the documented bound {@code n}, + * but no specific element positioning is required because PPL take() makes no + * ordering promise. + * + *

Unlike {@code list()}, {@code take()} does NOT filter nulls in PPL semantics + * (see {@code TakeAggFunction.add()} in sql/core) — null input values are kept in + * the output array. The assertions here tolerate (but don't require) null elements. + */ +public class CalcitePPLTakeCommandIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("calcs", "calcs"); + + private static boolean dataProvisioned = false; + + private void ensureDataProvisioned() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + /** All non-null {@code str2} values in the {@code calcs} dataset. */ + private static final Set STR2_NON_NULL_VALUES = Set.of( + "one", "two", "three", "five", "six", "eight", "nine", "ten", + "eleven", "twelve", "fourteen", "fifteen", "sixteen" + ); + + /** 17 docs in the calcs dataset → take(field) without n defaults to 10, so at most 10 items. */ + private static final int DEFAULT_TAKE_SIZE = 10; + + // ── take(field) — default size 10 ──────────────────────────────────────────── + + public void testTakeDefaultSize() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats take(str2)"); + List taken = extractListCell(response, "take(str2)"); + assertTrue( + "take(str2) without n must respect the default 10-element cap; got " + taken.size() + " items", + taken.size() <= DEFAULT_TAKE_SIZE + ); + assertContainsOnlyExpectedElements(taken, STR2_NON_NULL_VALUES, "take(str2)"); + } + + // ── take(field, n) — explicit bound ───────────────────────────────────────── + + public void testTakeExplicitSmallSize() throws IOException { + // Ask for 3 values — the result must have at most 3 elements. + Map response = executePpl("source=" + DATASET.indexName + " | stats take(str2, 3)"); + List taken = extractListCell(response, "take(str2, 3)"); + assertTrue("take(str2, 3) must have ≤ 3 elements; got " + taken.size(), taken.size() <= 3); + assertContainsOnlyExpectedElements(taken, STR2_NON_NULL_VALUES, "take(str2, 3)"); + } + + public void testTakeExplicitLargeSize() throws IOException { + // Ask for 100 values — the dataset only has 17 docs; must get ≤ 17 elements. + Map response = executePpl("source=" + DATASET.indexName + " | stats take(str2, 100)"); + List taken = extractListCell(response, "take(str2, 100)"); + assertTrue( + "take(str2, 100) cannot exceed the dataset size of 17; got " + taken.size(), + taken.size() <= 17 + ); + assertContainsOnlyExpectedElements(taken, STR2_NON_NULL_VALUES, "take(str2, 100)"); + } + + public void testTakeOnIntegerField() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats take(int0, 5)"); + List taken = extractListCell(response, "take(int0, 5)"); + assertTrue("take(int0, 5) must have ≤ 5 elements; got " + taken.size(), taken.size() <= 5); + Set int0NonNull = Set.of(1L, 7L, 3L, 8L, 4L, 10L, 11L); + for (Object v : taken) { + if (v != null) { + long actual = ((Number) v).longValue(); + assertTrue( + "take(int0, 5) element " + actual + " must come from the non-null int0 input set " + int0NonNull, + int0NonNull.stream().anyMatch(n -> n.longValue() == actual) + ); + } + } + } + + // ── stats ... by (GROUP BY) ────────────────────────────────────────────────── + + public void testTakeGroupedByBool0() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats take(str2, 2) by bool0"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertTrue("at least one group row expected", rows.size() >= 1); + for (List row : rows) { + assertEquals("two columns per row: take(str2, 2), bool0", 2, row.size()); + Object listCell = row.get(0); + assertTrue( + "take(str2, 2) cell must be an array, got " + (listCell == null ? "null" : listCell.getClass()), + listCell instanceof List + ); + @SuppressWarnings("unchecked") + List taken = (List) listCell; + assertTrue("each group's take(str2, 2) must have ≤ 2 elements; got " + taken.size(), taken.size() <= 2); + assertContainsOnlyExpectedElements(taken, STR2_NON_NULL_VALUES, "take(str2, 2) per-group"); + } + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + /** Every non-null element in {@code taken} must be in {@code allowedNonNull}. + * Null elements are tolerated (PPL take() does not filter nulls). */ + private static void assertContainsOnlyExpectedElements(List taken, Set allowedNonNull, String label) { + Set unexpected = new HashSet<>(); + for (Object v : taken) { + if (v != null && allowedNonNull.contains(v.toString()) == false) { + unexpected.add(v.toString()); + } + } + assertTrue( + label + " contains unexpected values " + unexpected + " not in input set " + allowedNonNull, + unexpected.isEmpty() + ); + } + + private static List extractListCell(Map response, String columnLabel) { + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertNotNull("Response missing 'rows'", rows); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + assertNotNull(columnLabel + " must not be null — dataset has non-null values", cell); + assertTrue(columnLabel + " must be a List, got " + cell.getClass(), cell instanceof List); + @SuppressWarnings("unchecked") + List list = (List) cell; + return list; + } + + private Map executePpl(String ppl) throws IOException { + ensureDataProvisioned(); + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLValuesCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLValuesCommandIT.java new file mode 100644 index 0000000000000..832982978f541 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CalcitePPLValuesCommandIT.java @@ -0,0 +1,219 @@ +/* + * 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.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * End-to-end IT for the PPL {@code values} aggregate routed through the analytics-engine + * REST path. PPL's {@code stats values(field)} collects the distinct values of + * {@code field} in ascending order. The DataFusion backend implements this via an + * AliasConfig entry that routes {@code values → array_agg} with forced DISTINCT and + * a forced ORDER BY on the operand itself. + * + *

Assertions (mirroring the pattern established by list / take): + *

    + *
  • Subset-containment: every non-null input value must appear in the + * collected array exactly once (distinct). Extra null entries are tolerated + * as a documented divergence (PPL {@code values()} filters nulls, DataFusion + * {@code array_agg DISTINCT} keeps at most one null).
    • + *
  • Distinct: no duplicate non-null values in the collected array.
    • + *
  • Ascending sort: the non-null portion of the collected array is in + * ascending order — the primary observable behavior VALUES promises that + * LIST doesn't.
    • + *
+ * + *

Semantic divergences deferred (same as LIST): + *

    + *
  • String-cast: PPL stringifies every element; DataFusion preserves input type.
    • + *
  • Plugin limit: PPL honors {@code plugins.ppl.values.max.limit}; DataFusion + * is unbounded. The calcs dataset has 17 docs so the cap is never hit anyway.
    • + *
+ */ +public class CalcitePPLValuesCommandIT extends AnalyticsRestTestCase { + + private static final Dataset DATASET = new Dataset("calcs", "calcs"); + + private static boolean dataProvisioned = false; + + private void ensureDataProvisioned() throws IOException { + if (dataProvisioned == false) { + DatasetProvisioner.provision(client(), DATASET); + dataProvisioned = true; + } + } + + /** All non-null {@code str2} values in the calcs dataset. */ + private static final Set STR2_NON_NULL_VALUES = Set.of( + "one", "two", "three", "five", "six", "eight", "nine", "ten", + "eleven", "twelve", "fourteen", "fifteen", "sixteen" + ); + + /** Distinct non-null int0 values. int0 has duplicates (4 appears 3x, 8 appears 3x). */ + private static final Set INT0_DISTINCT_NON_NULL = Set.of(1L, 3L, 4L, 7L, 8L, 10L, 11L); + + // ── scalar stats ──────────────────────────────────────────────────────────── + + public void testValuesOnStringField() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats values(str2)"); + List collected = extractListCell(response, "values(str2)"); + assertAllExpectedPresentExactlyOnce(collected, STR2_NON_NULL_VALUES, "values(str2)"); + assertNonNullPortionAscending(collected, "values(str2)"); + } + + public void testValuesOnIntegerField() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats values(int0)"); + List collected = extractListCell(response, "values(int0)"); + // Collect non-null longs from the result. + List collectedNonNull = new ArrayList<>(); + for (Object v : collected) { + if (v != null) { + collectedNonNull.add(((Number) v).longValue()); + } + } + // Every distinct non-null input must appear exactly once in the output. + Set collectedSet = new HashSet<>(collectedNonNull); + assertEquals( + "values(int0) must emit each non-null input exactly once; got " + collectedNonNull, + INT0_DISTINCT_NON_NULL.size(), + collectedSet.size() + ); + for (Long expected : INT0_DISTINCT_NON_NULL) { + assertTrue( + "values(int0) must contain " + expected + "; got " + collectedNonNull, + collectedSet.contains(expected) + ); + } + // Ascending order on the non-null portion. + for (int i = 1; i < collectedNonNull.size(); i++) { + assertTrue( + "values(int0) must be ascending; got " + collectedNonNull + " (violation at index " + i + ")", + collectedNonNull.get(i - 1) <= collectedNonNull.get(i) + ); + } + } + + // ── stats ... by (GROUP BY) ────────────────────────────────────────────────── + + public void testValuesGroupedByBool0() throws IOException { + Map response = executePpl("source=" + DATASET.indexName + " | stats values(str2) by bool0"); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertTrue("at least one group row expected", rows.size() >= 1); + for (List row : rows) { + assertEquals("two columns per row: values(str2), bool0", 2, row.size()); + Object listCell = row.get(0); + assertTrue( + "values(str2) cell must be an array, got " + (listCell == null ? "null" : listCell.getClass()), + listCell instanceof List + ); + @SuppressWarnings("unchecked") + List groupValues = (List) listCell; + // Per-group: no duplicates on non-null elements. + Set seen = new HashSet<>(); + for (Object v : groupValues) { + if (v != null) { + String s = v.toString(); + assertTrue( + "values(str2) per-group must be distinct on non-null; saw " + s + " twice in " + groupValues, + seen.add(s) + ); + assertTrue( + "values(str2) per-group element " + s + " must be in the known str2 set", + STR2_NON_NULL_VALUES.contains(s) + ); + } + } + // Per-group: ascending. + assertNonNullPortionAscending(groupValues, "values(str2) per-group"); + } + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + /** Every expected value must appear in the collected array, no duplicates among non-null + * elements. Null elements are tolerated (DataFusion array_agg DISTINCT keeps one null). + */ + private static void assertAllExpectedPresentExactlyOnce(List collected, Set expected, String label) { + List nonNullStrings = new ArrayList<>(); + for (Object v : collected) { + if (v != null) { + nonNullStrings.add(v.toString()); + } + } + Set asSet = new HashSet<>(nonNullStrings); + assertEquals( + label + " must emit each non-null input exactly once; got " + nonNullStrings, + nonNullStrings.size(), + asSet.size() + ); + for (String want : expected) { + assertTrue(label + " must contain " + want + "; got " + nonNullStrings, asSet.contains(want)); + } + } + + /** Assert the non-null portion of the collected list is in ascending order per + * String.compareTo / Number.doubleValue comparison. */ + private static void assertNonNullPortionAscending(List collected, String label) { + Object previous = null; + int index = 0; + for (Object v : collected) { + if (v != null) { + if (previous != null) { + int cmp = compare(previous, v); + assertTrue( + label + " non-null portion must be ascending; saw " + previous + " before " + v + + " at index " + index + " in " + collected, + cmp <= 0 + ); + } + previous = v; + } + index++; + } + } + + @SuppressWarnings({ "unchecked", "rawtypes" }) + private static int compare(Object a, Object b) { + if (a instanceof Number && b instanceof Number) { + return Double.compare(((Number) a).doubleValue(), ((Number) b).doubleValue()); + } + return ((Comparable) a).compareTo(b); + } + + private static List extractListCell(Map response, String columnLabel) { + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertNotNull("Response missing 'rows'", rows); + assertEquals("scalar agg → exactly 1 result row", 1, rows.size()); + assertEquals("scalar agg → exactly 1 column", 1, rows.get(0).size()); + Object cell = rows.get(0).get(0); + assertNotNull(columnLabel + " must not be null — dataset has non-null values", cell); + assertTrue(columnLabel + " must be a List, got " + cell.getClass(), cell instanceof List); + @SuppressWarnings("unchecked") + List list = (List) cell; + return list; + } + + private Map executePpl(String ppl) throws IOException { + ensureDataProvisioned(); + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +}