diff --git a/plugins/arrow-flight-rpc/build.gradle b/plugins/arrow-flight-rpc/build.gradle index a94c9301a4041..9e3a0b5dc3f98 100644 --- a/plugins/arrow-flight-rpc/build.gradle +++ b/plugins/arrow-flight-rpc/build.gradle @@ -36,6 +36,11 @@ dependencies { api "com.fasterxml.jackson.core:jackson-core:${versions.jackson}" api "com.fasterxml.jackson.core:jackson-databind:${versions.jackson}" api "com.fasterxml.jackson.core:jackson-annotations:${versions.jackson_annotations}" + // arrow-vector's JsonStringArrayList static-initializes a Jackson ObjectMapper that registers + // JavaTimeModule. Without jsr310 on arrow-flight-rpc's classpath, any reader of an Arrow + // ListVector (e.g. DataFusion's array-returning UDFs flowing through analytics-engine) hits + // a fatal NoClassDefFoundError that exits the JVM. + api "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${versions.jackson}" api "commons-codec:commons-codec:${versions.commonscodec}" // arrow flight dependencies. diff --git a/plugins/arrow-flight-rpc/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 b/plugins/arrow-flight-rpc/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 new file mode 100644 index 0000000000000..5bf925c777b5f --- /dev/null +++ b/plugins/arrow-flight-rpc/licenses/jackson-datatype-jsr310-2.21.3.jar.sha1 @@ -0,0 +1 @@ +a0958ebdaba836d31e5462ebc37b6349a0725ff9 diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java index fa013789d5a79..2a6a68a076d09 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldType.java @@ -55,7 +55,16 @@ public enum FieldType { NESTED("nested"), OBJECT("object"), FLAT_OBJECT("flat_object"), - COMPLETION("completion"); + COMPLETION("completion"), + /** + * Array-typed expression result. Used for the return-type slot of array-producing scalar + * functions (PPL {@code array(…)}, {@code array_slice}, {@code array_distinct}). Has no + * OpenSearch mapping equivalent — arrays in OpenSearch are multi-value fields with the + * underlying element type, not a separate type. The mapping string is {@code "array"} as a + * placeholder; {@link #fromMappingType} keeps working unchanged because no source + * advertises that mapping string. + */ + ARRAY("array"); private final String mappingType; @@ -117,6 +126,7 @@ public static FieldType fromSqlTypeName(SqlTypeName sqlTypeName) { case TIME, TIMESTAMP, TIMESTAMP_WITH_LOCAL_TIME_ZONE -> FieldType.DATE; case BOOLEAN -> FieldType.BOOLEAN; case BINARY, VARBINARY -> FieldType.BINARY; + case ARRAY -> FieldType.ARRAY; default -> null; }; } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java index af6f15e84935b..60dd70f629018 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ScalarFunction.java @@ -177,7 +177,60 @@ public enum ScalarFunction { JSON_EXTEND(Category.SCALAR, SqlKind.OTHER_FUNCTION), JSON_EXTRACT(Category.SCALAR, SqlKind.OTHER_FUNCTION), JSON_KEYS(Category.SCALAR, SqlKind.OTHER_FUNCTION), - JSON_SET(Category.SCALAR, SqlKind.OTHER_FUNCTION); + JSON_SET(Category.SCALAR, SqlKind.OTHER_FUNCTION), + + // ── Array ──────────────────────────────────────────────────────── + /** + * PPL {@code array(a, b, …)} constructor — resolves through the SQL plugin's + * {@code ArrayFunctionImpl} UDF named {@code "array"}. DataFusion's native + * equivalent is {@code make_array}, so a backend that supports this needs a + * name-mapping adapter (see {@code MakeArrayAdapter} in the DataFusion backend). + */ + ARRAY(Category.SCALAR, SqlKind.OTHER_FUNCTION), + ARRAY_LENGTH(Category.SCALAR, SqlKind.OTHER_FUNCTION), + ARRAY_SLICE(Category.SCALAR, SqlKind.OTHER_FUNCTION), + ARRAY_DISTINCT(Category.SCALAR, SqlKind.OTHER_FUNCTION), + /** + * Calcite's {@code ARRAY_JOIN} — joins array elements with a separator. PPL + * {@code mvjoin} is registered to this operator. DataFusion's native equivalent + * is named {@code array_to_string}, so the DataFusion backend rewrites to that + * via a name-mapping adapter. + */ + ARRAY_JOIN(Category.SCALAR, SqlKind.OTHER_FUNCTION), + /** + * Calcite's {@code SqlStdOperatorTable.ITEM} — element access ({@code arr[N]}). + * PPL's {@code mvindex(arr, N)} single-element form lowers through + * {@code MVIndexFunctionImp.resolveSingleElement} to ITEM with a 1-based index + * (already converted from PPL's 0-based input). DataFusion's native equivalent + * is {@code array_element}, also 1-based; the DataFusion backend renames via a + * name-mapping adapter. + */ + ITEM(Category.SCALAR, SqlKind.ITEM), + /** + * PPL {@code mvzip(left, right [, sep])} — element-wise zip of two arrays into an + * array of strings, joined per pair by a separator (default {@code ","}). Resolves + * through the SQL plugin's {@code MVZipFunctionImpl} UDF named {@code "mvzip"}. + * No DataFusion stdlib equivalent — the analytics-backend-datafusion plugin ships + * a custom Rust UDF (`udf::mvzip`) registered on its session context. + */ + MVZIP(Category.SCALAR, SqlKind.OTHER_FUNCTION), + /** + * PPL {@code mvfind(arr, regex)} — find the 0-based index of the first array + * element matching a regex, or NULL if no match. Resolves through the SQL + * plugin's {@code MVFindFunctionImpl} UDF named {@code "mvfind"}. No + * DataFusion stdlib equivalent — the analytics-backend-datafusion plugin + * ships a custom Rust UDF (`udf::mvfind`) registered on its session context. + */ + MVFIND(Category.SCALAR, SqlKind.OTHER_FUNCTION), + /** + * PPL {@code mvappend(arg1, arg2, …)} — flatten a mixed list of array and + * scalar arguments into one array, dropping null args and null elements. + * Resolves through the SQL plugin's {@code MVAppendFunctionImpl} UDF named + * {@code "mvappend"}. DataFusion's {@code array_concat} only accepts arrays + * and preserves nulls, so the analytics-backend-datafusion plugin ships a + * custom Rust UDF ({@code udf::mvappend}) registered on its session context. + */ + MVAPPEND(Category.SCALAR, SqlKind.OTHER_FUNCTION); /** * Category of scalar function. diff --git a/sandbox/plugins/analytics-backend-datafusion/build.gradle b/sandbox/plugins/analytics-backend-datafusion/build.gradle index 1d71f16082900..5e5175ac2a8f3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/build.gradle +++ b/sandbox/plugins/analytics-backend-datafusion/build.gradle @@ -73,6 +73,11 @@ dependencies { implementation "io.substrait:isthmus:0.89.1" implementation "io.substrait:core:0.89.1" implementation "com.fasterxml.jackson.datatype:jackson-datatype-jdk8:${versions.jackson}" + // jackson-datatype-jsr310 — added to arrow-flight-rpc (the parent plugin that bundles + // arrow-vector). arrow-vector's JsonStringArrayList eagerly registers JavaTimeModule on + // its ObjectMapper, so jsr310 must be visible to arrow-vector's defining classloader, + // not this plugin's. compileOnly here would also work; runtime is provided by parent. + compileOnly "com.fasterxml.jackson.datatype:jackson-datatype-jsr310:${versions.jackson}" calciteCompile "com.google.guava:guava:${versions.guava}" calciteTestCompile "com.google.guava:guava:${versions.guava}" diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index ba0f7a0c0226a..93c7f0dea7265 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml @@ -68,6 +68,8 @@ serde_json = { workspace = true, features = ["preserve_order"] } # multi-path JSON-array output. Moving to 1.x is a follow-up once we can # reproduce that distinction against the new API surface. jsonpath-rust = "0.7" +# mvfind UDF — regex matching against stringified array elements +regex = "1.10" [dev-dependencies] criterion = { workspace = true } 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 c226deabb2fbe..cb5eea3112b07 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -111,6 +111,10 @@ pub async unsafe fn create_session_context( .build(); let ctx = SessionContext::new_with_state(state); + // Register OpenSearch UDFs (mvappend, mvfind, mvzip, convert_tz, …) on this session + // so the substrait converter at execute_with_context can resolve their function names. + // Without this, fragment execution fails with "Unsupported function name" because + // df_execute_with_context reuses this handle's ctx instead of building a fresh one. crate::udf::register_all(&ctx); // Register default ListingTable for parquet scans diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs index cdbcb9db9bf83..66b4e20abc295 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs @@ -122,6 +122,9 @@ pub mod json_extend; pub mod json_extract; pub mod json_keys; pub mod json_set; +pub mod mvappend; +pub mod mvfind; +pub mod mvzip; pub mod tonumber; pub mod tostring; @@ -141,10 +144,13 @@ pub fn register_all(ctx: &SessionContext) { json_extract::register_all(ctx); json_keys::register_all(ctx); json_set::register_all(ctx); + mvzip::register_all(ctx); + mvfind::register_all(ctx); + mvappend::register_all(ctx); tonumber::register_all(ctx); tostring::register_all(ctx); log::info!( - "OpenSearch UDF register_all: convert_tz, json_append, json_array_length, json_delete, json_extend, json_extract, json_keys, json_set, tonumber, tostring registered" + "OpenSearch UDF register_all: convert_tz, json_append, json_array_length, json_delete, json_extend, json_extract, json_keys, json_set, mvzip, mvfind, mvappend, tonumber, tostring registered" ); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvappend.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvappend.rs new file mode 100644 index 0000000000000..a8e636de7bda3 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvappend.rs @@ -0,0 +1,531 @@ +/* + * 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. + */ + +//! `mvappend(arg1, arg2, …)` — flatten a mixed list of array and scalar args +//! into a single array, dropping null args AND null elements inside array args. +//! +//! Mirrors PPL's [`MVAppendFunctionImpl`] / [`MVAppendCore`] semantics: +//! +//! * For each argument, in order: +//! * NULL argument → skipped entirely. +//! * Array argument → each non-null element is appended to the output. +//! * Scalar argument → appended as a single element. +//! * Returns NULL if no non-null elements were collected (PPL convention — +//! distinguishes `mvappend(null)` from `mvappend()` from `mvappend([])`). +//! +//! ## Type homogeneity +//! +//! The Java adapter (`MvappendAdapter`) casts every scalar argument to the +//! call's array component type and every array argument to +//! `ARRAY` before this UDF runs, so by the time we see operands +//! they share a single element type. The element-conversion macro below +//! handles each supported scalar Arrow type explicitly; a list whose data +//! vector type isn't covered surfaces as a planning error rather than a +//! silent coercion. +//! +//! Mixed-type calls (`mvappend(1, 'text', 2.5)`) end up with Calcite type +//! `ARRAY` which substrait doesn't have an encoding for — those fail at +//! substrait conversion, before reaching this UDF, and aren't addressed here. + +use std::any::Any; +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, BooleanBuilder, Decimal128Array, Decimal128Builder, + Float32Array, Float32Builder, Float64Array, Float64Builder, GenericListArray, Int16Array, + Int16Builder, Int32Array, Int32Builder, Int64Array, Int64Builder, Int8Array, Int8Builder, + ListArray, ListBuilder, StringArray, StringBuilder, StringViewArray, StringViewBuilder, + UInt16Array, UInt16Builder, UInt32Array, UInt32Builder, UInt64Array, UInt64Builder, + UInt8Array, UInt8Builder, +}; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::common::plan_err; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(MvappendUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct MvappendUdf { + signature: Signature, +} + +impl MvappendUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for MvappendUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for MvappendUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "mvappend" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.is_empty() { + return plan_err!("mvappend expects at least 1 argument, got 0"); + } + // Java adapter pre-coerces every operand to either ARRAY or E for a + // single E. Use whichever element type we see first. + let element_type = element_type(arg_types) + .ok_or_else(|| DataFusionError::Plan("mvappend: unable to determine element type from operand types".to_string()))?; + Ok(DataType::List(Arc::new(Field::new( + "item", + element_type, + true, + )))) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.is_empty() { + return plan_err!("mvappend expects at least 1 argument, got 0"); + } + // Trust the Java adapter to have already coerced operands. coerce_types here + // exists only because Signature::user_defined demands an implementation; we + // pass each type through unchanged. + Ok(arg_types.to_vec()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let n = args.number_rows; + let element_type = element_type(&args.arg_fields.iter().map(|f| f.data_type().clone()).collect::>()) + .ok_or_else(|| DataFusionError::Internal("mvappend: lost element type at invoke".to_string()))?; + + // Materialize each operand as an ArrayRef whose Arrow type is either + // {element_type} or List. Scalar operands replicate to n rows. + let operand_arrays: Vec = args + .args + .iter() + .map(|a| a.clone().into_array(n)) + .collect::>>()?; + + macro_rules! build { + ($Builder:ty, $Scalar:ty, $List:ty) => {{ + let inner = <$Builder>::new(); + let mut builder = ListBuilder::new(inner); + for row in 0..n { + let mut any_value = false; + for arr in &operand_arrays { + if arr.is_null(row) { + continue; + } + if let Some(list_arr) = arr.as_any().downcast_ref::>() { + // Iterate elements of the list at this row. + let row_list = list_arr.value(row); + let typed = row_list + .as_any() + .downcast_ref::<$Scalar>() + .ok_or_else(|| DataFusionError::Internal(format!( + "mvappend: list element vector type mismatch ({:?})", + row_list.data_type() + )))?; + for i in 0..typed.len() { + if !typed.is_null(i) { + builder.values().append_value(typed.value(i)); + any_value = true; + } + } + } else if let Some(typed) = arr.as_any().downcast_ref::<$Scalar>() { + builder.values().append_value(typed.value(row)); + any_value = true; + } else { + return plan_err!( + "mvappend: unexpected operand vector type {:?}", + arr.data_type() + ); + } + } + if any_value { + builder.append(true); + } else { + builder.append_null(); + } + } + Arc::new(builder.finish()) as ArrayRef + }}; + } + + let result: ArrayRef = match &element_type { + DataType::Int8 => build!(Int8Builder, Int8Array, ListArray), + DataType::Int16 => build!(Int16Builder, Int16Array, ListArray), + DataType::Int32 => build!(Int32Builder, Int32Array, ListArray), + DataType::Int64 => build!(Int64Builder, Int64Array, ListArray), + DataType::UInt8 => build!(UInt8Builder, UInt8Array, ListArray), + DataType::UInt16 => build!(UInt16Builder, UInt16Array, ListArray), + DataType::UInt32 => build!(UInt32Builder, UInt32Array, ListArray), + DataType::UInt64 => build!(UInt64Builder, UInt64Array, ListArray), + DataType::Float32 => build!(Float32Builder, Float32Array, ListArray), + DataType::Float64 => build!(Float64Builder, Float64Array, ListArray), + DataType::Boolean => build!(BooleanBuilder, BooleanArray, ListArray), + // Decimal128 element type — needs a builder configured with the same precision + // and scale as the input. Calcite's leastRestrictive widening for INT + DECIMAL + // produces DECIMAL(p, s) which substrait converts to Decimal128(p, s); the Java + // adapter's CAST aligns every operand's element type to that. + DataType::Decimal128(precision, scale) => { + let inner = Decimal128Builder::new() + .with_precision_and_scale(*precision, *scale) + .map_err(|e| DataFusionError::Plan(format!("mvappend: decimal builder: {e}")))?; + let mut builder = ListBuilder::new(inner); + for row in 0..n { + let mut any_value = false; + for arr in &operand_arrays { + if arr.is_null(row) { + continue; + } + if let Some(list_arr) = arr.as_any().downcast_ref::>() { + let row_list = list_arr.value(row); + let typed = row_list + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal(format!( + "mvappend: list element vector type mismatch ({:?})", + row_list.data_type() + )))?; + for i in 0..typed.len() { + if !typed.is_null(i) { + builder.values().append_value(typed.value(i)); + any_value = true; + } + } + } else if let Some(typed) = arr.as_any().downcast_ref::() { + builder.values().append_value(typed.value(row)); + any_value = true; + } else { + return plan_err!( + "mvappend: unexpected operand vector type {:?}", + arr.data_type() + ); + } + } + if any_value { + builder.append(true); + } else { + builder.append_null(); + } + } + Arc::new(builder.finish()) as ArrayRef + } + // String element types — handled specially because list children may be any of + // {Utf8, LargeUtf8, Utf8View} depending on whether the operand is a string literal, + // a field read (DataFusion's substrait consumer uses Utf8View for column reads in + // 52+), or a computed expression. The output type must match what {@code return_type} + // declared (which is whatever {@code element_type()} returned, driven by the first + // operand) — DataFusion validates the actual output schema against the declared + // schema and rejects mismatches with "column types must match schema types". + DataType::Utf8 | DataType::LargeUtf8 => { + let mut builder = ListBuilder::new(StringBuilder::new()); + build_string_rows::(&operand_arrays, n, &mut builder)?; + Arc::new(builder.finish()) as ArrayRef + } + DataType::Utf8View => { + let mut builder = ListBuilder::new(StringViewBuilder::new()); + build_string_rows::(&operand_arrays, n, &mut builder)?; + Arc::new(builder.finish()) as ArrayRef + } + other => { + return plan_err!("mvappend: unsupported element type {other:?}"); + } + }; + + Ok(ColumnarValue::Array(result)) + } +} + +/// First operand type drives the element type. The Java adapter has already +/// normalized everything to that single element type — either bare scalar or +/// `List`. String element types are normalized to {@code Utf8} so the +/// match arm in {@link MvappendUdf::invoke_with_args} catches every flavor — +/// scalar literals come through as {@code Utf8}, field reads as {@code Utf8View}, +/// computed expressions as {@code Utf8} or {@code LargeUtf8}. +fn element_type(arg_types: &[DataType]) -> Option { + arg_types.iter().find_map(|t| match t { + DataType::List(field) | DataType::LargeList(field) | DataType::FixedSizeList(field, _) => { + Some(field.data_type().clone()) + } + DataType::Null => None, + other => Some(other.clone()), + }) +} + +/// Trait abstracting the difference between {@link StringBuilder} and +/// {@link StringViewBuilder} when appending {@code &str} values. Both expose +/// `append_value(&str)`, but they're concrete types with no shared trait, so +/// this glue lets {@link build_string_rows} drive either via the same code. +trait StrAppend { + fn append_str(&mut self, s: &str); +} +impl StrAppend for StringBuilder { + fn append_str(&mut self, s: &str) { + self.append_value(s); + } +} +impl StrAppend for StringViewBuilder { + fn append_str(&mut self, s: &str) { + self.append_value(s); + } +} + +/// Generic per-row writer for string-typed mvappend output. Dispatches list-child +/// downcasts across all three Utf8 flavors so the input doesn't need to match the +/// output type. +fn build_string_rows( + operand_arrays: &[ArrayRef], + n: usize, + builder: &mut ListBuilder, +) -> Result<()> { + for row in 0..n { + let mut any_value = false; + for arr in operand_arrays { + if arr.is_null(row) { + continue; + } + if let Some(list_arr) = arr.as_any().downcast_ref::>() { + let row_list = list_arr.value(row); + append_string_elements(row_list.as_ref(), builder.values(), &mut any_value)?; + } else { + append_string_scalar(arr.as_ref(), row, builder.values(), &mut any_value)?; + } + } + if any_value { + builder.append(true); + } else { + builder.append_null(); + } + } + Ok(()) +} + +/// Append all non-null string elements from a row's list to the output builder. +/// Handles list children typed as Utf8, LargeUtf8, or Utf8View. +fn append_string_elements( + row_list: &dyn Array, + out: &mut B, + any_value: &mut bool, +) -> Result<()> { + if let Some(typed) = row_list.as_any().downcast_ref::() { + for i in 0..typed.len() { + if !typed.is_null(i) { + out.append_str(typed.value(i)); + *any_value = true; + } + } + return Ok(()); + } + if let Some(typed) = row_list.as_any().downcast_ref::() { + for i in 0..typed.len() { + if !typed.is_null(i) { + out.append_str(typed.value(i)); + *any_value = true; + } + } + return Ok(()); + } + if let Some(large) = row_list.as_string_opt::() { + for i in 0..large.len() { + if !large.is_null(i) { + out.append_str(large.value(i)); + *any_value = true; + } + } + return Ok(()); + } + plan_err!( + "mvappend: list element vector type mismatch — expected string, got {:?}", + row_list.data_type() + ) +} + +/// Append a single scalar string operand at the given row to the output builder. +/// Handles operands typed as Utf8, LargeUtf8, or Utf8View. +fn append_string_scalar( + arr: &dyn Array, + row: usize, + out: &mut B, + any_value: &mut bool, +) -> Result<()> { + if let Some(typed) = arr.as_any().downcast_ref::() { + out.append_str(typed.value(row)); + *any_value = true; + return Ok(()); + } + if let Some(typed) = arr.as_any().downcast_ref::() { + out.append_str(typed.value(row)); + *any_value = true; + return Ok(()); + } + if let Some(large) = arr.as_string_opt::() { + out.append_str(large.value(row)); + *any_value = true; + return Ok(()); + } + plan_err!( + "mvappend: scalar operand vector type mismatch — expected string, got {:?}", + arr.data_type() + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::AsArray; + use datafusion::common::ScalarValue; + + fn run(args: Vec, n: usize) -> ArrayRef { + let arg_fields: Vec> = args + .iter() + .enumerate() + .map(|(i, cv)| { + let dt = match cv { + ColumnarValue::Array(a) => a.data_type().clone(), + ColumnarValue::Scalar(sv) => sv.data_type(), + }; + Arc::new(Field::new(format!("a{i}"), dt, true)) + }) + .collect(); + let return_field = Arc::new(Field::new( + "out", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )); + let result = MvappendUdf::new() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: n, + return_field, + config_options: Arc::new(datafusion::config::ConfigOptions::default()), + }) + .unwrap(); + match result { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(_) => panic!("expected array"), + } + } + + fn list_of_ints(rows: &[Option<&[Option]>]) -> ArrayRef { + let mut builder = ListBuilder::new(Int32Builder::new()); + for row in rows { + match row { + None => builder.append_null(), + Some(elems) => { + for e in *elems { + match e { + Some(v) => builder.values().append_value(*v), + None => builder.values().append_null(), + } + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn extract_row_ints(arr: &ArrayRef, row: usize) -> Vec { + let list = arr.as_any().downcast_ref::().unwrap(); + if list.is_null(row) { + return vec![]; + } + let inner = list.value(row); + let typed = inner.as_primitive::(); + (0..typed.len()).filter(|i| !typed.is_null(*i)).map(|i| typed.value(i)).collect() + } + + #[test] + fn three_scalar_ints() { + // mvappend(1, 2, 3) → [1, 2, 3] + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(3))), + ]; + let result = run(args, 1); + assert_eq!(extract_row_ints(&result, 0), vec![1, 2, 3]); + } + + #[test] + fn flattens_array_argument() { + // mvappend([1, 2], 3) → [1, 2, 3] + let arr = list_of_ints(&[Some(&[Some(1), Some(2)])]); + let args = vec![ + ColumnarValue::Array(arr), + ColumnarValue::Scalar(ScalarValue::Int32(Some(3))), + ]; + let result = run(args, 1); + assert_eq!(extract_row_ints(&result, 0), vec![1, 2, 3]); + } + + #[test] + fn drops_null_arg() { + // mvappend(NULL, 1, 2) → [1, 2] + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + let result = run(args, 1); + assert_eq!(extract_row_ints(&result, 0), vec![1, 2]); + } + + #[test] + fn drops_null_elements_inside_array() { + // mvappend([1, NULL, 2], 3) → [1, 2, 3] + let arr = list_of_ints(&[Some(&[Some(1), None, Some(2)])]); + let args = vec![ + ColumnarValue::Array(arr), + ColumnarValue::Scalar(ScalarValue::Int32(Some(3))), + ]; + let result = run(args, 1); + assert_eq!(extract_row_ints(&result, 0), vec![1, 2, 3]); + } + + #[test] + fn all_null_args_yield_null_row() { + // mvappend(NULL, NULL) → NULL + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ColumnarValue::Scalar(ScalarValue::Int32(None)), + ]; + let result = run(args, 1); + let list = result.as_any().downcast_ref::().unwrap(); + assert!(list.is_null(0)); + } + + #[test] + fn empty_array_in_args_contributes_nothing() { + // mvappend([], 1) → [1] + let arr = list_of_ints(&[Some(&[])]); + let args = vec![ + ColumnarValue::Array(arr), + ColumnarValue::Scalar(ScalarValue::Int32(Some(1))), + ]; + let result = run(args, 1); + assert_eq!(extract_row_ints(&result, 0), vec![1]); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs new file mode 100644 index 0000000000000..767f949c8c4f2 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs @@ -0,0 +1,366 @@ +/* + * 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. + */ + +//! `mvfind(arr, regex)` — find the 0-based index of the first array element +//! matching a regex, or NULL if no match. +//! +//! Mirrors PPL's [`MVFindFunctionImpl`] semantics: +//! +//! * Per-element regex match (Java's `Matcher.find` semantics — substring +//! match, not full-string anchored). The Rust `regex` crate's +//! `Regex::is_match` matches the same way (unanchored). +//! * NULL element → skipped (continues to next element). +//! * NULL array or NULL pattern → NULL result. +//! * Empty array → NULL result (no element to match). +//! * Returns Int32 — PPL's surface is `Integer`. +//! * Result type is consistent with the YAML declaration in +//! `opensearch_array_functions.yaml`. +//! +//! # Pattern compilation strategy +//! +//! When the pattern operand is a non-NULL Utf8 scalar literal we compile the +//! regex once up front (mirrors the SQL plugin's `tryCompileLiteralPattern` +//! plan-time optimization). Column-valued patterns are compiled per row; +//! invalid patterns yield NULL for that row (per PPL spec, dynamic-pattern +//! errors are non-fatal — bad rows just produce NULL). + +use std::any::Any; +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, Float32Array, Float64Array, GenericListArray, + Int16Array, Int32Array, Int32Builder, Int64Array, Int8Array, ListArray, StringArray, + UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{plan_err, ScalarValue}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use regex::Regex; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(MvfindUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct MvfindUdf { + signature: Signature, +} + +impl MvfindUdf { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for MvfindUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for MvfindUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "mvfind" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.len() != 2 { + return plan_err!("mvfind expects 2 arguments, got {}", arg_types.len()); + } + Ok(DataType::Int32) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.len() != 2 { + return plan_err!("mvfind expects 2 arguments, got {}", arg_types.len()); + } + if !matches!( + &arg_types[0], + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) { + return plan_err!("mvfind: arg 0 expected list type, got {:?}", arg_types[0]); + } + let pattern_t = match &arg_types[1] { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DataType::Utf8, + other => return plan_err!("mvfind: arg 1 expected string, got {other:?}"), + }; + Ok(vec![arg_types[0].clone(), pattern_t]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 2 { + return plan_err!("mvfind expects 2 arguments, got {}", args.args.len()); + } + let n = args.number_rows; + + // Fast path: pattern is a Utf8 scalar literal — compile once. + let scalar_regex: Option = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(p))) + | ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(p))) + | ColumnarValue::Scalar(ScalarValue::Utf8View(Some(p))) = &args.args[1] + { + // Plan-time invalid pattern → planning error so users see it instantly. + // Mirrors the SQL plugin's IllegalArgumentException for invalid literal regex. + match Regex::new(p) { + Ok(r) => Some(r), + Err(e) => return plan_err!("mvfind: invalid regex pattern '{p}': {e}"), + } + } else { + None + }; + + let arr_arr = args.args[0].clone().into_array(n)?; + let list = arr_arr.as_any().downcast_ref::>().ok_or_else(|| { + DataFusionError::Internal(format!( + "mvfind: expected ListArray, got {:?}", + arr_arr.data_type() + )) + })?; + + // Materialize a column-valued pattern up front; for scalar patterns we keep + // the pre-compiled regex. + let pattern_arr_ref: Option = if scalar_regex.is_none() { + Some(args.args[1].clone().into_array(n)?) + } else { + None + }; + let pattern_arr: Option<&StringArray> = pattern_arr_ref + .as_ref() + .and_then(|a| a.as_any().downcast_ref::()); + + let mut builder = Int32Builder::with_capacity(n); + for i in 0..n { + if list.is_null(i) { + builder.append_null(); + continue; + } + // Per-row regex (compile if column-valued; reuse the scalar compile otherwise). + let regex_for_row: Option = match (&scalar_regex, pattern_arr) { + (Some(r), _) => Some(r.clone()), + (None, Some(arr)) if !arr.is_null(i) => Regex::new(arr.value(i)).ok(), + _ => None, + }; + let regex = match regex_for_row { + Some(r) => r, + None => { + builder.append_null(); + continue; + } + }; + let row = list.value(i); + match find_first_match(row.as_ref(), ®ex) { + Some(idx) => builder.append_value(idx), + None => builder.append_null(), + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) + } +} + +/// Walk an Arrow array of any supported scalar type, return the 0-based index +/// of the first non-null element whose stringified form matches `regex`. +/// Returns None if no element matches. +fn find_first_match(arr: &dyn Array, regex: &Regex) -> Option { + let n = arr.len(); + macro_rules! scan { + ($A:ty, $fmt:expr) => {{ + let typed = arr.as_any().downcast_ref::<$A>()?; + for i in 0..n { + if typed.is_null(i) { + continue; + } + if regex.is_match(&$fmt(typed.value(i))) { + return Some(i as i32); + } + } + None + }}; + } + match arr.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + // String children may arrive as any of the three Utf8 flavors. + if let Some(typed) = arr.as_string_opt::() { + for i in 0..n { + if typed.is_null(i) { + continue; + } + if regex.is_match(typed.value(i)) { + return Some(i as i32); + } + } + None + } else { + let large = arr.as_string_opt::()?; + for i in 0..n { + if large.is_null(i) { + continue; + } + if regex.is_match(large.value(i)) { + return Some(i as i32); + } + } + None + } + } + DataType::Int8 => scan!(Int8Array, |v: i8| v.to_string()), + DataType::Int16 => scan!(Int16Array, |v: i16| v.to_string()), + DataType::Int32 => scan!(Int32Array, |v: i32| v.to_string()), + DataType::Int64 => scan!(Int64Array, |v: i64| v.to_string()), + DataType::UInt8 => scan!(UInt8Array, |v: u8| v.to_string()), + DataType::UInt16 => scan!(UInt16Array, |v: u16| v.to_string()), + DataType::UInt32 => scan!(UInt32Array, |v: u32| v.to_string()), + DataType::UInt64 => scan!(UInt64Array, |v: u64| v.to_string()), + DataType::Float32 => scan!(Float32Array, |v: f32| v.to_string()), + DataType::Float64 => scan!(Float64Array, |v: f64| v.to_string()), + DataType::Boolean => scan!(BooleanArray, |v: bool| v.to_string()), + DataType::Null => None, + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Int32Builder, ListBuilder, StringBuilder}; + use datafusion::arrow::datatypes::Field; + + fn run(arr: ArrayRef, pattern: &str) -> ArrayRef { + let n = arr.len(); + let return_field = Arc::new(Field::new("out", DataType::Int32, true)); + let arg_fields: Vec> = vec![ + Arc::new(Field::new("a", arr.data_type().clone(), true)), + Arc::new(Field::new("p", DataType::Utf8, true)), + ]; + let result = MvfindUdf::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(arr), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))), + ], + arg_fields, + number_rows: n, + return_field, + config_options: Arc::new(datafusion::config::ConfigOptions::default()), + }) + .unwrap(); + match result { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(_) => panic!("expected array"), + } + } + + fn list_of_strings(rows: &[Option<&[Option<&str>]>]) -> ArrayRef { + let mut builder = ListBuilder::new(StringBuilder::new()); + for row in rows { + match row { + None => builder.append_null(), + Some(elems) => { + for e in *elems { + match e { + Some(s) => builder.values().append_value(s), + None => builder.values().append_null(), + } + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn list_of_ints(rows: &[Option<&[Option]>]) -> ArrayRef { + let mut builder = ListBuilder::new(Int32Builder::new()); + for row in rows { + match row { + None => builder.append_null(), + Some(elems) => { + for e in *elems { + match e { + Some(v) => builder.values().append_value(*v), + None => builder.values().append_null(), + } + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn int_value(arr: &ArrayRef, row: usize) -> Option { + let typed = arr.as_any().downcast_ref::().unwrap(); + if typed.is_null(row) { + None + } else { + Some(typed.value(row)) + } + } + + #[test] + fn first_match_returns_zero_based_index() { + let arr = list_of_strings(&[Some(&[Some("apple"), Some("banana"), Some("apricot")])]); + let result = run(arr, "ban.*"); + assert_eq!(int_value(&result, 0), Some(1)); + } + + #[test] + fn no_match_returns_null() { + let arr = list_of_strings(&[Some(&[Some("apple"), Some("banana")])]); + let result = run(arr, "kiwi"); + assert_eq!(int_value(&result, 0), None); + } + + #[test] + fn null_array_returns_null() { + let arr = list_of_strings(&[None]); + let result = run(arr, "any"); + assert_eq!(int_value(&result, 0), None); + } + + #[test] + fn empty_array_returns_null() { + let arr = list_of_strings(&[Some(&[])]); + let result = run(arr, "any"); + assert_eq!(int_value(&result, 0), None); + } + + #[test] + fn null_element_skipped_index_still_zero_based() { + let arr = list_of_strings(&[Some(&[None, Some("banana")])]); + let result = run(arr, "ban.*"); + assert_eq!(int_value(&result, 0), Some(1)); + } + + #[test] + fn integer_array_stringified_for_regex() { + let arr = list_of_ints(&[Some(&[Some(10), Some(20), Some(30)])]); + let result = run(arr, "^2"); + assert_eq!(int_value(&result, 0), Some(1)); + } + + #[test] + fn substring_match_semantics_are_unanchored() { + // Java's Matcher.find: "banana" matches /an/ (unanchored). + let arr = list_of_strings(&[Some(&[Some("apple"), Some("banana")])]); + let result = run(arr, "an"); + assert_eq!(int_value(&result, 0), Some(1)); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvzip.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvzip.rs new file mode 100644 index 0000000000000..a0c8466be835e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvzip.rs @@ -0,0 +1,412 @@ +/* + * 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. + */ + +//! `mvzip(left, right [, separator])` — element-wise zip of two arrays into an +//! array of strings, joined per pair by a separator (default `,`). +//! +//! Mirrors PPL's [`MVZipFunctionImpl`] semantics: +//! +//! * Result length is `min(len(left), len(right))` (Python-`zip` truncation). +//! * Either array NULL → NULL result. +//! * Element NULLs are rendered as empty strings (matches the SQL plugin's +//! `Objects.toString(elem, "")`), so `mvzip([1, NULL], ["a", "b"])` yields +//! `["1,a", ",b"]`. +//! * The separator is a Utf8 scalar. Calling it as a column would require +//! per-row materialization; PPL's surface only exposes a literal so we +//! constrain to scalars here and produce a planning error otherwise. +//! +//! Result type is `List` regardless of the input element types — `mvzip` +//! is fundamentally a string-formatting operation. The Java side relies on the +//! `opensearch_array_functions.yaml` declaration to type the call before +//! substrait emission. + +use std::any::Any; +use std::sync::Arc; + +use datafusion::arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, Float32Array, Float64Array, GenericListArray, + Int16Array, Int32Array, Int64Array, Int8Array, ListArray, ListBuilder, StringArray, + StringBuilder, UInt16Array, UInt32Array, UInt64Array, UInt8Array, +}; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::common::{plan_err, ScalarValue}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +pub fn register_all(ctx: &SessionContext) { + ctx.register_udf(ScalarUDF::from(MvzipUdf::new())); +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct MvzipUdf { + signature: Signature, +} + +impl MvzipUdf { + pub fn new() -> Self { + // user_defined lets us accept ListArray with any element type and a + // string separator without enumerating every concrete combination. + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl Default for MvzipUdf { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for MvzipUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "mvzip" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.len() != 2 && arg_types.len() != 3 { + return plan_err!("mvzip expects 2 or 3 arguments, got {}", arg_types.len()); + } + Ok(DataType::List(Arc::new(Field::new( + "item", + DataType::Utf8, + true, + )))) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.len() != 2 && arg_types.len() != 3 { + return plan_err!("mvzip expects 2 or 3 arguments, got {}", arg_types.len()); + } + for (i, t) in arg_types.iter().take(2).enumerate() { + if !matches!(t, DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _)) { + return plan_err!("mvzip: arg {i} expected list type, got {t:?}"); + } + } + let mut coerced: Vec = arg_types[..2].to_vec(); + if arg_types.len() == 3 { + match &arg_types[2] { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + coerced.push(DataType::Utf8) + } + other => return plan_err!("mvzip: arg 2 (separator) expected string, got {other:?}"), + } + } + Ok(coerced) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + if args.args.len() != 2 && args.args.len() != 3 { + return plan_err!("mvzip expects 2 or 3 arguments, got {}", args.args.len()); + } + let n = args.number_rows; + + // Materialize left/right as ListArrays. Scalar list operands are unusual but + // ColumnarValue::into_array handles them by replicating to {row count}. + let left_arr = args.args[0].clone().into_array(n)?; + let right_arr = args.args[1].clone().into_array(n)?; + + let separator: String = if args.args.len() == 3 { + scalar_string(&args.args[2]) + .ok_or_else(|| { + DataFusionError::Plan( + "mvzip: separator must be a non-NULL string scalar literal".to_string(), + ) + })? + .to_string() + } else { + ",".to_string() + }; + + let left = downcast_list(&left_arr, "left")?; + let right = downcast_list(&right_arr, "right")?; + + // Output: List, one row per input row. + let mut builder = ListBuilder::new(StringBuilder::new()); + for i in 0..n { + if left.is_null(i) || right.is_null(i) { + builder.append_null(); + continue; + } + let left_row = left.value(i); + let right_row = right.value(i); + let take = left_row.len().min(right_row.len()); + let left_strs = elements_as_strings(left_row.as_ref())?; + let right_strs = elements_as_strings(right_row.as_ref())?; + for j in 0..take { + let l = left_strs[j].as_deref().unwrap_or(""); + let r = right_strs[j].as_deref().unwrap_or(""); + let mut joined = + String::with_capacity(l.len() + separator.len() + r.len()); + joined.push_str(l); + joined.push_str(&separator); + joined.push_str(r); + builder.values().append_value(joined); + } + builder.append(true); + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) + } +} + +/// Extract a non-NULL Utf8 scalar literal, or return None for everything else +/// (NULL scalar, column-valued, or non-string types). +fn scalar_string(cv: &ColumnarValue) -> Option<&str> { + if let ColumnarValue::Scalar(sv) = cv { + match sv { + ScalarValue::Utf8(Some(s)) + | ScalarValue::LargeUtf8(Some(s)) + | ScalarValue::Utf8View(Some(s)) => Some(s.as_str()), + _ => None, + } + } else { + None + } +} + +fn downcast_list<'a>(arr: &'a ArrayRef, slot: &str) -> Result<&'a ListArray> { + arr.as_any().downcast_ref::>().ok_or_else(|| { + DataFusionError::Internal(format!( + "mvzip: {slot} expected ListArray, got {:?}", + arr.data_type() + )) + }) +} + +/// Convert each element of an Arrow array of any supported scalar type to its +/// canonical string form (matching `Objects.toString(elem, "")` semantics on +/// the SQL plugin side — NULL elements become None which the caller renders as +/// the empty string). +fn elements_as_strings(arr: &dyn Array) -> Result>> { + let n = arr.len(); + let mut out: Vec> = Vec::with_capacity(n); + macro_rules! collect { + ($A:ty, $fmt:expr) => {{ + let typed = arr.as_any().downcast_ref::<$A>().ok_or_else(|| { + DataFusionError::Internal(format!( + "mvzip: failed to downcast element vector to {}", + stringify!($A) + )) + })?; + for i in 0..n { + if typed.is_null(i) { + out.push(None); + } else { + out.push(Some($fmt(typed.value(i)))); + } + } + }}; + } + match arr.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + // String children may arrive as any of the three Utf8 flavors; AsArray + // handles the dispatch for us. + let s = arr.as_string_opt::(); + if let Some(typed) = s { + for i in 0..n { + if typed.is_null(i) { + out.push(None); + } else { + out.push(Some(typed.value(i).to_string())); + } + } + } else { + let large = arr.as_string_opt::().ok_or_else(|| { + DataFusionError::Internal(format!( + "mvzip: string element downcast failed for {:?}", + arr.data_type() + )) + })?; + for i in 0..n { + if large.is_null(i) { + out.push(None); + } else { + out.push(Some(large.value(i).to_string())); + } + } + } + } + DataType::Int8 => collect!(Int8Array, |v: i8| v.to_string()), + DataType::Int16 => collect!(Int16Array, |v: i16| v.to_string()), + DataType::Int32 => collect!(Int32Array, |v: i32| v.to_string()), + DataType::Int64 => collect!(Int64Array, |v: i64| v.to_string()), + DataType::UInt8 => collect!(UInt8Array, |v: u8| v.to_string()), + DataType::UInt16 => collect!(UInt16Array, |v: u16| v.to_string()), + DataType::UInt32 => collect!(UInt32Array, |v: u32| v.to_string()), + DataType::UInt64 => collect!(UInt64Array, |v: u64| v.to_string()), + DataType::Float32 => collect!(Float32Array, |v: f32| v.to_string()), + DataType::Float64 => collect!(Float64Array, |v: f64| v.to_string()), + DataType::Boolean => collect!(BooleanArray, |v: bool| v.to_string()), + DataType::Null => { + // Element vector with Null type — every cell is NULL by definition, so + // emit None for each. Reachable when the input list is empty and its + // declared element type is Null/UNKNOWN (e.g. PPL `array()` no-arg + // before the SQL-plugin VARCHAR-default kicks in). + for _ in 0..n { + out.push(None); + } + } + other => { + return Err(DataFusionError::NotImplemented(format!( + "mvzip: unsupported list element type {other:?}" + ))); + } + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int32Builder; + use datafusion::arrow::array::StringBuilder as ArrowStringBuilder; + + fn run(left: ArrayRef, right: ArrayRef, sep: Option<&str>) -> ArrayRef { + let mut args = vec![ColumnarValue::Array(left.clone()), ColumnarValue::Array(right.clone())]; + if let Some(s) = sep { + args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(s.to_string())))); + } + let n = left.len(); + let return_field = Arc::new(Field::new( + "out", + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + true, + )); + let arg_fields: Vec> = args + .iter() + .enumerate() + .map(|(i, _)| Arc::new(Field::new(format!("a{i}"), DataType::Utf8, true))) + .collect(); + let result = MvzipUdf::new() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: n, + return_field, + config_options: Arc::new(datafusion::config::ConfigOptions::default()), + }) + .unwrap(); + match result { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(_) => panic!("expected array"), + } + } + + fn list_of_strings(rows: &[Option<&[Option<&str>]>]) -> ArrayRef { + let mut builder = ListBuilder::new(ArrowStringBuilder::new()); + for row in rows { + match row { + None => builder.append_null(), + Some(elems) => { + for e in *elems { + match e { + Some(s) => builder.values().append_value(s), + None => builder.values().append_null(), + } + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn list_of_ints(rows: &[Option<&[Option]>]) -> ArrayRef { + let mut builder = ListBuilder::new(Int32Builder::new()); + for row in rows { + match row { + None => builder.append_null(), + Some(elems) => { + for e in *elems { + match e { + Some(v) => builder.values().append_value(*v), + None => builder.values().append_null(), + } + } + builder.append(true); + } + } + } + Arc::new(builder.finish()) + } + + fn extract_row_strings(arr: &ArrayRef, row: usize) -> Vec { + let list = arr.as_any().downcast_ref::().unwrap(); + let inner = list.value(row); + let strs = inner.as_string::(); + (0..strs.len()).map(|i| strs.value(i).to_string()).collect() + } + + #[test] + fn basic_two_string_arrays_with_default_separator() { + let left = list_of_strings(&[Some(&[Some("a"), Some("b")])]); + let right = list_of_strings(&[Some(&[Some("1"), Some("2")])]); + let result = run(left, right, None); + assert_eq!(extract_row_strings(&result, 0), vec!["a,1", "b,2"]); + } + + #[test] + fn custom_separator() { + let left = list_of_strings(&[Some(&[Some("x"), Some("y")])]); + let right = list_of_strings(&[Some(&[Some("1"), Some("2")])]); + let result = run(left, right, Some("-")); + assert_eq!(extract_row_strings(&result, 0), vec!["x-1", "y-2"]); + } + + #[test] + fn truncate_to_shorter_array() { + let left = list_of_strings(&[Some(&[Some("a"), Some("b"), Some("c")])]); + let right = list_of_strings(&[Some(&[Some("1")])]); + let result = run(left, right, None); + assert_eq!(extract_row_strings(&result, 0), vec!["a,1"]); + } + + #[test] + fn null_element_renders_as_empty_string() { + let left = list_of_strings(&[Some(&[Some("a"), None])]); + let right = list_of_strings(&[Some(&[Some("1"), Some("2")])]); + let result = run(left, right, None); + assert_eq!(extract_row_strings(&result, 0), vec!["a,1", ",2"]); + } + + #[test] + fn null_array_yields_null_row() { + let left = list_of_strings(&[None]); + let right = list_of_strings(&[Some(&[Some("1")])]); + let result = run(left, right, None); + let list = result.as_any().downcast_ref::().unwrap(); + assert!(list.is_null(0)); + } + + #[test] + fn empty_array_yields_empty_result() { + let left = list_of_strings(&[Some(&[])]); + let right = list_of_strings(&[Some(&[Some("1")])]); + let result = run(left, right, None); + assert_eq!(extract_row_strings(&result, 0), Vec::::new()); + } + + #[test] + fn integer_arrays_are_stringified() { + let left = list_of_ints(&[Some(&[Some(10), Some(20)])]); + let right = list_of_ints(&[Some(&[Some(1), Some(2)])]); + let result = run(left, right, None); + assert_eq!(extract_row_strings(&result, 0), vec!["10,1", "20,2"]); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayElementAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayElementAdapter.java new file mode 100644 index 0000000000000..02476ad222a4e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayElementAdapter.java @@ -0,0 +1,85 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +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.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.util.List; + +/** + * Adapter for Calcite's {@link SqlStdOperatorTable#ITEM} operator (element + * access — {@code arr[N]}). PPL's {@code mvindex(arr, N)} single-element form + * lowers through {@code MVIndexFunctionImp.resolveSingleElement} to ITEM with + * a 1-based index (already converted from PPL's 0-based input). + * + *

Two transforms before substrait emission: + * + *

    + *
  1. Rename to {@code array_element}. DataFusion's native single-element + * array accessor is named {@code array_element} (also 1-based), declared + * in {@code opensearch_array_functions.yaml}. Calcite's ITEM operator name + * is {@code "ITEM"} which doesn't resolve to anything in the substrait + * extension catalog. + *
  2. Coerce the index to {@code BIGINT}. PPL's parser types positive + * integer literals as {@code DECIMAL(20,0)}; DataFusion's + * {@code array_element} signature accepts only integer indexes. + *
+ * + * @opensearch.internal + */ +class ArrayElementAdapter implements ScalarFunctionAdapter { + + /** + * Locally-declared target operator. Name matches DataFusion's native + * {@code array_element}. Return-type inference is a placeholder — the + * adapt method explicitly carries the original ITEM call's return type + * (the element type). + */ + static final SqlOperator LOCAL_ARRAY_ELEMENT_OP = new SqlFunction( + "array_element", + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.ANY_ANY, + SqlFunctionCategory.SYSTEM + ); + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + RelDataTypeFactory typeFactory = cluster.getTypeFactory(); + List operands = original.getOperands(); + if (operands.size() != 2) { + return rexBuilder.makeCall(original.getType(), LOCAL_ARRAY_ELEMENT_OP, operands); + } + RexNode array = operands.get(0); + RexNode index = operands.get(1); + if (index.getType().getSqlTypeName() != SqlTypeName.BIGINT) { + RelDataType bigint = typeFactory.createSqlType(SqlTypeName.BIGINT); + RelDataType nullableBigint = typeFactory.createTypeWithNullability(bigint, index.getType().isNullable()); + index = rexBuilder.makeCast(nullableBigint, index, true, false); + } + return rexBuilder.makeCall(original.getType(), LOCAL_ARRAY_ELEMENT_OP, List.of(array, index)); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArraySliceAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArraySliceAdapter.java new file mode 100644 index 0000000000000..15202620d10ee --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArraySliceAdapter.java @@ -0,0 +1,113 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.fun.SqlLibraryOperators; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.List; + +/** + * Adapter for Calcite's {@link SqlLibraryOperators#ARRAY_SLICE}. Two transforms + * are needed before substrait emission: + * + *
    + *
  1. Index coercion to {@code BIGINT}. PPL's parser types positive + * integer literals as {@code DECIMAL(20,0)} (precision wide enough to + * hold any 64-bit unsigned value), but DataFusion's {@code array_slice} + * signature accepts only integer indexes and refuses to coerce decimal + * arguments — failing with {@code "No function matches the given name + * and argument types 'array_slice(List(Int32), Decimal128(20, 0), + * Decimal128(22, 0))'"}. + *
  2. Semantic conversion: 0-based {@code (start, length)} → + * 1-based {@code (start, end)} inclusive. Calcite's + * {@link SqlLibraryOperators#ARRAY_SLICE} (used by PPL's + * {@code MVIndexFunctionImp.resolveRange}) is the Spark / Hive flavor + * with 0-based start and a length-of-elements third arg. DataFusion's + * native {@code array_slice} is 1-based with an inclusive end-index + * third arg. Without this conversion, {@code mvindex(arr=[1..5], 1, 3)} + * would emit {@code ARRAY_SLICE(arr, 1, 3)} → DataFusion returns + * {@code [1, 2, 3]}, but the PPL expectation is {@code [2, 3, 4]} + * (0-based positions 1..3 inclusive). + *

    The conversion is purely arithmetic on the operands: + *

      + *
    • {@code start' = start + 1} + *
    • {@code end' = start + length} (which is {@code start + 1 + + * (length - 1)} = the 1-based inclusive end) + *
    + * Negative indexes have already been normalized to non-negative + * 0-based positions by {@code MVIndexFunctionImp} before this adapter + * runs (it uses {@code arrayLen + idx} for both start and end), so the + * arithmetic above applies uniformly. + *
+ * + *

The 2-arg form {@code ARRAY_SLICE(arr, start)} (single-element extract) + * is not produced by PPL's {@code MVIndexFunctionImp} (single-element access + * lowers through {@code INTERNAL_ITEM} instead), so this adapter handles + * only the 3-arg form. + * + * @opensearch.internal + */ +class ArraySliceAdapter implements ScalarFunctionAdapter { + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + RelDataTypeFactory typeFactory = cluster.getTypeFactory(); + RelDataType bigint = typeFactory.createSqlType(SqlTypeName.BIGINT); + List operands = original.getOperands(); + if (operands.size() != 3) { + // Defensive: unexpected arity. Fall through with BIGINT coercion only — the substrait + // converter will surface a missing-signature error with a clear message. + return rexBuilder.makeCall( + original.getType(), + original.getOperator(), + coerceIndexes(rexBuilder, typeFactory, bigint, operands) + ); + } + List coerced = coerceIndexes(rexBuilder, typeFactory, bigint, operands); + RexNode array = coerced.get(0); + RexNode start = coerced.get(1); + RexNode length = coerced.get(2); + RexNode one = rexBuilder.makeExactLiteral(BigDecimal.ONE, bigint); + RexNode oneBasedStart = rexBuilder.makeCall(SqlStdOperatorTable.PLUS, start, one); + RexNode endInclusive = rexBuilder.makeCall(SqlStdOperatorTable.PLUS, start, length); + return rexBuilder.makeCall(original.getType(), original.getOperator(), List.of(array, oneBasedStart, endInclusive)); + } + + private static List coerceIndexes( + RexBuilder rexBuilder, + RelDataTypeFactory typeFactory, + RelDataType bigint, + List operands + ) { + List coerced = new ArrayList<>(operands.size()); + for (int i = 0; i < operands.size(); i++) { + RexNode operand = operands.get(i); + if (i == 0 || operand.getType().getSqlTypeName() == SqlTypeName.BIGINT) { + coerced.add(operand); + } else { + RelDataType nullableBigint = typeFactory.createTypeWithNullability(bigint, operand.getType().isNullable()); + coerced.add(rexBuilder.makeCast(nullableBigint, operand, true, false)); + } + } + return coerced; + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayToStringAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayToStringAdapter.java new file mode 100644 index 0000000000000..258b47a75440e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrayToStringAdapter.java @@ -0,0 +1,45 @@ +/* + * 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.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.AbstractNameMappingAdapter; + +import java.util.List; + +/** + * Rename adapter for Calcite's {@code ARRAY_JOIN(arr, sep)} — used by PPL's + * {@code mvjoin} via {@code SqlLibraryOperators.ARRAY_JOIN}. DataFusion's native + * equivalent is named {@code array_to_string} (same semantics: join array + * elements with a separator). Rewrites to a locally-declared {@link SqlFunction} + * named {@code array_to_string}; isthmus emits a Substrait scalar function call + * with that name and DataFusion's substrait consumer resolves it natively. + * + * @opensearch.internal + */ +class ArrayToStringAdapter extends AbstractNameMappingAdapter { + + static final SqlOperator LOCAL_ARRAY_TO_STRING_OP = new SqlFunction( + "array_to_string", + SqlKind.OTHER_FUNCTION, + ReturnTypes.VARCHAR_NULLABLE, + null, + OperandTypes.ANY_ANY, + SqlFunctionCategory.SYSTEM + ); + + ArrayToStringAdapter() { + super(LOCAL_ARRAY_TO_STRING_OP, List.of(), List.of()); + } +} 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 4e5b5955cb92c..176de49872dd1 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 @@ -228,7 +228,45 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP ScalarFunction.JSON_EXTEND, ScalarFunction.JSON_EXTRACT, ScalarFunction.JSON_KEYS, - ScalarFunction.JSON_SET + ScalarFunction.JSON_SET, + // Array functions whose RETURN type is element-typed (not ARRAY itself), so the + // capability lookup at OpenSearchProjectRule resolves the call's return type to a + // standard scalar FieldType and matches against SUPPORTED_FIELD_TYPES. + // ARRAY_LENGTH returns BIGINT → FieldType.LONG; ARRAY_JOIN returns VARCHAR → + // FieldType.KEYWORD (renamed to DataFusion `array_to_string` via {@link ArrayToStringAdapter}). + // ITEM returns the array's element type (any of the supported scalar types) — used by + // PPL `mvindex(arr, N)` single-element form. + ScalarFunction.ARRAY_LENGTH, + ScalarFunction.ARRAY_JOIN, + ScalarFunction.ITEM, + // PPL `mvfind` returns INTEGER (the 0-based index of the first match, or NULL); backed + // by a custom Rust UDF on the DataFusion session context (`udf::mvfind`), routed via + // {@link MvfindAdapter}. + ScalarFunction.MVFIND + ); + + /** + * Project-side scalar functions whose return type is {@code ARRAY}. Registered separately + * because the capability lookup keys on the call's return type, and for these the lookup + * resolves to {@link FieldType#ARRAY} — which is intentionally not in + * {@link #SUPPORTED_FIELD_TYPES} (filter and aggregate operators have no meaningful semantics + * over array-typed values, so we don't want them claiming viability there). + * + *

{@code ARRAY} (PPL {@code array(a, b, …)} constructor) renames to DataFusion's + * {@code make_array} via {@link MakeArrayAdapter}. {@code ARRAY_SLICE} and + * {@code ARRAY_DISTINCT} pass through by name (Calcite stdlib operator names match + * DataFusion's native names — isthmus default catalog binds them). + */ + private static final Set ARRAY_RETURNING_PROJECT_OPS = Set.of( + ScalarFunction.ARRAY, + ScalarFunction.ARRAY_SLICE, + ScalarFunction.ARRAY_DISTINCT, + // PPL `mvzip` returns ARRAY; backed by a custom Rust UDF on the DataFusion + // session context (`udf::mvzip`), routed via {@link MvzipAdapter}. + ScalarFunction.MVZIP, + // PPL `mvappend` returns ARRAY; backed by a custom Rust UDF + // (`udf::mvappend`), routed via {@link MvappendAdapter}. + ScalarFunction.MVAPPEND ); private static final Set AGG_FUNCTIONS = Set.of( @@ -289,6 +327,9 @@ public Set projectCapabilities() { for (ScalarFunction op : STANDARD_PROJECT_OPS) { caps.add(new ProjectCapability.Scalar(op, Set.copyOf(SUPPORTED_FIELD_TYPES), formats, true)); } + for (ScalarFunction op : ARRAY_RETURNING_PROJECT_OPS) { + caps.add(new ProjectCapability.Scalar(op, Set.of(FieldType.ARRAY), formats, true)); + } return Set.copyOf(caps); } @@ -319,6 +360,13 @@ public Map scalarFunctionAdapters() { DateTimeAdapters.CurrentDateAdapter currentDate = new DateTimeAdapters.CurrentDateAdapter(); DateTimeAdapters.CurrentTimeAdapter currentTime = new DateTimeAdapters.CurrentTimeAdapter(); return Map.ofEntries( + Map.entry(ScalarFunction.ARRAY, new MakeArrayAdapter()), + Map.entry(ScalarFunction.ARRAY_JOIN, new ArrayToStringAdapter()), + Map.entry(ScalarFunction.ARRAY_SLICE, new ArraySliceAdapter()), + Map.entry(ScalarFunction.ITEM, new ArrayElementAdapter()), + Map.entry(ScalarFunction.MVFIND, new MvfindAdapter()), + Map.entry(ScalarFunction.MVZIP, new MvzipAdapter()), + Map.entry(ScalarFunction.MVAPPEND, new MvappendAdapter()), Map.entry(ScalarFunction.CONCAT, new ConcatFunctionAdapter()), Map.entry(ScalarFunction.CONVERT_TZ, new ConvertTzAdapter()), Map.entry(ScalarFunction.COSH, new HyperbolicOperatorAdapter(SqlLibraryOperators.COSH)), 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 a06d2a4bb20d3..2161abc08b7a2 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 @@ -153,7 +153,23 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { FunctionMappings.s(JsonFunctionAdapters.JsonSetAdapter.LOCAL_JSON_SET_OP, "json_set"), FunctionMappings.s(SqlLibraryOperators.REGEXP_CONTAINS, "regex_match"), FunctionMappings.s(SqlStdOperatorTable.REPLACE, "replace"), - FunctionMappings.s(SqlLibraryOperators.REGEXP_REPLACE_3, "regexp_replace") + FunctionMappings.s(SqlLibraryOperators.REGEXP_REPLACE_3, "regexp_replace"), + // Array S0 ladder — see DataFusionAnalyticsBackendPlugin.STANDARD_PROJECT_OPS / + // ARRAY_RETURNING_PROJECT_OPS for the capability registration. ARRAY_LENGTH / + // ARRAY_SLICE / ARRAY_DISTINCT pass through under their Calcite-stdlib names + // (DataFusion's substrait consumer resolves them natively). MakeArrayAdapter / + // ArrayToStringAdapter / ArrayElementAdapter rewrite PPL `array(...)` / + // `mvjoin(...)` / `mvindex(...)` single-element to locally-declared SqlFunctions + // so isthmus emits Substrait calls with DataFusion's native function names. + FunctionMappings.s(SqlLibraryOperators.ARRAY_LENGTH, "array_length"), + FunctionMappings.s(SqlLibraryOperators.ARRAY_SLICE, "array_slice"), + FunctionMappings.s(SqlLibraryOperators.ARRAY_DISTINCT, "array_distinct"), + FunctionMappings.s(MakeArrayAdapter.LOCAL_MAKE_ARRAY_OP, "make_array"), + FunctionMappings.s(ArrayToStringAdapter.LOCAL_ARRAY_TO_STRING_OP, "array_to_string"), + FunctionMappings.s(ArrayElementAdapter.LOCAL_ARRAY_ELEMENT_OP, "array_element"), + FunctionMappings.s(MvzipAdapter.LOCAL_MVZIP_OP, "mvzip"), + FunctionMappings.s(MvfindAdapter.LOCAL_MVFIND_OP, "mvfind"), + FunctionMappings.s(MvappendAdapter.LOCAL_MVAPPEND_OP, "mvappend") ); private final SimpleExtension.ExtensionCollection extensions; @@ -214,7 +230,19 @@ private byte[] convertToSubstrait(RelNode fragment) { RelNode preprocessed = UntypedNullPreprocessor.rewrite(fragment); RelRoot root = RelRoot.of(preprocessed, SqlKind.SELECT); SubstraitRelVisitor visitor = createVisitor(preprocessed); - Rel substraitRel = visitor.apply(root.rel); + Rel substraitRel; + try { + substraitRel = visitor.apply(root.rel); + } catch (AssertionError e) { + // Substrait validators (e.g. VariadicParameterConsistencyValidator, + // RelOptUtil.eq via Litmus.THROW) throw AssertionError directly via Java + // code rather than via the `assert` keyword, so JVM -da doesn't gate them. + // If one fires inside a search thread, OpenSearchUncaughtExceptionHandler + // exits the cluster JVM. Convert to IllegalStateException so the analytics- + // engine error path treats it as a normal per-query failure (HTTP 500 with + // a bucketable message) instead of taking down the cluster. + throw new IllegalStateException("Substrait conversion rejected the plan: " + e.getMessage(), e); + } List fieldNames = root.fields.stream().map(field -> field.getValue()).toList(); 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 f6713b005393f..d72823015dae5 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 @@ -171,7 +171,8 @@ private static SimpleExtension.ExtensionCollection loadSubstraitExtensions() { t.setContextClassLoader(DataFusionPlugin.class.getClassLoader()); SimpleExtension.ExtensionCollection delegationExtensions = SimpleExtension.load(List.of("/delegation_functions.yaml")); SimpleExtension.ExtensionCollection scalarExtensions = SimpleExtension.load(List.of("/opensearch_scalar_functions.yaml")); - return DefaultExtensionCatalog.DEFAULT_COLLECTION.merge(delegationExtensions).merge(scalarExtensions); + SimpleExtension.ExtensionCollection arrayExtensions = SimpleExtension.load(List.of("/opensearch_array_functions.yaml")); + return DefaultExtensionCatalog.DEFAULT_COLLECTION.merge(delegationExtensions).merge(scalarExtensions).merge(arrayExtensions); } finally { t.setContextClassLoader(previous); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeArrayAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeArrayAdapter.java new file mode 100644 index 0000000000000..672433d87a8b1 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MakeArrayAdapter.java @@ -0,0 +1,89 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.util.ArrayList; +import java.util.List; + +/** + * Rename adapter for PPL's {@code array(a, b, …)} constructor — rewrites to a + * locally-declared {@link SqlFunction} named {@code make_array}, which is + * DataFusion's native array constructor (no UDF registration required on the + * Rust side; isthmus emits a Substrait scalar function call with that name and + * DataFusion's substrait consumer maps it to {@code make_array} natively). + * + *

Unlike {@link org.opensearch.analytics.spi.AbstractNameMappingAdapter}, + * this adapter also CASTs each operand to the array's inferred element type + * before emission. PPL's {@code ArrayFunctionImpl} returns + * {@code ARRAY} (Calcite type-widens to find the common + * element type), but it does NOT widen the individual operand types — so a + * call like {@code array(1, 1.5)} produces a RexCall whose operand types are + * {@code (INTEGER, DECIMAL(2,1))} but whose return type is {@code ARRAY}. + * Substrait's variadic {@code make_array(any1)} signature requires consistent + * argument types ({@link io.substrait.expression.VariadicParameterConsistencyValidator}) + * and throws an AssertionError that fatally exits the JVM otherwise — so we + * widen each operand to the call's component type before substrait sees it. + * + *

Same machinery as {@link UnixTimestampAdapter}: locally-declared operator + * is the referent of the {@link io.substrait.isthmus.expression.FunctionMappings.Sig} + * in {@link DataFusionFragmentConvertor#ADDITIONAL_SCALAR_SIGS}. + * + * @opensearch.internal + */ +class MakeArrayAdapter implements ScalarFunctionAdapter { + + /** + * Locally-declared target operator. Name matches DataFusion's native {@code make_array}. + * Return type inference is a placeholder — {@link #adapt} explicitly carries the + * original call's array return type forward. + */ + static final SqlOperator LOCAL_MAKE_ARRAY_OP = new SqlFunction( + "make_array", + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.SYSTEM + ); + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + RelDataType arrayType = original.getType(); + RelDataType elementType = arrayType.getComponentType(); + if (elementType == null) { + // Defensive — Calcite's array() always infers a component type. If somehow + // missing, fall through with original operands and let substrait fail. + return rexBuilder.makeCall(arrayType, LOCAL_MAKE_ARRAY_OP, original.getOperands()); + } + List widened = new ArrayList<>(original.getOperands().size()); + for (RexNode operand : original.getOperands()) { + if (operand.getType().equals(elementType)) { + widened.add(operand); + } else { + widened.add(rexBuilder.makeCast(elementType, operand, true, false)); + } + } + return rexBuilder.makeCall(arrayType, LOCAL_MAKE_ARRAY_OP, widened); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvappendAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvappendAdapter.java new file mode 100644 index 0000000000000..ac6dcb3ff4e81 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvappendAdapter.java @@ -0,0 +1,97 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.util.ArrayList; +import java.util.List; + +/** + * Rename + operand-coerce adapter for PPL's {@code mvappend(arg1, arg2, …)}. + * + *

The Rust UDF (`udf::mvappend`) handles operands as a uniform stream where + * every operand is either {@code element_type} (scalar) or + * {@code List} (array) for a single inferred {@code element_type}. + * The Calcite call's return type is {@code ARRAY}; this adapter + * casts each scalar operand to {@code componentType} and each array operand to + * {@code ARRAY} before substrait emission, so the UDF sees a + * single element type across all positions. + * + *

Mixed-type {@code mvappend} calls (PPL widens to {@code ARRAY}) end + * up with a Calcite {@code ANY} component type which substrait can't serialize + * — those fail at substrait conversion before reaching this adapter, and + * aren't handled by it. + * + *

Same templated machinery as {@link MvzipAdapter} / {@link MvfindAdapter}: + * the locally-declared operator is the referent of the + * {@link io.substrait.isthmus.expression.FunctionMappings.Sig} entry in + * {@link DataFusionFragmentConvertor#ADDITIONAL_SCALAR_SIGS}. + * + * @opensearch.internal + */ +class MvappendAdapter implements ScalarFunctionAdapter { + + static final SqlOperator LOCAL_MVAPPEND_OP = new SqlFunction( + "mvappend", + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.SYSTEM + ); + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + RelDataType arrayType = original.getType(); + RelDataType componentType = arrayType.getComponentType(); + if (componentType == null) { + return rexBuilder.makeCall(arrayType, LOCAL_MVAPPEND_OP, original.getOperands()); + } + // Substrait's variadic {@code any1} parameter requires every operand at the same + // variadic position to share a type. PPL's {@code mvappend(arg, …)} accepts a mix + // of bare scalars and arrays, which substrait's signature matcher rejects with + // {@code Unable to convert call mvappend(list<…>, scalar, …)}. Normalize every + // operand to {@code ARRAY} — array operands cast their element + // type if it differs; scalar operands wrap in a {@code make_array(…)} singleton + // call. The Rust UDF then sees a uniform {@code list} variadic. + RelDataType targetArrayType = cluster.getTypeFactory().createArrayType(componentType, -1); + List coerced = new ArrayList<>(original.getOperands().size()); + for (RexNode operand : original.getOperands()) { + RelDataType operandType = operand.getType(); + if (operandType.getComponentType() != null) { + // Array operand — cast to ARRAY if its element type differs. + if (operandType.equals(targetArrayType)) { + coerced.add(operand); + } else { + coerced.add(rexBuilder.makeCast(targetArrayType, operand, true, false)); + } + } else { + // Scalar operand — first cast to componentType (so the singleton array's + // element type matches), then wrap in make_array so substrait sees a list. + RexNode casted = operandType.equals(componentType) ? operand : rexBuilder.makeCast(componentType, operand, true, false); + coerced.add(rexBuilder.makeCall(targetArrayType, MakeArrayAdapter.LOCAL_MAKE_ARRAY_OP, List.of(casted))); + } + } + return rexBuilder.makeCall(arrayType, LOCAL_MVAPPEND_OP, coerced); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvfindAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvfindAdapter.java new file mode 100644 index 0000000000000..3a441bbf52b5f --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvfindAdapter.java @@ -0,0 +1,67 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.util.List; + +/** + * Rename adapter for PPL's {@code mvfind(arr, regex)} — rewrites the Calcite + * UDF call (PPL's {@code MVFindFunctionImpl} registered under the function + * name {@code "mvfind"}) to a locally-declared {@link SqlFunction} also named + * {@code mvfind}. The locally-declared op is the referent of the + * {@link io.substrait.isthmus.expression.FunctionMappings.Sig} entry in + * {@link DataFusionFragmentConvertor#ADDITIONAL_SCALAR_SIGS}, so isthmus + * emits a Substrait scalar function call with that exact name. The + * analytics-backend-datafusion plugin's Rust crate (`udf::mvfind`) registers + * a matching ScalarUDF on the DataFusion session context with the same name, + * which the substrait consumer resolves natively. + * + *

The PPL UDF's Calcite-side return type is already {@code INTEGER NULLABLE} + * ({@code MVFindFunctionImpl.getReturnTypeInference()} returns + * {@code ReturnTypes.INTEGER_NULLABLE}), matching the {@code i32?} declared + * in {@code opensearch_array_functions.yaml}. No operand widening is needed — + * the Rust UDF accepts any list element type and any string flavor for the + * regex pattern. + * + * @opensearch.internal + */ +class MvfindAdapter implements ScalarFunctionAdapter { + + /** + * Locally-declared target operator. Name matches the Rust UDF + * {@code MvfindUdf::name()}. + */ + static final SqlOperator LOCAL_MVFIND_OP = new SqlFunction( + "mvfind", + SqlKind.OTHER_FUNCTION, + ReturnTypes.INTEGER_NULLABLE, + null, + OperandTypes.ANY_ANY, + SqlFunctionCategory.SYSTEM + ); + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + return rexBuilder.makeCall(original.getType(), LOCAL_MVFIND_OP, original.getOperands()); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvzipAdapter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvzipAdapter.java new file mode 100644 index 0000000000000..22164425fb34f --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/MvzipAdapter.java @@ -0,0 +1,68 @@ +/* + * 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.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlFunction; +import org.apache.calcite.sql.SqlFunctionCategory; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.type.OperandTypes; +import org.apache.calcite.sql.type.ReturnTypes; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.ScalarFunctionAdapter; + +import java.util.List; + +/** + * Rename adapter for PPL's {@code mvzip(left, right [, sep])} — rewrites the + * Calcite UDF call (PPL's {@code MVZipFunctionImpl} registered under the + * function name {@code "mvzip"}) to a locally-declared {@link SqlFunction} + * also named {@code mvzip}. The locally-declared op is the referent of the + * {@link io.substrait.isthmus.expression.FunctionMappings.Sig} entry in + * {@link DataFusionFragmentConvertor#ADDITIONAL_SCALAR_SIGS}, so isthmus + * emits a Substrait scalar function call with that exact name. The + * analytics-backend-datafusion plugin's Rust crate (`udf::mvzip`) registers a + * matching ScalarUDF on the DataFusion session context with the same name, + * which the substrait consumer resolves natively. + * + *

The PPL UDF's Calcite-side return type is already + * {@code ARRAY<VARCHAR>} (set by {@code MVZipFunctionImpl.getReturnTypeInference}), + * matching the {@code list<string?>} declared in + * {@code opensearch_array_functions.yaml}. No operand widening is needed — + * mvzip accepts any pair of array element types and emits strings. + * + * @opensearch.internal + */ +class MvzipAdapter implements ScalarFunctionAdapter { + + /** + * Locally-declared target operator. Name matches the Rust UDF + * {@code MvzipUdf::name()}. Return-type inference here is a placeholder — + * the call's original return type ({@code ARRAY<VARCHAR>}) is carried + * forward explicitly in {@link #adapt}. + */ + static final SqlOperator LOCAL_MVZIP_OP = new SqlFunction( + "mvzip", + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0, + null, + OperandTypes.VARIADIC, + SqlFunctionCategory.SYSTEM + ); + + @Override + public RexNode adapt(RexCall original, List fieldStorage, RelOptCluster cluster) { + RexBuilder rexBuilder = cluster.getRexBuilder(); + return rexBuilder.makeCall(original.getType(), LOCAL_MVZIP_OP, original.getOperands()); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_array_functions.yaml b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_array_functions.yaml new file mode 100644 index 0000000000000..41361ea3a4acc --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_array_functions.yaml @@ -0,0 +1,158 @@ +%YAML 1.2 +--- +# Substrait extension declaring the array-producing and array-consuming scalar +# functions DataFusion's native runtime can execute. Substrait's standard +# extension catalog has no array_* entries, so isthmus' RexExpressionConverter +# would fail with "Unable to convert call …" until we declare them here. +# +# DataFusion's `datafusion-substrait` consumer resolves these names to native +# DataFusion implementations (datafusion-functions-array crate): +# make_array → array constructor +# array_length → array length +# array_slice → array slice (1-based, inclusive) +# array_distinct → array distinct elements +# array_to_string → join array elements with a separator +urn: extension:org.opensearch:array_functions +scalar_functions: + - name: make_array + description: >- + Construct an array literal from variadic operands. All operands must share + a common element type (Calcite type-widens at the operator level before + emission). Returns a list of that element type. + impls: + - args: + - value: any1 + name: element + variadic: + min: 0 + return: "list" + + - name: array_length + description: >- + Return the number of elements in the array, or NULL if the array is NULL. + Calcite's {@code SqlLibraryOperators.ARRAY_LENGTH} lowers to this name. + impls: + - args: + - value: "list" + name: array + return: "i64?" + + - name: array_slice + description: >- + Return a sub-array slice [from, to] (1-based, inclusive on both ends). + Calcite's {@code SqlLibraryOperators.ARRAY_SLICE} lowers to this name. + impls: + - args: + - value: "list" + name: array + - value: "i64" + name: from + - value: "i64" + name: to + return: "list" + - args: + - value: "list" + name: array + - value: "i32" + name: from + - value: "i32" + name: to + return: "list" + + - name: array_distinct + description: >- + Return the array with duplicate elements removed (preserving first occurrence). + Calcite's {@code SqlLibraryOperators.ARRAY_DISTINCT} lowers to this name. + impls: + - args: + - value: "list" + name: array + return: "list" + + - name: array_element + description: >- + Return the element at the given 1-based position. Calcite's + {@code SqlStdOperatorTable.ITEM} (used by PPL's {@code mvindex(arr, N)} + single-element form via {@code MVIndexFunctionImp.resolveSingleElement}) + renames to this for DataFusion. Returns null if the index is out of range. + impls: + - args: + - value: "list" + name: array + - value: "i64" + name: index + return: "any1?" + + - name: mvappend + description: >- + Flatten a list of arrays into one array, dropping null arrays and null + elements within array arguments. Returns NULL if no non-null elements + were collected. PPL surface is {@code mvappend(arg1, arg2, …)} which + accepts mixed scalar+array operands; the Java adapter wraps each + scalar in a singleton {@code make_array(…)} call so by the time the + Rust UDF sees the operands they're uniformly arrays. Backed by a custom + Rust UDF on the analytics-backend-datafusion plugin (DataFusion's + array_concat preserves nulls — different semantics). + impls: + - args: + - value: "list" + name: arg + variadic: + min: 1 + return: "list" + + - name: mvfind + description: >- + Find the 0-based index of the first array element matching a regex pattern, + or NULL if no match. NULL elements are skipped (not matched). PPL surface is + {@code mvfind(arr, regex)}; registered as a custom Rust UDF on the + analytics-backend-datafusion plugin (no DataFusion stdlib equivalent). + impls: + - args: + - value: "list" + name: array + - value: "string" + name: pattern + return: "i32?" + + - name: mvzip + description: >- + Element-wise zip of two arrays into a list of strings, joined per pair + by a separator (default ","). Result length is min(len(left), len(right)) + (Python-zip truncation). Element NULLs render as empty strings; either + array NULL → NULL result. PPL surface is {@code mvzip(left, right [, sep])}; + registered as a custom Rust UDF on the analytics-backend-datafusion plugin + (no DataFusion stdlib equivalent). + impls: + - args: + - value: "list" + name: left + - value: "list" + name: right + return: "list" + - args: + - value: "list" + name: left + - value: "list" + name: right + - value: "string" + name: separator + return: "list" + + - name: array_to_string + description: >- + Join array elements into a single string using a separator. Calcite's + {@code SqlLibraryOperators.ARRAY_JOIN} renames to this for DataFusion. + impls: + - args: + - value: "list" + name: array + - value: "string" + name: separator + return: "string?" + - args: + - value: "list" + name: array + - value: "varchar" + name: separator + return: "string?" diff --git a/sandbox/plugins/analytics-engine/build.gradle b/sandbox/plugins/analytics-engine/build.gradle index 0058df4b6eb68..78fe8da9d709d 100644 --- a/sandbox/plugins/analytics-engine/build.gradle +++ b/sandbox/plugins/analytics-engine/build.gradle @@ -104,9 +104,14 @@ dependencies { // so bundle it into analytics-engine's own zip. runtimeOnly "org.apache.commons:commons-math3:3.6.1" - // commons-text — Calcite's SqlFunctions. references - // org.apache.commons.text.similarity.LevenshteinDistance. Must be loaded via - // the same classloader as calcite-core so that succeeds. + // commons-text — Calcite's SqlFunctions class statically references + // org.apache.commons.text.similarity.LevenshteinDistance (used by SQL fuzzy-match + // helpers, also pulled in transitively when constant-folding array literals via + // ReduceExpressionsRule). Must be loaded via the same classloader as calcite-core + // so that SqlFunctions. succeeds; otherwise it throws NoClassDefFoundError + // on first use and poisons every subsequent Calcite operation in the JVM — symptom + // is a single failing analytics query taking the cluster's planner thread offline + // for the rest of the run. runtimeOnly "org.apache.commons:commons-text:1.11.0" // httpcore5/httpclient5 — Avatica's BuiltInConnectionProperty static initializer references diff --git a/sandbox/plugins/analytics-engine/licenses/commons-text-1.11.0.jar.sha1 b/sandbox/plugins/analytics-engine/licenses/commons-text-1.11.0.jar.sha1 index 6f090739df8c6..c7b597f6550e0 100644 --- a/sandbox/plugins/analytics-engine/licenses/commons-text-1.11.0.jar.sha1 +++ b/sandbox/plugins/analytics-engine/licenses/commons-text-1.11.0.jar.sha1 @@ -1 +1 @@ -2bb044b7717ec2eccaf9ea7769c1509054b50e9a \ No newline at end of file +2bb044b7717ec2eccaf9ea7769c1509054b50e9a 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 6b0ee31226f39..2a944451363cd 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,9 +10,12 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.complex.ListVector; import org.apache.arrow.vector.util.Text; 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 @@ -35,10 +38,22 @@ public static Object toJavaValue(FieldVector vector, int index) { if (vector instanceof VarCharVector v) { return new String(v.get(index), StandardCharsets.UTF_8); } - Object obj = vector.getObject(index); - if (obj instanceof Text t) { + Object value = vector.getObject(index); + if (vector instanceof ListVector && value instanceof List raw) { + // ListVector.getObject returns a JsonStringArrayList whose elements are the + // child vector's typed values. For VarCharVector children that's Arrow's + // Text, which downstream consumers (e.g. {@code ExprValueUtils.fromObjectValue}) + // don't recognize and reject as "unsupported object class". Mirror the + // top-level VarCharVector branch above and substitute Java strings. + List normalized = new ArrayList<>(raw.size()); + for (Object element : raw) { + normalized.add(element instanceof Text t ? t.toString() : element); + } + return normalized; + } + if (value instanceof Text t) { return t.toString(); } - return obj; + return value; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ArrowSchemaFromCalcite.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ArrowSchemaFromCalcite.java index 2c599b96dc531..f81a9bd1e2951 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ArrowSchemaFromCalcite.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ArrowSchemaFromCalcite.java @@ -40,12 +40,35 @@ private ArrowSchemaFromCalcite() {} public static Schema arrowSchemaFromRowType(RelDataType rowType) { List fields = new ArrayList<>(); for (RelDataTypeField f : rowType.getFieldList()) { - ArrowType arrowType = toArrowType(f.getType().getSqlTypeName()); - fields.add(new Field(f.getName(), new FieldType(true, arrowType, null), null)); + fields.add(toArrowField(f.getName(), f.getType())); } return new Schema(fields); } + /** + * Build an Arrow {@link Field} from a Calcite type. For scalar types this is a + * leaf field with the appropriate {@link ArrowType}; for ARRAY this is a + * {@code List} whose single child is the recursively-converted element type + * (Arrow names the child {@code $data$} by convention — kept here for parity with + * Arrow's own builders so downstream tooling that walks list children by name + * doesn't break). + */ + private static Field toArrowField(String name, RelDataType type) { + SqlTypeName sqlTypeName = type.getSqlTypeName(); + if (sqlTypeName == SqlTypeName.ARRAY) { + RelDataType elementType = type.getComponentType(); + if (elementType == null) { + throw new IllegalArgumentException( + "ARRAY type with no component type for field [" + name + "]; cannot derive list element schema" + ); + } + Field elementField = toArrowField("$data$", elementType); + return new Field(name, new FieldType(true, ArrowType.List.INSTANCE, null), List.of(elementField)); + } + ArrowType arrowType = toArrowType(sqlTypeName); + return new Field(name, new FieldType(true, arrowType, null), null); + } + private static ArrowType toArrowType(SqlTypeName sqlTypeName) { switch (sqlTypeName) { case BIGINT: 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..97aa10ff82ca2 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; @@ -28,6 +30,7 @@ import java.nio.charset.StandardCharsets; import java.util.ArrayList; +import java.util.Collections; import java.util.List; /** @@ -60,8 +63,7 @@ public VectorSchemaRoot decode(FragmentExecutionResponse response, BufferAllocat // Infer Arrow type per column from the first non-null value 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); @@ -86,29 +88,76 @@ public VectorSchemaRoot decode(FragmentExecutionResponse response, BufferAllocat } /** - * 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 - * are null or the Java type is unrecognized. + * Infers the Arrow {@link Field} for a column by scanning rows for the first + * non-null value. Falls back to a nullable {@code Utf8} (VarChar) field if all + * values are null or the Java type is unrecognized. + * + *

For {@link List} cells (produced by analytics-engine routes that emit + * array-typed values — PPL {@code array(...)}, {@code array_slice}, …), this + * returns a {@code List} field where the inner element type is inferred + * from the first non-null element. Without this branch, list values fall + * through to the {@code Utf8} fallback and {@link #setVectorValue} produces + * {@code value.toString()} (e.g. {@code "[2,3,4]"} as a JSON-like string) + * instead of a typed array. */ - static ArrowType inferArrowType(List rows, int col) { + static Field inferField(String name, List rows, int col) { for (Object[] row : rows) { Object value = row[col]; if (value == null) continue; - if (value instanceof Long) return new ArrowType.Int(64, true); - if (value instanceof Integer) return new ArrowType.Int(32, true); - if (value instanceof Short) return new ArrowType.Int(16, true); - if (value instanceof Byte) return new ArrowType.Int(8, true); - if (value instanceof Double) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE); - if (value instanceof Float) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE); - 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 Number) return new ArrowType.Int(64, true); - break; + if (value instanceof List list) { + ArrowType elementType = inferElementArrowType(list, rows, col); + Field elementField = new Field("$data$", FieldType.nullable(elementType), null); + return new Field(name, FieldType.nullable(ArrowType.List.INSTANCE), Collections.singletonList(elementField)); + } + return new Field(name, FieldType.nullable(scalarArrowType(value)), null); + } + return new Field(name, FieldType.nullable(ArrowType.Utf8.INSTANCE), null); + } + + /** + * Best-effort inference for a list element type. Looks at the first non-null + * element of the given list, then falls back to scanning later rows of the + * same column if this list is empty or all-null. Defaults to {@code Utf8}. + */ + private static ArrowType inferElementArrowType(List list, List rows, int col) { + for (Object element : list) { + if (element != null) return scalarArrowType(element); + } + for (Object[] row : rows) { + Object value = row[col]; + if (value instanceof List other) { + for (Object element : other) { + if (element != null) return scalarArrowType(element); + } + } } return ArrowType.Utf8.INSTANCE; } + /** Maps a Java scalar value to the corresponding Arrow scalar type. */ + private static ArrowType scalarArrowType(Object value) { + if (value instanceof Long) return new ArrowType.Int(64, true); + if (value instanceof Integer) return new ArrowType.Int(32, true); + if (value instanceof Short) return new ArrowType.Int(16, true); + if (value instanceof Byte) return new ArrowType.Int(8, true); + if (value instanceof Double) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE); + if (value instanceof Float) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE); + // BigDecimal must be checked before the Number fallback below — a BigDecimal that + // would round to a long but actually carries fractional precision (e.g. PPL + // {@code array(1, -1.5)} where the common element type is DECIMAL and + // {@code ArrayImplementor.internalCast} produces BigDecimal cells) would otherwise + // get encoded as an integer Arrow vector and lose its fractional digits. Promote to + // DOUBLE — the same path the v2 engine takes for decimal-typed PPL results. + if (value instanceof java.math.BigDecimal) return new ArrowType.FloatingPoint( + org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE + ); + 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 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,64 @@ 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) { + writeListValue(listVector, index, (List) value); } else { throw new IllegalArgumentException("Unsupported Arrow vector type: " + vector.getClass().getSimpleName()); } } + + /** + * Writes a Java {@link List} into an Arrow {@link ListVector} at the given row + * index. Bypasses {@link UnionListWriter} entirely — writes directly to the + * list's offset / validity buffers and to the inner data vector via the inner + * vector's own typed setter. The writer-based API requires an + * {@link org.apache.arrow.memory.ArrowBuf} per varchar element which Arrow + * couples to a release lifecycle that's tricky to get right (early close → + * use-after-free, no close → leak); the direct path avoids that altogether. + * + *

The inner data vector's type was decided in {@link #inferField} from the + * first non-null list element, so the {@code instanceof} dispatch here simply + * needs to match. + */ + private static void writeListValue(ListVector listVector, int index, List list) { + FieldVector dataVector = listVector.getDataVector(); + int startOffset = listVector.getOffsetBuffer().getInt((long) index * ListVector.OFFSET_WIDTH); + int writePos = startOffset; + for (Object element : list) { + // Grow the data vector if needed before writing. setSafe-style helpers handle this + // automatically per type, but we still need to pre-position the cursor. + if (element == null) { + dataVector.setNull(writePos); + } else if (dataVector instanceof BigIntVector v) { + v.setSafe(writePos, ((Number) element).longValue()); + } else if (dataVector instanceof IntVector v) { + v.setSafe(writePos, ((Number) element).intValue()); + } else if (dataVector instanceof SmallIntVector v) { + v.setSafe(writePos, ((Number) element).shortValue()); + } else if (dataVector instanceof TinyIntVector v) { + v.setSafe(writePos, ((Number) element).byteValue()); + } else if (dataVector instanceof Float8Vector v) { + v.setSafe(writePos, ((Number) element).doubleValue()); + } else if (dataVector instanceof Float4Vector v) { + v.setSafe(writePos, ((Number) element).floatValue()); + } else if (dataVector instanceof BitVector v) { + v.setSafe(writePos, ((Boolean) element) ? 1 : 0); + } else if (dataVector instanceof VarCharVector v) { + v.setSafe(writePos, element.toString().getBytes(StandardCharsets.UTF_8)); + } else if (dataVector instanceof VarBinaryVector v) { + v.setSafe(writePos, (byte[]) element); + } else { + throw new IllegalArgumentException("Unsupported list element vector type: " + dataVector.getClass().getSimpleName()); + } + writePos++; + } + // Mark this row's list as non-null and update its end offset. + listVector.setNotNull(index); + listVector.getOffsetBuffer().setInt((long) (index + 1) * ListVector.OFFSET_WIDTH, writePos); + // Keep the data vector's value count in sync so subsequent reads see the new tail. + if (writePos > dataVector.getValueCount()) { + dataVector.setValueCount(writePos); + } + } } diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ArrayFunctionIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ArrayFunctionIT.java new file mode 100644 index 0000000000000..19cb0b076809b --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ArrayFunctionIT.java @@ -0,0 +1,311 @@ +/* + * 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.Arrays; +import java.util.List; +import java.util.Map; + +/** + * End-to-end coverage for the PPL array-construction and multivalue (mv*) + * functions on the analytics-engine route (PPL → CalciteRelNodeVisitor → + * Substrait → DataFusion). Mirrors the SQL plugin's + * {@code CalciteArrayFunctionIT} one-test-method-to-one for the subset of tests + * the analytics-engine path supports today. + * + *

Function surface exercised: + *

    + *
  • {@code array(...)} → DataFusion {@code make_array} via + * {@link org.opensearch.be.datafusion.MakeArrayAdapter}.
    • + *
  • {@code array_length} → DataFusion native {@code array_length}.
    • + *
  • {@code mvindex(arr, from, to)} (range form) → DataFusion {@code array_slice} + * via {@link org.opensearch.be.datafusion.ArraySliceAdapter} (BIGINT index + * coerce + 0-based-{@code (start, length)} → 1-based-{@code (start, end)}).
    • + *
  • {@code mvindex(arr, N)} (single-element form) → DataFusion {@code array_element} + * via {@link org.opensearch.be.datafusion.ArrayElementAdapter}.
    • + *
  • {@code mvdedup(arr)} → DataFusion native {@code array_distinct}.
    • + *
  • {@code mvjoin(arr, sep)} → DataFusion {@code array_to_string} via + * {@link org.opensearch.be.datafusion.ArrayToStringAdapter}.
    • + *
  • {@code mvzip(left, right [, sep])} → custom Rust UDF {@code udf::mvzip}.
    • + *
  • {@code mvfind(arr, regex)} → custom Rust UDF {@code udf::mvfind}.
    • + *
  • {@code split(str, delim)} (returns array) → DataFusion {@code string_to_array}.
    • + *
+ * + *

The {@code calcs} dataset is used as a scan target; most tests build literal + * arrays inside {@code eval} so the field types don't matter — what matters is + * that the source is a parquet-backed index the analytics-engine planner can + * scan. + * + *

Tests for lambda-based functions ({@code transform}, {@code mvmap}, + * {@code reduce}, {@code forall}, {@code exists}, {@code filter}) are + * intentionally absent: substrait extension YAML doesn't support declaring + * {@code func<…>} lambda-typed arguments, so those don't ship through the + * analytics-engine route in this PR. Empty-array tests are also absent — + * {@code array()} defaults to {@code ARRAY[UNKNOWN]} which substrait can't + * encode without the SQL companion {@code #5421} default to {@code VARCHAR}. + */ +public class ArrayFunctionIT 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; + } + } + + /** Base query template: pin to one row so every assertion runs against a single result row. */ + private String oneRow() { + return "source=" + DATASET.indexName + " | head 1 "; + } + + // ── array(...) constructor ────────────────────────────────────────────── + + /** Mixed-numeric literal array — exercises the BigDecimal → Double row-codec + * promotion (without it, decimal cells truncate to integers). */ + public void testArray() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = array(1, -1.5, 2, 1.0) | fields result", + Arrays.asList(1.0, -1.5, 2.0, 1.0)); + } + + /** Mixed int+string literal array — Calcite widens to {@code ARRAY} + * via {@code ArrayFunctionImpl.internalCast}. */ + public void testArrayWithString() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = array(1, 'demo') | fields result", + Arrays.asList("1", "demo")); + } + + // ── array_length ──────────────────────────────────────────────────────── + + public void testArrayLength() throws IOException { + assertFirstRowDouble( + oneRow() + "| eval arr = array(1, -1.5, 2, 1.0) | eval len = array_length(arr) | fields len", + 4.0); + } + + // ── mvindex range (array_slice) ───────────────────────────────────────── + + /** {@code mvindex(arr, 1, 3)} — 0-based-(start, length) → DataFusion 1-based-(start, end inclusive) + * via {@link org.opensearch.be.datafusion.ArraySliceAdapter}. Without the rewrite the result + * would be {@code [1, 2, 3]} instead of the expected {@code [2, 3, 4]}. */ + public void testMvindexRangePositive() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array(1, 2, 3, 4, 5) | eval result = mvindex(arr, 1, 3) | fields result", + Arrays.asList(2, 3, 4)); + } + + /** Negative indices — DataFusion's array_slice supports them natively. */ + public void testMvindexRangeNegative() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array(1, 2, 3, 4, 5) | eval result = mvindex(arr, -3, -1) | fields result", + Arrays.asList(3, 4, 5)); + } + + public void testMvindexRangeFirstThree() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array(10, 20, 30, 40, 50) | eval result = mvindex(arr, 0, 2) | fields result", + Arrays.asList(10, 20, 30)); + } + + // ── mvindex single (array_element) ────────────────────────────────────── + + /** {@code mvindex(arr, N)} with a single index — PPL emits Calcite's + * {@code SqlStdOperatorTable.ITEM} which {@link org.opensearch.be.datafusion.ArrayElementAdapter} + * renames to DataFusion {@code array_element} with a BIGINT-coerced 1-based index. */ + public void testMvindexSingleElementPositive() throws IOException { + assertFirstRowDouble( + oneRow() + "| eval arr = array(10, 20, 30) | eval result = mvindex(arr, 1) | fields result", + 20.0); + } + + public void testMvindexSingleElementNegative() throws IOException { + assertFirstRowDouble( + oneRow() + "| eval arr = array(10, 20, 30) | eval result = mvindex(arr, -1) | fields result", + 30.0); + } + + // ── mvdedup (array_distinct) ──────────────────────────────────────────── + + public void testMvdedupWithDuplicates() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array(1, 2, 2, 3, 3, 3) | eval result = mvdedup(arr) | fields result", + Arrays.asList(1, 2, 3)); + } + + public void testMvdedupWithStrings() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array('a', 'b', 'a', 'c', 'b') | eval result = mvdedup(arr) | fields result", + Arrays.asList("a", "b", "c")); + } + + public void testMvdedupAllDuplicates() throws IOException { + assertFirstRowList( + oneRow() + "| eval arr = array(7, 7, 7) | eval result = mvdedup(arr) | fields result", + Arrays.asList(7)); + } + + // ── mvjoin (array_to_string) ──────────────────────────────────────────── + + public void testMvjoinWithStringArray() throws IOException { + assertFirstRowString( + oneRow() + "| eval result = mvjoin(array('a', 'b', 'c'), ',') | fields result", + "a,b,c"); + } + + public void testMvjoinWithStringifiedNumbers() throws IOException { + assertFirstRowString( + oneRow() + "| eval result = mvjoin(array('1', '2', '3'), ' | ') | fields result", + "1 | 2 | 3"); + } + + public void testMvjoinWithSpecialDelimiters() throws IOException { + assertFirstRowString( + oneRow() + "| eval result = mvjoin(array('x', 'y'), '-->') | fields result", + "x-->y"); + } + + // ── mvzip (Rust UDF) ──────────────────────────────────────────────────── + + public void testMvzipBasic() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = mvzip(array('a', 'b', 'c'), array('1', '2', '3')) | fields result", + Arrays.asList("a,1", "b,2", "c,3")); + } + + public void testMvzipWithCustomDelimiter() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = mvzip(array('a', 'b'), array('1', '2'), '-') | fields result", + Arrays.asList("a-1", "b-2")); + } + + public void testMvzipNested() throws IOException { + assertFirstRowList( + oneRow() + + "| eval r = mvzip(mvzip(array('a','b'), array('1','2')), array('x','y')) | fields r", + Arrays.asList("a,1,x", "b,2,y")); + } + + // ── mvfind (Rust UDF) ─────────────────────────────────────────────────── + + /** Returns the 0-based index of the first array element matching the regex. */ + public void testMvfindWithMatch() throws IOException { + assertFirstRowDouble( + oneRow() + "| eval result = mvfind(array('apple', 'banana', 'cherry'), 'ban.*') | fields result", + 1.0); + } + + public void testMvfindWithNoMatch() throws IOException { + assertFirstRowNull( + oneRow() + "| eval result = mvfind(array('apple', 'banana'), 'zzz') | fields result"); + } + + /** Dynamic regex — exercises the {@code SqlLibraryOperators.CONCAT_FUNCTION} → substrait + * {@code concat} Sig bridge added in this PR. Without that bridge the call fails substrait + * conversion with {@code Unable to convert call CONCAT(string, string)}. */ + public void testMvfindWithDynamicRegex() throws IOException { + assertFirstRowDouble( + oneRow() + + "| eval result = mvfind(array('apple', 'banana', 'cherry'), concat('ban', '.*')) | fields result", + 1.0); + } + + // ── split (returns array of strings) ───────────────────────────────── + + public void testSplitWithSemicolonDelimiter() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = split('a;b;c', ';') | fields result", + Arrays.asList("a", "b", "c")); + } + + public void testSplitWithMultiCharDelimiter() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = split('a::b::c', '::') | fields result", + Arrays.asList("a", "b", "c")); + } + + // ── helpers ───────────────────────────────────────────────────────────── + + /** Numeric-tolerant list comparison — Jackson parses JSON numbers as + * Integer/Long/Double interchangeably, so equality on cross-type numbers + * fails even when values match. Compare via {@link Double#compare} on + * numeric pairs and {@link Object#equals} otherwise. */ + private void assertFirstRowList(String ppl, List expected) throws IOException { + Object cell = firstRowFirstCell(ppl); + assertNotNull("Expected non-null array result for query [" + ppl + "]", cell); + assertTrue( + "Expected list result for query [" + ppl + "] but got: " + cell + " (" + cell.getClass() + ")", + cell instanceof List); + List actual = (List) cell; + assertEquals( + "Length mismatch for query [" + ppl + "]: expected " + expected + " but got " + actual, + expected.size(), + actual.size()); + for (int i = 0; i < expected.size(); i++) { + assertCellEquals(expected.get(i), actual.get(i)); + } + } + + private void assertFirstRowDouble(String ppl, double expected) throws IOException { + Object cell = firstRowFirstCell(ppl); + assertTrue("Expected numeric result for query [" + ppl + "] but got: " + cell, cell instanceof Number); + assertEquals("Value mismatch for query: " + ppl, expected, ((Number) cell).doubleValue(), 1e-9); + } + + private void assertFirstRowString(String ppl, String expected) throws IOException { + Object cell = firstRowFirstCell(ppl); + assertEquals("Value mismatch for query: " + ppl, expected, cell); + } + + private void assertFirstRowNull(String ppl) throws IOException { + Object cell = firstRowFirstCell(ppl); + assertNull("Expected null result for query [" + ppl + "] but got: " + cell, cell); + } + + private static void assertCellEquals(Object expected, Object actual) { + if (expected == null || actual == null) { + assertEquals(expected, actual); + return; + } + if (expected instanceof Number && actual instanceof Number) { + assertEquals( + "Numeric value mismatch", + ((Number) expected).doubleValue(), + ((Number) actual).doubleValue(), + 1e-9); + return; + } + assertEquals(expected, actual); + } + + private Object firstRowFirstCell(String ppl) throws IOException { + Map response = executePpl(ppl); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertNotNull("Response missing 'rows' for query: " + ppl, rows); + assertTrue("Expected at least one row for query: " + ppl, rows.size() >= 1); + return rows.get(0).get(0); + } + + 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/MVAppendFunctionIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MVAppendFunctionIT.java new file mode 100644 index 0000000000000..c4ada7cf538c7 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MVAppendFunctionIT.java @@ -0,0 +1,180 @@ +/* + * 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.Arrays; +import java.util.List; +import java.util.Map; + +/** + * End-to-end coverage for PPL {@code mvappend(arg1, arg2, …)} on the + * analytics-engine route. Mirrors the SQL plugin's + * {@code CalciteMVAppendFunctionIT} one-test-method-to-one for the subset of + * tests that pass on the analytics-engine path. + * + *

{@code mvappend} flattens an arbitrary mix of scalar and array operands + * into a single array, dropping null elements. Onboarded as a custom Rust UDF + * ({@code udf::mvappend}) registered at session-context creation; the Java + * adapter ({@link org.opensearch.be.datafusion.MvappendAdapter}) reshapes scalar + * operands into singleton {@code make_array} calls so substrait's variadic-{@code any1} + * shape sees a uniform {@code list[componentType]} across every position. + * + *

Tests covering genuinely heterogeneous mvappend signatures + * ({@code mvappend(1, 'text', 2.5)}, {@code mvappend(age, 'years', 'old')}, + * {@code mvappend('test', nullif(1,1), 2)}) are absent because Calcite legitimately + * widens those to {@code ARRAY[ANY]} — substrait can't encode {@code ANY}, and + * Arrow's Union arrays aren't operated on by {@code datafusion-functions-array}. + * Empty-array operand tests are also absent — the empty {@code array()} default + * surfaces as {@code ARRAY[UNKNOWN]}/{@code ARRAY[VARCHAR]} in the column ref, + * which type-inference can't reach back through the project chain to ignore. + * + *

The {@code testMvappendInWhereClause} variant (filter predicate on an + * ARRAY field) is also absent because the analytics-engine planner's filter + * rule rejects {@code EQUALS} on an ARRAY field without walking into the + * predicate tree — that's a separate planner refactor tracked under #21554's + * "What's left" section. + */ +public class MVAppendFunctionIT 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; + } + } + + private String oneRow() { + return "source=" + DATASET.indexName + " | head 1 "; + } + + // ── uniform-typed scalar variadic ─────────────────────────────────────── + + public void testMvappendWithMultipleElements() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = mvappend(1, 2, 3) | fields result", + Arrays.asList(1, 2, 3)); + } + + public void testMvappendWithSingleElement() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = mvappend(42) | fields result", + Arrays.asList(42)); + } + + public void testMvappendWithStringValues() throws IOException { + assertFirstRowList( + oneRow() + "| eval result = mvappend('hello', 'world') | fields result", + Arrays.asList("hello", "world")); + } + + // ── array operands (uniform element type) ─────────────────────────────── + + public void testMvappendWithArrayFlattening() throws IOException { + assertFirstRowList( + oneRow() + + "| eval arr1 = array(1, 2), arr2 = array(3, 4), result = mvappend(arr1, arr2) | fields result", + Arrays.asList(1, 2, 3, 4)); + } + + public void testMvappendWithNestedArrays() throws IOException { + assertFirstRowList( + oneRow() + + "| eval arr1 = array('a', 'b'), arr2 = array('c'), arr3 = array('d', 'e')," + + " result = mvappend(arr1, arr2, arr3) | fields result", + Arrays.asList("a", "b", "c", "d", "e")); + } + + // ── field references ──────────────────────────────────────────────────── + + /** Two VARCHAR field references → uniform {@code ARRAY[VARCHAR]}. Anchored + * to a specific row by filtering on {@code key} so the assertion is + * deterministic. */ + public void testMvappendWithRealFields() throws IOException { + assertFirstRowList( + "source=" + DATASET.indexName + + " | where key='key00' | head 1 | eval result = mvappend(str0, str1) | fields result", + // calcs row key00: str0='FURNITURE', str1='CLAMP ON LAMPS' + Arrays.asList("FURNITURE", "CLAMP ON LAMPS")); + } + + // ── tests gated on SQL companion #5424 ────────────────────────────────── + // The following SQL-side tests are intentionally absent until + // opensearch-project/sql#5424 (the {@code MVAppendFunctionImpl} widening + // via {@code leastRestrictive} + DECIMAL → DOUBLE promotion + operand + // pre-cast in {@code MVAppendImplementor}) is merged and republished as + // {@code unified-query-core:3.7.0.0-SNAPSHOT}. Without it, these collapse + // to {@code ARRAY[ANY]} which substrait can't encode: + // + // testMvappendWithMixedArrayAndScalar — array(1,2), 3, 4 (nullability bridge) + // testMvappendWithNumericArrays — array(1.5,2.5), array(3.5), 4.5 (nullability bridge) + // testMvappendWithIntAndDouble — 1, 2.5 (DECIMAL → DOUBLE promotion + pre-cast) + // testMvappendWithComplexExpression — array(int0), array(int0*2), int0+10 (nullability bridge) + // + // Add them back once #5424 lands. Their SQL-side counterparts are verified + // in CalciteMVAppendFunctionIT against the analytics-engine route. + + // ── helpers ───────────────────────────────────────────────────────────── + + private void assertFirstRowList(String ppl, List expected) throws IOException { + Object cell = firstRowFirstCell(ppl); + assertNotNull("Expected non-null array result for query [" + ppl + "]", cell); + assertTrue( + "Expected list result for query [" + ppl + "] but got: " + cell + " (" + cell.getClass() + ")", + cell instanceof List); + List actual = (List) cell; + assertEquals( + "Length mismatch for query [" + ppl + "]: expected " + expected + " but got " + actual, + expected.size(), + actual.size()); + for (int i = 0; i < expected.size(); i++) { + assertCellEquals(expected.get(i), actual.get(i)); + } + } + + private static void assertCellEquals(Object expected, Object actual) { + if (expected == null || actual == null) { + assertEquals(expected, actual); + return; + } + if (expected instanceof Number && actual instanceof Number) { + assertEquals( + "Numeric value mismatch", + ((Number) expected).doubleValue(), + ((Number) actual).doubleValue(), + 1e-9); + return; + } + assertEquals(expected, actual); + } + + private Object firstRowFirstCell(String ppl) throws IOException { + Map response = executePpl(ppl); + @SuppressWarnings("unchecked") + List> rows = (List>) response.get("rows"); + assertNotNull("Response missing 'rows' for query: " + ppl, rows); + assertTrue("Expected at least one row for query: " + ppl, rows.size() >= 1); + return rows.get(0).get(0); + } + + 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); + } +}