diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java index 22b755a73772a..3495e67d60659 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkContext.java @@ -9,7 +9,6 @@ package org.opensearch.analytics.spi; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.types.pojo.Schema; import java.util.List; @@ -32,9 +31,10 @@ *
  • {@code childInputs} — one entry per child stage. Each entry carries * the child's stage id (used by the backend to register a per-child * input partition under a stable name like {@code "input-"}) - * and the Arrow schema of the batches the child will feed in. For - * single-input shapes this list has size 1; for {@code UNION}-style - * multi-input shapes it has one entry per Union branch.
    • + * and the producer-side plan bytes (e.g. partial-aggregate substrait) + * the backend lowers to derive the input schema. For single-input + * shapes this list has size 1; for {@code UNION}-style multi-input + * shapes it has one entry per Union branch. *
  • {@code downstream} — sink the backend drains its reduced output * into. The backend owns {@code downstream}'s lifecycle: it must * feed every produced batch and close it when draining is complete.
    • @@ -45,21 +45,11 @@ public record ExchangeSinkContext(String queryId, int stageId, byte[] fragmentBytes, BufferAllocator allocator, List< ChildInput> childInputs, ExchangeSink downstream) implements CommonExecutionContext { - /** Per-child input descriptor: the child stage id and the schema of its outgoing batches. */ - public record ChildInput(int childStageId, Schema schema) { - } - /** - * Convenience for single-input back-compat. Returns the schema of the sole - * child input. Throws when {@link #childInputs} contains more than one entry — - * multi-input callers must inspect {@link #childInputs} directly. + * Per-child input descriptor: the child stage id and the producer-side plan bytes the + * backend lowers when it registers the child's input partition. The actual Arrow schema + * is learned at registration time, not declared here. */ - public Schema inputSchema() { - if (childInputs.size() != 1) { - throw new IllegalStateException( - "inputSchema() requires exactly one child input; got " + childInputs.size() + " — use childInputs() instead" - ); - } - return childInputs.get(0).schema(); + public record ChildInput(int childStageId, byte[] producerPlanBytes) { } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkProvider.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkProvider.java index dcef3717354cd..d56115fbe1d6f 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkProvider.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSinkProvider.java @@ -25,8 +25,13 @@ public interface ExchangeSinkProvider { /** * Creates a sink for coordinator-side execution. The backend implementation * uses {@link ExchangeSinkContext#fragmentBytes()} as the serialized plan - * (produced by {@link FragmentConvertor#convertFinalAggFragment}) and - * writes its reduced output into {@link ExchangeSinkContext#downstream()}. + * (produced by {@code FragmentConvertor#convertFragment}) and writes its + * reduced output into {@link ExchangeSinkContext#downstream()}. + * + *

    The schema of each child's batches is learned at the backend boundary + * (not pre-declared on {@link ExchangeSinkContext.ChildInput}) — the backend + * derives it when it registers the child input on its native session, since + * the producer-side plan bytes already encode the producer schema. * * @param context core-provided context carrying plan bytes, allocator, child inputs, and downstream sink * @param backendContext backend-opaque state produced by instruction handlers (e.g. diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java index 6fec046d03d29..2dd7c7cf0cc6c 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java @@ -21,9 +21,9 @@ * *

    Composable pipeline for multi-shard aggregate with sort at coordinator: *

      - *
    1. {@code convertShardScanFragment(tableName, Filter(Scan))} → data node inner bytes
      2. + *
    2. {@code convertFragment(Filter(Scan))} → data node inner bytes
      3. *
    3. {@code attachPartialAggOnTop(PartialAgg, innerBytes)} → data node bytes
      4. - *
    4. {@code convertFinalAggFragment(FinalAgg(StageInputScan))} → reduce stage inner bytes
      5. + *
    5. {@code convertFragment(FinalAgg(StageInputScan))} → reduce stage inner bytes
      6. *
    6. {@code attachFragmentOnTop(Sort, innerBytes)} → reduce stage bytes
      7. *
    * @@ -35,16 +35,25 @@ public interface FragmentConvertor { /** - * Converts a fragment whose leaf is a native physical shard scan, containing - * everything below a partial aggregate (e.g. Filter(Scan), Scan). - * The backend handles all operators natively — no delegation, no shuffle. + * Converts a resolved RelNode fragment (annotations stripped) into + * backend-specific serialized plan bytes. The fragment may be: + * * - * @param tableName named table the fragment's scan references - * @param fragment resolved RelNode fragment (annotations stripped) + *

    TODO: revisit placement of FragmentConvertor — it references Calcite RelNode + * and is called only by analytics-engine. Consider moving to analytics-engine and + * removing getFragmentConvertor() from AnalyticsSearchBackendPlugin SPI. + * + * @param fragment resolved RelNode fragment * @return backend-specific serialized plan bytes */ - default byte[] convertShardScanFragment(String tableName, RelNode fragment) { - throw new UnsupportedOperationException("convertShardScanFragment not implemented for this backend"); + default byte[] convertFragment(RelNode fragment) { + throw new UnsupportedOperationException("convertFragment not implemented for this backend"); } /** @@ -53,31 +62,13 @@ default byte[] convertShardScanFragment(String tableName, RelNode fragment) { * aggregate execution node. * * @param partialAggFragment the partial aggregate RelNode (annotations stripped, no children) - * @param innerBytes serialized bytes from a prior {@code convert*} call + * @param innerBytes serialized bytes from a prior {@link #convertFragment} call * @return serialized plan bytes with partial aggregate attached on top */ default byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerBytes) { throw new UnsupportedOperationException("attachPartialAggOnTop not implemented for this backend"); } - /** - * Converts the final aggregate fragment at the reduce stage. - * The leaf is a StageInputScan placeholder used for schema inference — - * replaced at execution time with a streaming Arrow batch source - * (e.g. StreamingTableExec in DataFusion). - * - *

    TODO: revisit placement of FragmentConvertor — it references Calcite RelNode - * and is called only by analytics-engine. Consider moving to analytics-engine and - * removing getFragmentConvertor() from AnalyticsSearchBackendPlugin SPI. - * - * @param fragment resolved final aggregate RelNode (annotations stripped, - * ExchangeReducer removed, StageInputScan as leaf) - * @return backend-specific serialized plan bytes - */ - default byte[] convertFinalAggFragment(RelNode fragment) { - throw new UnsupportedOperationException("convertFinalAggFragment not implemented for this backend"); - } - /** * Attaches a generic fragment (Sort, Project, etc.) on top of already-converted * inner bytes. The backend deserializes the inner plan and wraps it with the diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index 7a5cefa33567f..da1e73a67ff84 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -31,12 +31,10 @@ //! - `stream_get_schema`, `stream_close` must NOT be called //! concurrently on the same stream pointer. -use std::io::Cursor; use std::num::NonZeroUsize; use std::path::PathBuf; use std::sync::Arc; -use arrow::ipc::reader::StreamReader; use arrow_array::ffi::FFI_ArrowArray; use arrow_array::RecordBatch; use arrow_array::{Array, StructArray}; @@ -534,6 +532,186 @@ pub unsafe fn sql_to_substrait( }) } +/// Lowers a partial-aggregate Substrait plan against a throwaway session and +/// returns its physical output schema. NamedTable references are resolved +/// against empty MemTables built from the substrait base_schema — the plan +/// itself is the source of truth for the producer side, so no view-type or +/// timestamp-precision rewrites are applied here. The plan is dropped at +/// function exit; only the schema is returned. +fn derive_schema_from_partial_plan( + substrait_bytes: &[u8], +) -> Result { + use datafusion::datasource::MemTable; + use datafusion::prelude::SessionContext; + use datafusion_substrait::extensions::Extensions; + use datafusion_substrait::logical_plan::consumer::{ + from_substrait_named_struct, from_substrait_plan, DefaultSubstraitConsumer, + }; + use prost::Message; + use substrait::proto::{read_rel::ReadType, Plan}; + + let plan = Plan::decode(substrait_bytes).map_err(|e| { + DataFusionError::Execution(format!("derive_schema_from_partial_plan: decode failed: {}", e)) + })?; + + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new()) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + crate::udf::register_all(&ctx); + + let extensions = Extensions::default(); + let session_state = ctx.state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let mut reads = Vec::new(); + for plan_rel in &plan.relations { + if let Some(rel) = root_rel(plan_rel) { + collect_reads(&rel, &mut reads); + } + } + for read in &reads { + let Some(ReadType::NamedTable(nt)) = read.read_type.as_ref() else { + continue; + }; + let table_name = nt.names.last().cloned().unwrap_or_default(); + let base_schema = read.base_schema.as_ref().ok_or_else(|| { + DataFusionError::Execution("ReadRel missing base_schema".to_string()) + })?; + let df_schema = from_substrait_named_struct(&consumer, base_schema)?; + let arrow_schema = df_schema.as_arrow().clone(); + + // Mirror the two transformations the data-node session applies to its + // parquet-read leaf, so the synthetic leaf we register here matches. + // Without these, HashAggregateExec lowers over Utf8 / Timestamp(Second) + // on this throwaway session while the real data-node session lowers + // over Utf8View / Timestamp(Millisecond) — same plan, divergent + // physical outputs, runtime schema-mismatch on the wire. + // + // No data conversion happens at runtime — these only configure the + // coordinator's StreamingTable so it accepts the producer's batches + // without reinterpretation. Long-term plan: have the data node embed + // its lowered output schema as substrait extension metadata so the + // coordinator skips this throwaway lowering and both mirrors evaporate. + let view_types = ctx + .copied_config() + .options() + .execution + .parquet + .schema_force_view_types; + let arrow_schema = if view_types { + datafusion::datasource::file_format::parquet::transform_schema_to_view(&arrow_schema) + } else { + arrow_schema + }; + let arrow_schema = coerce_unsupported_timestamp_precision(&arrow_schema); + + let table = MemTable::try_new(Arc::new(arrow_schema), vec![vec![]])?; + // Plan may scan the same table twice; the second register is a no-op. + let _ = ctx.register_table(&table_name, Arc::new(table)); + } + + let logical_plan = futures::executor::block_on(from_substrait_plan(&session_state, &plan))?; + let physical_plan = futures::executor::block_on(session_state.create_physical_plan(&logical_plan))?; + Ok(physical_plan.schema()) +} + +/// Encodes a Schema as Arrow IPC stream-format bytes (a schema-only message +/// followed by the stream EOS marker). This is the wire format Java reads via +/// `MessageChannelReader` / `ArrowStreamReader`. +fn schema_to_ipc_bytes(schema: &arrow::datatypes::Schema) -> Result, DataFusionError> { + use arrow::ipc::writer::StreamWriter; + let mut buf: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema) + .map_err(|e| DataFusionError::Execution(format!("StreamWriter::try_new: {}", e)))?; + writer + .finish() + .map_err(|e| DataFusionError::Execution(format!("StreamWriter::finish: {}", e)))?; + } + Ok(buf) +} + +/// Mirror parquet's coercion of Arrow Timestamp precisions it cannot +/// represent in its logical type system. Parquet's TIMESTAMP supports +/// MILLIS / MICROS / NANOS only — `Timestamp(Second)` is silently +/// promoted to `Timestamp(Millisecond)` by the data-node parquet round +/// trip (Arrow's `TimeUnit` enum is closed at four variants, so this +/// is the only precision that needs coercion). +fn coerce_unsupported_timestamp_precision( + schema: &arrow::datatypes::Schema, +) -> arrow::datatypes::Schema { + use arrow::datatypes::{DataType, Field, TimeUnit}; + let fields: Vec = schema + .fields() + .iter() + .map(|f| match f.data_type() { + DataType::Timestamp(TimeUnit::Second, tz) => Field::new( + f.name(), + DataType::Timestamp(TimeUnit::Millisecond, tz.clone()), + f.is_nullable(), + ) + .with_metadata(f.metadata().clone()), + _ => f.as_ref().clone(), + }) + .collect(); + arrow::datatypes::Schema::new_with_metadata(fields, schema.metadata().clone()) +} + +fn root_rel(root: &substrait::proto::PlanRel) -> Option { + match root.rel_type.as_ref()? { + substrait::proto::plan_rel::RelType::Rel(r) => Some(r.clone()), + substrait::proto::plan_rel::RelType::Root(rr) => rr.input.as_ref().cloned(), + } +} + +fn collect_reads(rel: &substrait::proto::Rel, out: &mut Vec) { + use substrait::proto::rel::RelType; + match rel.rel_type.as_ref() { + Some(RelType::Read(r)) => out.push((**r).clone()), + Some(RelType::Filter(f)) => { + if let Some(input) = &f.input { + collect_reads(input, out); + } + } + Some(RelType::Project(p)) => { + if let Some(input) = &p.input { + collect_reads(input, out); + } + } + Some(RelType::Aggregate(a)) => { + if let Some(input) = &a.input { + collect_reads(input, out); + } + } + Some(RelType::Sort(s)) => { + if let Some(input) = &s.input { + collect_reads(input, out); + } + } + Some(RelType::Fetch(f)) => { + if let Some(input) = &f.input { + collect_reads(input, out); + } + } + Some(RelType::Join(j)) => { + if let Some(left) = &j.left { + collect_reads(left, out); + } + if let Some(right) = &j.right { + collect_reads(right, out); + } + } + Some(RelType::Set(s)) => { + for input in &s.inputs { + collect_reads(input, out); + } + } + _ => {} + } +} + // --------------------------------------------------------------------------- // Coordinator-reduce local execution API // @@ -570,18 +748,15 @@ pub unsafe fn close_local_session(ptr: i64) { } } -/// Registers a streaming input on the session under `input_id`, using the -/// Arrow schema decoded from the IPC stream bytes. +/// Registers a streaming input on the session under `input_id`. The schema is +/// derived by lowering `partial_plan_bytes` (the producer side's substrait) to +/// a physical plan and reading its output schema — that is the schema the +/// producer will actually emit, so we eliminate any divergence between +/// declared and physical types. /// -/// The IPC bytes are expected to be a single schema message produced by -/// Arrow's streaming IPC writer (e.g. Java's `MessageSerializer.serializeMetadata` -/// or an `ArrowStreamWriter` flush of just the schema). Only the schema is -/// read — any payload in the buffer is ignored. -/// -/// Returns a heap-allocated pointer (as i64) to a [`PartitionStreamSender`]. -/// Caller must call `sender_close` exactly once to free it (closing the -/// sender signals EOF to the receiver side, so the native execute driver -/// naturally completes). +/// Returns `(sender_ptr, schema_ipc_bytes)`. The IPC bytes are written so the +/// Java tripwire (`typesMatch` in DatafusionReduceSink) can validate batches +/// against the same schema the native session is registered with. /// /// # Safety /// `session_ptr` must be a valid, non-zero pointer returned by @@ -589,19 +764,13 @@ pub unsafe fn close_local_session(ptr: i64) { pub unsafe fn register_partition_stream( session_ptr: i64, input_id: &str, - schema_ipc: &[u8], -) -> Result { + partial_plan_bytes: &[u8], +) -> Result<(i64, Vec), DataFusionError> { let session = &mut *(session_ptr as *mut LocalSession); - let mut cursor = Cursor::new(schema_ipc); - let reader = StreamReader::try_new(&mut cursor, None).map_err(|e| { - DataFusionError::Execution(format!( - "Failed to decode Arrow IPC schema for '{}': {}", - input_id, e - )) - })?; - let schema = reader.schema(); + let schema = derive_schema_from_partial_plan(partial_plan_bytes)?; + let schema_ipc = schema_to_ipc_bytes(schema.as_ref())?; let sender = session.register_partition(input_id, schema)?; - Ok(Box::into_raw(Box::new(sender)) as i64) + Ok((Box::into_raw(Box::new(sender)) as i64, schema_ipc)) } /// Executes a Substrait plan against a `LocalSession` and returns a @@ -710,6 +879,10 @@ pub unsafe fn sender_close(sender_ptr: i64) { /// Imports a batch of Arrow C Data structures into a [`Vec`] and /// registers them as an in-memory table on the given session under `input_id`. /// +/// The schema is derived by lowering `partial_plan_bytes` (the producer side's +/// substrait) the same way `register_partition_stream` does. Returns the +/// schema as IPC bytes so the Java side can validate fed batches against it. +/// /// The Java side has accumulated all shard responses, exported each /// `VectorSchemaRoot` to a paired `FFI_ArrowArray` / `FFI_ArrowSchema`, and /// passed the raw pointers as two parallel slices. Rust takes ownership of @@ -726,10 +899,10 @@ pub unsafe fn sender_close(sender_ptr: i64) { pub unsafe fn register_memtable( session_ptr: i64, input_id: &str, - schema_ipc: &[u8], + partial_plan_bytes: &[u8], array_ptrs: &[i64], schema_ptrs: &[i64], -) -> Result<(), DataFusionError> { +) -> Result, DataFusionError> { if array_ptrs.len() != schema_ptrs.len() { return Err(DataFusionError::Execution(format!( "register_memtable: array_ptrs.len()={} != schema_ptrs.len()={}", @@ -739,22 +912,13 @@ pub unsafe fn register_memtable( } let session = &mut *(session_ptr as *mut LocalSession); - let mut cursor = Cursor::new(schema_ipc); - let reader = StreamReader::try_new(&mut cursor, None).map_err(|e| { - DataFusionError::Execution(format!( - "Failed to decode Arrow IPC schema for '{}': {}", - input_id, e - )) - })?; - let table_schema = reader.schema(); - - // The IPC schema is what the substrait plan was compiled against — same as the streaming - // sink registers. The exported VSRs may arrive with batch-level schemas that differ in - // nullability/metadata/field-naming details; the streaming sink tolerates this because - // DataFusion's streaming source addresses columns by index. `MemTable::try_new` instead - // checks each batch's schema against the table schema. To stay compatible with both - // shapes, rebuild each imported batch with `table_schema` — the column data is reused - // verbatim, but the schema header is the planner's. + let table_schema = derive_schema_from_partial_plan(partial_plan_bytes)?; + let schema_ipc = schema_to_ipc_bytes(table_schema.as_ref())?; + + // Exported VSRs may arrive with batch-level schemas that differ in + // nullability/metadata/field-naming details; rebuild each imported batch + // with `table_schema` so MemTable::try_new sees uniform headers. Column + // data is reused verbatim. let mut batches = Vec::with_capacity(array_ptrs.len()); for (&array_ptr, &schema_ptr) in array_ptrs.iter().zip(schema_ptrs.iter()) { let ffi_array = FFI_ArrowArray::from_raw(array_ptr as *mut FFI_ArrowArray); @@ -774,5 +938,6 @@ pub unsafe fn register_memtable( batches.push(aligned); } - session.register_memtable(input_id, table_schema, batches) + session.register_memtable(input_id, table_schema, batches)?; + Ok(schema_ipc) } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 9460ca84a4d8f..2c88828002a4b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -223,9 +223,24 @@ pub unsafe extern "C" fn df_sql_to_substrait( str_from_raw(sql_ptr, sql_len).map_err(|e| format!("df_sql_to_substrait: sql: {}", e))?; let bytes = api::sql_to_substrait(shard_view_ptr, table_name, sql, runtime_ptr, &mgr) .map_err(|e| e.to_string())?; + write_out_buffer(&bytes, out_ptr, out_cap, out_len, "substrait plan")?; + Ok(0) +} + +/// Copies `bytes` into a caller-allocated `(out_ptr, out_cap)` buffer and writes +/// the byte count through `out_len` (when non-null). Returns `Err` when the +/// buffer is too small — the caller can re-allocate and retry. +unsafe fn write_out_buffer( + bytes: &[u8], + out_ptr: *mut u8, + out_cap: i64, + out_len: *mut i64, + label: &str, +) -> Result<(), String> { if bytes.len() > out_cap as usize { return Err(format!( - "substrait plan size {} exceeds buffer capacity {}", + "{} size {} exceeds buffer capacity {}", + label, bytes.len(), out_cap )); @@ -234,7 +249,7 @@ pub unsafe extern "C" fn df_sql_to_substrait( if !out_len.is_null() { *out_len = bytes.len() as i64; } - Ok(0) + Ok(()) } // --------------------------------------------------------------------------- @@ -273,19 +288,30 @@ pub unsafe extern "C" fn df_destroy_custom_cache_manager(ptr: i64) { } } +/// Registers a streaming partition input on the session. Schema is derived by +/// lowering the producer-side substrait `partial_plan_bytes`; the resulting +/// IPC-encoded schema is written into the caller-allocated `out_ptr/out_cap` +/// buffer with the byte count written through `out_len`. Returns the sender +/// pointer (negated error pointer on failure). #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn df_register_partition_stream( session_ptr: i64, input_id_ptr: *const u8, input_id_len: i64, - schema_ipc_ptr: *const u8, - schema_ipc_len: i64, + partial_plan_ptr: *const u8, + partial_plan_len: i64, + out_ptr: *mut u8, + out_cap: i64, + out_len: *mut i64, ) -> i64 { let input_id = str_from_raw(input_id_ptr, input_id_len) .map_err(|e| format!("df_register_partition_stream: input_id: {}", e))?; - let schema_ipc = slice::from_raw_parts(schema_ipc_ptr, schema_ipc_len as usize); - api::register_partition_stream(session_ptr, input_id, schema_ipc).map_err(|e| e.to_string()) + let partial_plan = slice::from_raw_parts(partial_plan_ptr, partial_plan_len as usize); + let (sender_ptr, schema_ipc) = + api::register_partition_stream(session_ptr, input_id, partial_plan).map_err(|e| e.to_string())?; + write_out_buffer(&schema_ipc, out_ptr, out_cap, out_len, "register_partition_stream schema IPC")?; + Ok(sender_ptr) } #[ffm_safe] @@ -352,15 +378,18 @@ pub unsafe extern "C" fn df_register_memtable( session_ptr: i64, input_id_ptr: *const u8, input_id_len: i64, - schema_ipc_ptr: *const u8, - schema_ipc_len: i64, + partial_plan_ptr: *const u8, + partial_plan_len: i64, array_ptrs: *const i64, schema_ptrs: *const i64, n_batches: i64, + out_ptr: *mut u8, + out_cap: i64, + out_len: *mut i64, ) -> i64 { let input_id = str_from_raw(input_id_ptr, input_id_len) .map_err(|e| format!("df_register_memtable: input_id: {}", e))?; - let schema_ipc = slice::from_raw_parts(schema_ipc_ptr, schema_ipc_len as usize); + let partial_plan = slice::from_raw_parts(partial_plan_ptr, partial_plan_len as usize); let n = n_batches as usize; let array_slice: &[i64] = if n == 0 { &[] @@ -372,9 +401,11 @@ pub unsafe extern "C" fn df_register_memtable( } else { slice::from_raw_parts(schema_ptrs, n) }; - api::register_memtable(session_ptr, input_id, schema_ipc, array_slice, schema_slice) - .map(|_| 0) - .map_err(|e| e.to_string()) + let schema_ipc = + api::register_memtable(session_ptr, input_id, partial_plan, array_slice, schema_slice) + .map_err(|e| e.to_string())?; + write_out_buffer(&schema_ipc, out_ptr, out_cap, out_len, "register_memtable schema IPC")?; + Ok(0) } #[ffm_safe] diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/convert_tz.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/convert_tz.rs index 6ae7c4199640d..343f5b8341359 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/convert_tz.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/convert_tz.rs @@ -41,8 +41,10 @@ use std::sync::Arc; use chrono::{DateTime, NaiveDateTime, Offset, TimeZone, Utc}; use chrono_tz::Tz; use datafusion::arrow::array::{ - Array, ArrayRef, StringArray, TimestampMillisecondArray, TimestampMillisecondBuilder, + Array, ArrayRef, TimestampMillisecondArray, TimestampMillisecondBuilder, }; + +use super::json_common::StringArrayView; use datafusion::arrow::datatypes::{DataType, TimeUnit}; use datafusion::common::{plan_err, ScalarValue}; use datafusion::error::{DataFusionError, Result}; @@ -132,19 +134,21 @@ impl ScalarUDFImpl for ConvertTzUdf { // Only materialize column-valued tz operands; for scalars the parsed // TzSpec is already in hand. Keep the ArrayRef alive alongside the - // downcast reference — StringArray borrows from the underlying buffer. + // view — StringArrayView borrows from the underlying buffer. let from_arr_ref: Option = if from_scalar.is_none() && matches!(&args.args[1], ColumnarValue::Array(_)) { - Some(materialize_string_array(&args.args[1], n, "from_tz")?) + Some(args.args[1].clone().into_array(n)?) } else { None }; let to_arr_ref: Option = if to_scalar.is_none() && matches!(&args.args[2], ColumnarValue::Array(_)) { - Some(materialize_string_array(&args.args[2], n, "to_tz")?) + Some(args.args[2].clone().into_array(n)?) } else { None }; - let from_array: Option<&StringArray> = from_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); - let to_array: Option<&StringArray> = to_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); + let from_array: Option> = + from_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; + let to_array: Option> = + to_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; let mut builder = TimestampMillisecondBuilder::with_capacity(n); for i in 0..n { @@ -152,9 +156,9 @@ impl ScalarUDFImpl for ConvertTzUdf { builder.append_null(); continue; } - let from = match (&from_scalar, from_array) { + let from = match (&from_scalar, from_array.as_ref().and_then(|a| a.cell(i))) { (Some(tz), _) => tz.clone(), - (None, Some(arr)) if !arr.is_null(i) => match parse_tz(arr.value(i)) { + (None, Some(s)) => match parse_tz(s) { Some(tz) => tz, None => { builder.append_null(); @@ -166,9 +170,9 @@ impl ScalarUDFImpl for ConvertTzUdf { continue; } }; - let to = match (&to_scalar, to_array) { + let to = match (&to_scalar, to_array.as_ref().and_then(|a| a.cell(i))) { (Some(tz), _) => tz.clone(), - (None, Some(arr)) if !arr.is_null(i) => match parse_tz(arr.value(i)) { + (None, Some(s)) => match parse_tz(s) { Some(tz) => tz, None => { builder.append_null(); @@ -203,18 +207,6 @@ fn scalar_tz(cv: &ColumnarValue) -> Option { None } -fn materialize_string_array(cv: &ColumnarValue, n: usize, label: &'static str) -> Result { - let arr = cv.clone().into_array(n)?; - if arr.as_any().downcast_ref::().is_none() { - return Err(DataFusionError::Internal(format!( - "convert_tz: {} expected Utf8, got {:?}", - label, - arr.data_type() - ))); - } - Ok(arr) -} - /// Parse timezone string (IANA name or `±HH:MM` offset). #[derive(Clone)] enum TzSpec { @@ -322,6 +314,7 @@ fn offset_seconds_at_instant( #[cfg(test)] mod tests { use super::*; + use datafusion::arrow::array::StringArray; // ±HH:MM offsets parse to the expected second counts. #[test] diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs index 8b3608c27fcb1..85bb3bdeca80c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_append.rs @@ -11,15 +11,13 @@ //! delegates to `JsonFunctions.jsonInsert` + `.meaningless_key` trick so Jayway //! routes to `Collection.add`). Non-array / missing targets are silent no-ops; //! any-NULL-arg / odd trailing arg / malformed-doc / malformed-path → NULL. -//! -//! Values always push as `Value::String` — every UDF arg is coerced to Utf8 -//! upstream, so nested `json_object` / `json_array` results arrive already -//! stringified and append as strings, matching legacy. +//! Values always push as `Value::String` — nested `json_object` / `json_array` +//! results arrive already stringified and append as strings, matching legacy. use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -29,7 +27,7 @@ use datafusion::logical_expr::{ }; use serde_json::Value; -use super::json_common::{as_utf8_array, parse, parse_ppl_segments, walk_mut, Segment}; +use super::json_common::{parse, parse_ppl_segments, scalar_utf8, walk_mut, Segment, StringArrayView}; use super::{coerce_slot, CoerceMode}; const NAME: &str = "json_append"; @@ -103,16 +101,16 @@ impl ScalarUDFImpl for JsonAppendUdf { .iter() .map(|v| v.clone().into_array(n)) .collect::>()?; - let columns: Vec<&datafusion::arrow::array::StringArray> = - arrays.iter().map(as_utf8_array).collect::>()?; + let columns: Vec> = + arrays.iter().map(StringArrayView::from_array).collect::>()?; let mut b = StringBuilder::with_capacity(n, n * 16); let mut rest: Vec> = Vec::with_capacity(columns.len() - 1); for i in 0..n { - let doc = cell(columns[0], i); + let doc = columns[0].cell(i); rest.clear(); for col in &columns[1..] { - rest.push(cell(col, i)); + rest.push(col.cell(i)); } match append(doc, &rest) { Some(s) => b.append_value(&s), @@ -123,23 +121,6 @@ impl ScalarUDFImpl for JsonAppendUdf { } } -fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { - match v { - ColumnarValue::Scalar( - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), - ) => s.as_deref(), - _ => None, - } -} - -fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } -} - /// Apply each (path, value) pair to a fresh parse of `doc`. Push-only: /// non-array targets (scalar, object) are silent no-ops, matching legacy /// `jsonInsert`'s Collection-parent branch skip. @@ -227,8 +208,8 @@ mod tests { fn nested_path_appends_to_inner_array() { // testJsonAppend case c — a pre-stringified JSON array is appended as // a single string element (legacy calls gson/jackson on the outer doc - // but NOT on the value; our Utf8-coerced arg arrives already - // stringified and is pushed as-is). + // but NOT on the value; the arg arrives already stringified and is + // pushed as-is). assert_eq!( append( Some(r#"{"school":{"teacher":["Alice"]}}"#), @@ -310,21 +291,6 @@ mod tests { assert!(append(Some(r#"{"a":[1]}"#), &[Some("a{"), Some("v")]).is_none()); } - #[test] - fn coerce_types_enforces_string_on_every_slot() { - let udf = JsonAppendUdf::new(); - assert_eq!( - udf.coerce_types(&[DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]) - .unwrap(), - vec![DataType::Utf8, DataType::Utf8, DataType::Utf8] - ); - let err = udf - .coerce_types(&[DataType::Utf8, DataType::Int32, DataType::Utf8]) - .unwrap_err() - .to_string(); - assert!(err.contains("expected string")); - } - #[test] fn return_type_is_utf8() { assert_eq!( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs index c8f60647d0ea0..a23ede7b70560 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_array_length.rs @@ -14,16 +14,17 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, Int32Builder, StringArray}; +use datafusion::arrow::array::{ArrayRef, Int32Builder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::{plan_err, ScalarValue}; -use datafusion::error::{DataFusionError, Result}; +use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; use serde_json::Value; +use super::json_common::StringArrayView; use super::{coerce_args, CoerceMode}; pub fn register_all(ctx: &SessionContext) { @@ -37,7 +38,6 @@ pub struct JsonArrayLengthUdf { impl JsonArrayLengthUdf { pub fn new() -> Self { - // user_defined + coerce_types lets DF cast LargeUtf8 / Utf8View → Utf8. Self { signature: Signature::user_defined(Volatility::Immutable), } @@ -100,20 +100,11 @@ impl ScalarUDFImpl for JsonArrayLengthUdf { } let arr = args.args[0].clone().into_array(n)?; - let strings = arr.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Internal(format!( - "json_array_length: expected Utf8, got {:?}", - arr.data_type() - )) - })?; + let strings = StringArrayView::from_array(&arr)?; let mut builder = Int32Builder::with_capacity(n); for i in 0..n { - if strings.is_null(i) { - builder.append_null(); - continue; - } - match json_array_len(strings.value(i)) { + match strings.cell(i).and_then(json_array_len) { Some(len) => builder.append_value(len), None => builder.append_null(), } @@ -137,7 +128,7 @@ fn json_array_len(s: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::Int32Array; + use datafusion::arrow::array::{Array, Int32Array, StringArray}; use datafusion::arrow::datatypes::Field; #[test] @@ -164,22 +155,6 @@ mod tests { assert_eq!(json_array_len(""), None); } - #[test] - fn coerce_types_accepts_string_variants() { - let udf = JsonArrayLengthUdf::new(); - for t in [DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View] { - let out = udf.coerce_types(std::slice::from_ref(&t)).unwrap(); - assert_eq!(out, vec![DataType::Utf8], "input {t:?} should coerce to Utf8"); - } - } - - #[test] - fn coerce_types_rejects_non_string() { - let udf = JsonArrayLengthUdf::new(); - let err = udf.coerce_types(&[DataType::Int64]).unwrap_err(); - assert!(err.to_string().contains("expected string")); - } - #[test] fn coerce_types_rejects_wrong_arity() { let udf = JsonArrayLengthUdf::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs index 80e77431d9c3e..c5ed75afa5c07 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_common.rs @@ -10,8 +10,11 @@ //! strings and typed segment vectors), a segment-based mutation walker used by //! the write UDFs, and a malformed-to-`None` JSON parser. -use datafusion::arrow::array::{ArrayRef, StringArray}; +use datafusion::arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::ScalarValue; use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::ColumnarValue; use serde_json::Value; /// Convert a PPL-style path (`a.b{0}.c{}`) to a JSONPath expression @@ -166,22 +169,59 @@ where /// Standard arity guard. pub(crate) fn check_arity(udf: &str, observed: usize, expected: usize) -> Result<()> { - (observed == expected) - .then_some(()) - .ok_or_else(|| plan_err_msg(format!("{udf} expects {expected} arguments, got {observed}"))) + (observed == expected).then_some(()).ok_or_else(|| { + DataFusionError::Plan(format!("{udf} expects {expected} arguments, got {observed}")) + }) } -fn plan_err_msg(msg: String) -> DataFusionError { - DataFusionError::Plan(msg) +/// View over an Arrow string array of any logical width (`Utf8`, `LargeUtf8`, +/// `Utf8View`). Dispatch happens once in [`Self::from_array`]; per-row access +/// via [`Self::cell`] is a small enum match. +pub(crate) enum StringArrayView<'a> { + Utf8(&'a StringArray), + LargeUtf8(&'a LargeStringArray), + Utf8View(&'a StringViewArray), } -/// Downcast an `ArrayRef` to `StringArray`. `coerce_types` with `CoerceMode::Utf8` -/// canonicalizes every string input to `Utf8` before this point, so a failure -/// indicates a planner bug rather than bad user input. -pub(crate) fn as_utf8_array(arr: &ArrayRef) -> Result<&StringArray> { - arr.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Internal(format!("expected Utf8, got {:?}", arr.data_type())) - }) +impl<'a> StringArrayView<'a> { + pub(crate) fn from_array(arr: &'a ArrayRef) -> Result { + match arr.data_type() { + DataType::Utf8 => Ok(Self::Utf8( + arr.as_any().downcast_ref::().expect("Utf8 downcast"), + )), + DataType::LargeUtf8 => Ok(Self::LargeUtf8( + arr.as_any().downcast_ref::().expect("LargeUtf8 downcast"), + )), + DataType::Utf8View => Ok(Self::Utf8View( + arr.as_any().downcast_ref::().expect("Utf8View downcast"), + )), + other => Err(DataFusionError::Internal(format!( + "expected string array (Utf8/LargeUtf8/Utf8View), got {other:?}" + ))), + } + } + + /// Returns `Some(value)` for non-null rows, `None` for nulls. + #[inline] + pub(crate) fn cell(&self, i: usize) -> Option<&str> { + match self { + Self::Utf8(a) => (!a.is_null(i)).then(|| a.value(i)), + Self::LargeUtf8(a) => (!a.is_null(i)).then(|| a.value(i)), + Self::Utf8View(a) => (!a.is_null(i)).then(|| a.value(i)), + } + } +} + +/// Extract `&str` from a string-typed `ColumnarValue::Scalar`. Returns `None` +/// for non-scalars, non-string scalars, and `NULL` scalars. Shared by every +/// json-family UDF's all-scalar fast path. +pub(crate) fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { + match v { + ColumnarValue::Scalar( + ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), + ) => s.as_deref(), + _ => None, + } } #[cfg(test)] diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs index ff189a7bcc4bf..d1f3b3d6224f9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_delete.rs @@ -16,7 +16,7 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -26,7 +26,7 @@ use datafusion::logical_expr::{ }; use serde_json::Value; -use super::json_common::{as_utf8_array, parse, parse_ppl_segments, walk_mut, Segment}; +use super::json_common::{parse, parse_ppl_segments, scalar_utf8, walk_mut, Segment, StringArrayView}; use super::{coerce_slot, CoerceMode}; const NAME: &str = "json_delete"; @@ -99,16 +99,16 @@ impl ScalarUDFImpl for JsonDeleteUdf { .iter() .map(|v| v.clone().into_array(n)) .collect::>()?; - let columns: Vec<&datafusion::arrow::array::StringArray> = - arrays.iter().map(as_utf8_array).collect::>()?; + let columns: Vec> = + arrays.iter().map(StringArrayView::from_array).collect::>()?; let mut b = StringBuilder::with_capacity(n, n * 16); let mut path_buf: Vec> = Vec::with_capacity(columns.len() - 1); for i in 0..n { - let doc = cell(columns[0], i); + let doc = columns[0].cell(i); path_buf.clear(); for col in &columns[1..] { - path_buf.push(cell(col, i)); + path_buf.push(col.cell(i)); } match delete(doc, &path_buf) { Some(s) => b.append_value(&s), @@ -119,23 +119,6 @@ impl ScalarUDFImpl for JsonDeleteUdf { } } -fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { - match v { - ColumnarValue::Scalar( - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), - ) => s.as_deref(), - _ => None, - } -} - -fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } -} - /// Apply every path's delete to a fresh parse of `doc`. Returns `None` for /// any-NULL arg / malformed doc / malformed path; otherwise the mutated /// document serialized back to a string. @@ -261,21 +244,6 @@ mod tests { ); } - #[test] - fn coerce_types_enforces_string_on_every_slot() { - let udf = JsonDeleteUdf::new(); - assert_eq!( - udf.coerce_types(&[DataType::LargeUtf8, DataType::Utf8View]) - .unwrap(), - vec![DataType::Utf8, DataType::Utf8] - ); - let err = udf - .coerce_types(&[DataType::Utf8, DataType::Int32]) - .unwrap_err() - .to_string(); - assert!(err.contains("expected string")); - } - #[test] fn return_type_is_utf8() { assert_eq!( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs index d1f61289a61ed..e84737f389b32 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extend.rs @@ -21,7 +21,7 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -31,7 +31,7 @@ use datafusion::logical_expr::{ }; use serde_json::Value; -use super::json_common::{as_utf8_array, parse, parse_ppl_segments, walk_mut, Segment}; +use super::json_common::{parse, parse_ppl_segments, scalar_utf8, walk_mut, Segment, StringArrayView}; use super::{coerce_slot, CoerceMode}; const NAME: &str = "json_extend"; @@ -101,16 +101,16 @@ impl ScalarUDFImpl for JsonExtendUdf { .iter() .map(|v| v.clone().into_array(n)) .collect::>()?; - let columns: Vec<&datafusion::arrow::array::StringArray> = - arrays.iter().map(as_utf8_array).collect::>()?; + let columns: Vec> = + arrays.iter().map(StringArrayView::from_array).collect::>()?; let mut b = StringBuilder::with_capacity(n, n * 16); let mut rest: Vec> = Vec::with_capacity(columns.len() - 1); for i in 0..n { - let doc = cell(columns[0], i); + let doc = columns[0].cell(i); rest.clear(); for col in &columns[1..] { - rest.push(cell(col, i)); + rest.push(col.cell(i)); } match extend(doc, &rest) { Some(s) => b.append_value(&s), @@ -121,23 +121,6 @@ impl ScalarUDFImpl for JsonExtendUdf { } } -fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { - match v { - ColumnarValue::Scalar( - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), - ) => s.as_deref(), - _ => None, - } -} - -fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } -} - /// Classify the raw value string into the push-list the terminal closure /// should apply. A successful JSON-array parse expands to the array /// elements; anything else (scalar, object, malformed JSON, plain string) @@ -323,21 +306,6 @@ mod tests { assert!(extend(Some(r#"{"a":[1]}"#), &[Some("a{"), Some("v")]).is_none()); } - #[test] - fn coerce_types_enforces_string_on_every_slot() { - let udf = JsonExtendUdf::new(); - assert_eq!( - udf.coerce_types(&[DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]) - .unwrap(), - vec![DataType::Utf8, DataType::Utf8, DataType::Utf8] - ); - let err = udf - .coerce_types(&[DataType::Utf8, DataType::Int32, DataType::Utf8]) - .unwrap_err() - .to_string(); - assert!(err.contains("expected string")); - } - #[test] fn return_type_is_utf8() { assert_eq!( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs index e046fc759fa98..ad4ad5f66fd88 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_extract.rs @@ -16,7 +16,7 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -27,7 +27,7 @@ use datafusion::logical_expr::{ use jsonpath_rust::{JsonPath, JsonPathValue}; use serde_json::Value; -use super::json_common::{as_utf8_array, convert_ppl_path, parse}; +use super::json_common::{convert_ppl_path, parse, scalar_utf8, StringArrayView}; use super::{coerce_slot, CoerceMode}; const NAME: &str = "json_extract"; @@ -69,7 +69,7 @@ impl ScalarUDFImpl for JsonExtractUdf { Ok(DataType::Utf8) } fn coerce_types(&self, args: &[DataType]) -> Result> { - // Homogeneous string variadic — every slot canonicalizes to Utf8. + // Homogeneous string variadic. args.iter() .enumerate() .map(|(i, ty)| coerce_slot(NAME, i, ty, CoerceMode::Utf8)) @@ -97,23 +97,22 @@ impl ScalarUDFImpl for JsonExtractUdf { return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(out))); } - // Columnar path: materialize each operand to a StringArray and walk - // row-by-row. This is the branch production traffic takes. + // Columnar path: walk row-by-row. Branch production traffic takes. let arrays: Vec = args .args .iter() .map(|v| v.clone().into_array(n)) .collect::>()?; - let columns: Vec<&datafusion::arrow::array::StringArray> = - arrays.iter().map(as_utf8_array).collect::>()?; + let columns: Vec> = + arrays.iter().map(StringArrayView::from_array).collect::>()?; let mut b = StringBuilder::with_capacity(n, n * 16); let mut path_buf: Vec> = Vec::with_capacity(columns.len() - 1); for i in 0..n { - let doc = cell(columns[0], i); + let doc = columns[0].cell(i); path_buf.clear(); for col in &columns[1..] { - path_buf.push(cell(col, i)); + path_buf.push(col.cell(i)); } match extract(doc, &path_buf) { Some(s) => b.append_value(&s), @@ -124,23 +123,6 @@ impl ScalarUDFImpl for JsonExtractUdf { } } -fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { - match v { - ColumnarValue::Scalar( - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), - ) => s.as_deref(), - _ => None, - } -} - -fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } -} - /// Core extraction. Returns `None` for the legacy NULL-producing cases /// (any-null arg, malformed doc, malformed path, no match, explicit-null match) /// and a `Some(String)` for every matched case. @@ -282,12 +264,12 @@ mod tests { } #[test] - fn coerce_types_enforces_string_on_every_slot() { + fn coerce_types_threads_each_slot_independently() { let udf = JsonExtractUdf::new(); assert_eq!( udf.coerce_types(&[DataType::LargeUtf8, DataType::Utf8View]) .unwrap(), - vec![DataType::Utf8, DataType::Utf8] + vec![DataType::LargeUtf8, DataType::Utf8View] ); let err = udf .coerce_types(&[DataType::Utf8, DataType::Int32]) diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs index 5971f45aff3c9..342e450d0c15d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_keys.rs @@ -13,7 +13,7 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -23,7 +23,7 @@ use datafusion::logical_expr::{ }; use serde_json::Value; -use super::json_common::{as_utf8_array, check_arity, parse}; +use super::json_common::{check_arity, parse, StringArrayView}; use super::{coerce_args, CoerceMode}; const NAME: &str = "json_keys"; @@ -81,14 +81,10 @@ impl ScalarUDFImpl for JsonKeysUdf { } let arr = args.args[0].clone().into_array(n)?; - let strings = as_utf8_array(&arr)?; + let strings = StringArrayView::from_array(&arr)?; let mut b = StringBuilder::with_capacity(n, n * 16); for i in 0..n { - if strings.is_null(i) { - b.append_null(); - continue; - } - match json_keys(strings.value(i)) { + match strings.cell(i).and_then(json_keys) { Some(s) => b.append_value(&s), None => b.append_null(), } @@ -146,10 +142,11 @@ mod tests { } #[test] - fn coerce_types_enforces_string_arity() { + fn coerce_types_rejects_wrong_arity() { + // `coerce_args` covers the per-variant accept/reject contract centrally; + // here we only need the UDF-specific arity guard. let udf = JsonKeysUdf::new(); - assert_eq!(udf.coerce_types(&[DataType::LargeUtf8]).unwrap(), vec![DataType::Utf8]); - assert!(udf.coerce_types(&[DataType::Int64]).unwrap_err().to_string().contains("expected string")); assert!(udf.coerce_types(&[]).is_err()); + assert!(udf.coerce_types(&[DataType::Utf8, DataType::Utf8]).is_err()); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs index daf33e8d2f6bb..32e997730c2fa 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/json_set.rs @@ -10,15 +10,13 @@ //! each path match (parity with legacy `JsonSetFunctionImpl` → Jayway //! `ctx.set` guarded by `ctx.read != null`: *replace-only*, never inserts). //! Missing paths are no-ops; any-NULL-arg / odd trailing arg / malformed-doc / -//! malformed-path → NULL. -//! -//! Values always store as JSON strings because every UDF arg is coerced to -//! Utf8 upstream — matching the legacy fixture `"b":"3"` (not `"b":3`). +//! malformed-path → NULL. Values always store as JSON strings, matching the +//! legacy fixture `"b":"3"` (not `"b":3`). use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::ScalarValue; use datafusion::error::Result; @@ -28,7 +26,7 @@ use datafusion::logical_expr::{ }; use serde_json::Value; -use super::json_common::{as_utf8_array, parse, parse_ppl_segments, walk_mut, Segment}; +use super::json_common::{parse, parse_ppl_segments, scalar_utf8, walk_mut, Segment, StringArrayView}; use super::{coerce_slot, CoerceMode}; const NAME: &str = "json_set"; @@ -102,16 +100,16 @@ impl ScalarUDFImpl for JsonSetUdf { .iter() .map(|v| v.clone().into_array(n)) .collect::>()?; - let columns: Vec<&datafusion::arrow::array::StringArray> = - arrays.iter().map(as_utf8_array).collect::>()?; + let columns: Vec> = + arrays.iter().map(StringArrayView::from_array).collect::>()?; let mut b = StringBuilder::with_capacity(n, n * 16); let mut rest: Vec> = Vec::with_capacity(columns.len() - 1); for i in 0..n { - let doc = cell(columns[0], i); + let doc = columns[0].cell(i); rest.clear(); for col in &columns[1..] { - rest.push(cell(col, i)); + rest.push(col.cell(i)); } match set(doc, &rest) { Some(s) => b.append_value(&s), @@ -122,23 +120,6 @@ impl ScalarUDFImpl for JsonSetUdf { } } -fn scalar_utf8(v: &ColumnarValue) -> Option<&str> { - match v { - ColumnarValue::Scalar( - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s), - ) => s.as_deref(), - _ => None, - } -} - -fn cell(arr: &datafusion::arrow::array::StringArray, i: usize) -> Option<&str> { - if arr.is_null(i) { - None - } else { - Some(arr.value(i)) - } -} - /// Apply each (path, value) pair to a fresh parse of `doc`. Replace-only: /// missing paths are no-ops, matching legacy `jsonSet`'s `ctx.read != null` /// guard. @@ -256,21 +237,6 @@ mod tests { assert!(set(Some(r#"{"a":1}"#), &[Some("a{"), Some("v")]).is_none()); } - #[test] - fn coerce_types_enforces_string_on_every_slot() { - let udf = JsonSetUdf::new(); - assert_eq!( - udf.coerce_types(&[DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View]) - .unwrap(), - vec![DataType::Utf8, DataType::Utf8, DataType::Utf8] - ); - let err = udf - .coerce_types(&[DataType::Utf8, DataType::Int32, DataType::Utf8]) - .unwrap_err() - .to_string(); - assert!(err.contains("expected string")); - } - #[test] fn return_type_is_utf8() { assert_eq!( 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 901432055b2ed..2749adb0aeafe 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mod.rs @@ -77,21 +77,23 @@ pub(crate) fn coerce_slot( ), }, CoerceMode::Int64 => match observed { - Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | Float32 | Float64 - | Decimal128(_, _) | Decimal256(_, _) => Ok(Int64), + Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | Float32 | Float64 => { + Ok(Int64) + } other => { plan_err!("{udf_name}: arg {slot_index} expected integer or float, got {other:?}") } }, CoerceMode::Float64 => match observed { - Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | Float32 | Float64 - | Decimal128(_, _) | Decimal256(_, _) => Ok(Float64), + Int8 | Int16 | Int32 | Int64 | UInt8 | UInt16 | UInt32 | UInt64 | Float32 | Float64 => { + Ok(Float64) + } other => { plan_err!("{udf_name}: arg {slot_index} expected integer or float, got {other:?}") } }, CoerceMode::Utf8 => match observed { - Utf8 | LargeUtf8 | Utf8View => Ok(Utf8), + Utf8 | LargeUtf8 | Utf8View => Ok(observed.clone()), other => plan_err!("{udf_name}: arg {slot_index} expected string, got {other:?}"), }, } @@ -283,17 +285,6 @@ mod tests { assert!(err.to_string().contains("expected integer or float")); } - #[test] - fn int64_accepts_decimal_types() { - // PPL emits Decimal128(p,s) literals (e.g. `span=2.5` becomes - // Decimal128(2, 1)). The Int64 coerce-mode must accept and canonicalize. - for observed in [DataType::Decimal128(2, 1), DataType::Decimal256(10, 3)] { - let result = coerce_slot("i", 0, &observed, CoerceMode::Int64).unwrap(); - assert_eq!(result, DataType::Int64); - } - } - - // ── Float64 ──────────────────────────────────────────────────────────── #[test] fn float64_accepts_every_number() { @@ -315,22 +306,12 @@ mod tests { assert!(err.to_string().contains("expected integer or float")); } - #[test] - fn float64_accepts_decimal_types() { - // Decimal128 flows in for fractional literals like `span=2.5`. - for observed in [DataType::Decimal128(2, 1), DataType::Decimal256(10, 3)] { - let result = coerce_slot("f", 0, &observed, CoerceMode::Float64).unwrap(); - assert_eq!(result, DataType::Float64); - } - } - - // ── Utf8 ─────────────────────────────────────────────────────────────── #[test] - fn utf8_accepts_every_string_variant() { + fn utf8_passes_string_variant_through_unchanged() { for observed in [DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View] { let result = coerce_slot("s", 0, &observed, CoerceMode::Utf8).unwrap(); - assert_eq!(result, DataType::Utf8); + assert_eq!(result, observed, "CoerceMode::Utf8 should pass the variant through"); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs index 767f949c8c4f2..647f1f1d830f3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/mvfind.rs @@ -33,10 +33,12 @@ 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, + Array, ArrayRef, BooleanArray, Float32Array, Float64Array, GenericListArray, Int16Array, + Int32Array, Int32Builder, Int64Array, Int8Array, UInt16Array, UInt32Array, UInt64Array, + UInt8Array, }; + +use super::json_common::StringArrayView; use datafusion::arrow::datatypes::DataType; use datafusion::common::{plan_err, ScalarValue}; use datafusion::error::{DataFusionError, Result}; @@ -142,9 +144,8 @@ impl ScalarUDFImpl for MvfindUdf { } else { None }; - let pattern_arr: Option<&StringArray> = pattern_arr_ref - .as_ref() - .and_then(|a| a.as_any().downcast_ref::()); + let pattern_arr: Option> = + pattern_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; let mut builder = Int32Builder::with_capacity(n); for i in 0..n { @@ -153,9 +154,9 @@ impl ScalarUDFImpl for MvfindUdf { continue; } // Per-row regex (compile if column-valued; reuse the scalar compile otherwise). - let regex_for_row: Option = match (&scalar_regex, pattern_arr) { + let regex_for_row: Option = match (&scalar_regex, pattern_arr.as_ref().and_then(|a| a.cell(i))) { (Some(r), _) => Some(r.clone()), - (None, Some(arr)) if !arr.is_null(i) => Regex::new(arr.value(i)).ok(), + (None, Some(s)) => Regex::new(s).ok(), _ => None, }; let regex = match regex_for_row { @@ -195,30 +196,32 @@ fn find_first_match(arr: &dyn Array, regex: &Regex) -> Option { }}; } 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); - } + DataType::Utf8 => { + let typed = arr.as_any().downcast_ref::()?; + for i in 0..n { + if !typed.is_null(i) && 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::LargeUtf8 => { + let typed = arr.as_any().downcast_ref::()?; + for i in 0..n { + if !typed.is_null(i) && regex.is_match(typed.value(i)) { + return Some(i as i32); } - None } + None + } + DataType::Utf8View => { + let typed = arr.as_any().downcast_ref::()?; + for i in 0..n { + if !typed.is_null(i) && regex.is_match(typed.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()), diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract.rs index 131cb2ad18712..704a4186f5af3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract.rs @@ -41,7 +41,7 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringArray, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; use datafusion::arrow::datatypes::DataType; use datafusion::common::{plan_err, ScalarValue}; use datafusion::error::{DataFusionError, Result}; @@ -51,6 +51,7 @@ use datafusion::logical_expr::{ }; use regex::Regex; +use super::json_common::StringArrayView; use super::{coerce_args, CoerceMode}; pub fn register_all(ctx: &SessionContext) { @@ -120,19 +121,11 @@ impl ScalarUDFImpl for RexExtractUdf { None => None, }; - let input = args.args[0].clone().into_array(n)?; - let input = input - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "rex_extract: expected Utf8 input, got {:?}", - input.data_type() - )) - })?; + let input_arr = args.args[0].clone().into_array(n)?; + let input = StringArrayView::from_array(&input_arr)?; // Materialize column-valued operands lazily; keep the ArrayRef alive - // alongside any downcast reference (the StringArray borrows its buffer). + // alongside the StringArrayView (the view borrows the underlying buffers). let pattern_arr_ref: Option = if pattern_scalar.is_none() && matches!(&args.args[1], ColumnarValue::Array(_)) { Some(args.args[1].clone().into_array(n)?) } else { @@ -143,22 +136,24 @@ impl ScalarUDFImpl for RexExtractUdf { } else { None }; - let pattern_array: Option<&StringArray> = pattern_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); - let group_array: Option<&StringArray> = group_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); + let pattern_array: Option> = + pattern_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; + let group_array: Option> = + group_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; let mut builder = StringBuilder::with_capacity(n, n * 16); for i in 0..n { - if input.is_null(i) { + let Some(input_value) = input.cell(i) else { builder.append_null(); continue; - } + }; // Resolve the regex — scalar fast-path or per-row column lookup. let regex_owned; - let regex: &Regex = match (&scalar_regex, pattern_array) { + let regex: &Regex = match (&scalar_regex, pattern_array.as_ref().and_then(|a| a.cell(i))) { (Some(r), _) => r, - (None, Some(arr)) if !arr.is_null(i) => { - regex_owned = compile_pattern(arr.value(i))?; + (None, Some(s)) => { + regex_owned = compile_pattern(s)?; ®ex_owned } _ => { @@ -168,16 +163,16 @@ impl ScalarUDFImpl for RexExtractUdf { }; // Resolve the group name (or numeric index, if it parses as one). - let group_name: &str = match (&group_scalar, group_array) { + let group_name: &str = match (&group_scalar, group_array.as_ref().and_then(|a| a.cell(i))) { (Some(g), _) => g.as_str(), - (None, Some(arr)) if !arr.is_null(i) => arr.value(i), + (None, Some(s)) => s, _ => { builder.append_null(); continue; } }; - match extract_group(regex, input.value(i), group_name) { + match extract_group(regex, input_value, group_name) { Some(s) => builder.append_value(s), None => builder.append_null(), } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract_multi.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract_multi.rs index ee1011dec68e8..0f7bd3422d53c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract_multi.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_extract_multi.rs @@ -26,12 +26,12 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{ - Array, ArrayRef, Int32Array, ListBuilder, StringArray, StringBuilder, -}; +use datafusion::arrow::array::{Array, ArrayRef, Int32Array, ListBuilder, StringBuilder}; + +use super::json_common::StringArrayView; use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::common::{plan_err, ScalarValue}; -use datafusion::error::{DataFusionError, Result}; +use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, @@ -112,16 +112,8 @@ impl ScalarUDFImpl for RexExtractMultiUdf { None => None, }; - let input = args.args[0].clone().into_array(n)?; - let input = input - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "rex_extract_multi: expected Utf8 input, got {:?}", - input.data_type() - )) - })?; + let input_arr = args.args[0].clone().into_array(n)?; + let input = StringArrayView::from_array(&input_arr)?; let pattern_arr_ref: Option = if pattern_scalar.is_none() && matches!(&args.args[1], ColumnarValue::Array(_)) { Some(args.args[1].clone().into_array(n)?) @@ -138,8 +130,10 @@ impl ScalarUDFImpl for RexExtractMultiUdf { } else { None }; - let pattern_array: Option<&StringArray> = pattern_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); - let group_array: Option<&StringArray> = group_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); + let pattern_array: Option> = + pattern_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; + let group_array: Option> = + group_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; // After coerce_types(Int64) the array may arrive as Int32 in some plans; // accept either by widening on read. let max_match_array_i32: Option<&Int32Array> = max_match_arr_ref @@ -152,16 +146,16 @@ impl ScalarUDFImpl for RexExtractMultiUdf { ))); for i in 0..n { - if input.is_null(i) { + let Some(input_value) = input.cell(i) else { builder.append_null(); continue; - } + }; let regex_owned; - let regex: &Regex = match (&scalar_regex, pattern_array) { + let regex: &Regex = match (&scalar_regex, pattern_array.as_ref().and_then(|a| a.cell(i))) { (Some(r), _) => r, - (None, Some(arr)) if !arr.is_null(i) => { - regex_owned = compile_pattern(arr.value(i))?; + (None, Some(s)) => { + regex_owned = compile_pattern(s)?; ®ex_owned } _ => { @@ -170,9 +164,9 @@ impl ScalarUDFImpl for RexExtractMultiUdf { } }; - let group_name: &str = match (&group_scalar, group_array) { + let group_name: &str = match (&group_scalar, group_array.as_ref().and_then(|a| a.cell(i))) { (Some(g), _) => g.as_str(), - (None, Some(arr)) if !arr.is_null(i) => arr.value(i), + (None, Some(s)) => s, _ => { builder.append_null(); continue; @@ -188,7 +182,7 @@ impl ScalarUDFImpl for RexExtractMultiUdf { } }; - let matches = collect_matches(regex, input.value(i), group_name, max_match); + let matches = collect_matches(regex, input_value, group_name, max_match); if matches.is_empty() { builder.append_null(); } else { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_offset.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_offset.rs index 1cb0a636469a2..6d521fcdfd851 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_offset.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/rex_offset.rs @@ -33,10 +33,12 @@ use std::any::Any; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, StringArray, StringBuilder}; +use datafusion::arrow::array::{ArrayRef, StringBuilder}; + +use super::json_common::StringArrayView; use datafusion::arrow::datatypes::DataType; use datafusion::common::{plan_err, ScalarValue}; -use datafusion::error::{DataFusionError, Result}; +use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, @@ -107,35 +109,28 @@ impl ScalarUDFImpl for RexOffsetUdf { None => None, }; - let input = args.args[0].clone().into_array(n)?; - let input = input - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Internal(format!( - "rex_offset: expected Utf8 input, got {:?}", - input.data_type() - )) - })?; + let input_arr = args.args[0].clone().into_array(n)?; + let input = StringArrayView::from_array(&input_arr)?; let pattern_arr_ref: Option = if pattern_scalar.is_none() && matches!(&args.args[1], ColumnarValue::Array(_)) { Some(args.args[1].clone().into_array(n)?) } else { None }; - let pattern_array: Option<&StringArray> = pattern_arr_ref.as_ref().and_then(|a| a.as_any().downcast_ref::()); + let pattern_array: Option> = + pattern_arr_ref.as_ref().map(StringArrayView::from_array).transpose()?; let mut builder = StringBuilder::with_capacity(n, n * 32); for i in 0..n { - if input.is_null(i) { + let Some(input_value) = input.cell(i) else { builder.append_null(); continue; - } + }; let regex_owned; - let regex: &Regex = match (&scalar_regex, pattern_array) { + let regex: &Regex = match (&scalar_regex, pattern_array.as_ref().and_then(|a| a.cell(i))) { (Some(r), _) => r, - (None, Some(arr)) if !arr.is_null(i) => { - regex_owned = compile_pattern(arr.value(i))?; + (None, Some(s)) => { + regex_owned = compile_pattern(s)?; ®ex_owned } _ => { @@ -144,7 +139,7 @@ impl ScalarUDFImpl for RexOffsetUdf { } }; - match calculate_offsets(regex, input.value(i)) { + match calculate_offsets(regex, input_value) { Some(s) => builder.append_value(s), None => builder.append_null(), } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/sha1.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/sha1.rs index a300a35b3d5c4..c37d181402ebb 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/sha1.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/sha1.rs @@ -17,7 +17,7 @@ use std::any::Any; use std::hash::{Hash, Hasher}; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, AsArray, StringArray}; +use datafusion::arrow::array::{Array, ArrayRef, StringArray}; use datafusion::arrow::datatypes::DataType; use datafusion::common::{exec_err, Result, ScalarValue}; use datafusion::execution::context::SessionContext; @@ -43,6 +43,7 @@ impl Sha1Udf { vec![ TypeSignature::Exact(vec![DataType::Utf8]), TypeSignature::Exact(vec![DataType::LargeUtf8]), + TypeSignature::Exact(vec![DataType::Utf8View]), ], Volatility::Immutable, ), @@ -92,27 +93,17 @@ impl ScalarUDFImpl for Sha1Udf { return exec_err!("sha1 expects exactly 1 argument, got {}", args.args.len()); } match &args.args[0] { - ColumnarValue::Scalar(ScalarValue::Utf8(opt)) - | ColumnarValue::Scalar(ScalarValue::LargeUtf8(opt)) => Ok(ColumnarValue::Scalar( - ScalarValue::Utf8(opt.as_ref().map(|s| hash_hex(s.as_bytes()))), - )), + ColumnarValue::Scalar( + ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) | ScalarValue::Utf8View(opt), + ) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8( + opt.as_ref().map(|s| hash_hex(s.as_bytes())), + ))), ColumnarValue::Scalar(other) => exec_err!("sha1: expected Utf8 input, got {other:?}"), ColumnarValue::Array(arr) => { - let out: StringArray = match arr.data_type() { - DataType::Utf8 => arr - .as_string::() - .iter() - .map(|opt| opt.map(|s| hash_hex(s.as_bytes()))) - .collect(), - DataType::LargeUtf8 => arr - .as_string::() - .iter() - .map(|opt| opt.map(|s| hash_hex(s.as_bytes()))) - .collect(), - other => { - return exec_err!("sha1: expected Utf8 array, got {other:?}"); - } - }; + let view = super::json_common::StringArrayView::from_array(arr)?; + let out: StringArray = (0..arr.len()) + .map(|i| view.cell(i).map(|s| hash_hex(s.as_bytes()))) + .collect(); Ok(ColumnarValue::Array(Arc::new(out) as ArrayRef)) } } @@ -135,7 +126,7 @@ fn hash_hex(value: &[u8]) -> String { #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::{Array, StringArray}; + use datafusion::arrow::array::{Array, AsArray, StringArray}; use datafusion::arrow::datatypes::Field; fn udf() -> Sha1Udf { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/strftime.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/strftime.rs index 89f667dbac2b9..7bd4f50972848 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/strftime.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/strftime.rs @@ -82,7 +82,7 @@ impl ScalarUDFImpl for StrftimeUdf { other => return plan_err!("strftime: arg 0 expected numeric/timestamp/date/string, got {other:?}"), }; let format = match &arg_types[1] { - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => DataType::Utf8, + t @ (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View) => t.clone(), other => return plan_err!("strftime: arg 1 expected string, got {other:?}"), }; Ok(vec![value, format]) @@ -150,24 +150,17 @@ impl ScalarUDFImpl for StrftimeUdf { fn scalar_utf8(s: &ScalarValue) -> Result> { match s { - ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) => Ok(opt.clone()), + ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) | ScalarValue::Utf8View(opt) => { + Ok(opt.clone()) + } other => exec_err!("strftime: format must be VARCHAR, got {other:?}"), } } fn format_at(array: &ArrayRef, row: usize) -> Result> { - let (is_null, value) = match array.data_type() { - DataType::Utf8 => { - let a = array.as_string::(); - (a.is_null(row), a.value(row).to_string()) - } - DataType::LargeUtf8 => { - let a = array.as_string::(); - (a.is_null(row), a.value(row).to_string()) - } - other => return exec_err!("strftime: expected string format array, got {other:?}"), - }; - Ok(if is_null { None } else { Some(value) }) + let view = super::json_common::StringArrayView::from_array(array) + .map_err(|e| datafusion::common::DataFusionError::Execution(format!("strftime: {e}")))?; + Ok(view.cell(row).map(|s| s.to_string())) } fn render_from_seconds(value: f64, format: Option<&str>) -> Option { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tonumber.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tonumber.rs index 765056a176774..518cad35bbb12 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tonumber.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tonumber.rs @@ -17,10 +17,11 @@ use std::any::Any; use std::hash::{Hash, Hasher}; use std::sync::Arc; -use datafusion::arrow::array::{Array, ArrayRef, Float64Array, Float64Builder, StringArray}; +use datafusion::arrow::array::{ArrayRef, Float64Array, Float64Builder}; + +use super::json_common::StringArrayView; use datafusion::arrow::datatypes::DataType; use datafusion::common::{exec_err, Result, ScalarValue}; -use datafusion::error::DataFusionError; use datafusion::execution::context::SessionContext; use datafusion::logical_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TypeSignature, @@ -84,7 +85,7 @@ enum BaseMode { /// How the `value` argument is supplied enum ValueSource<'a> { Scalar(Option<&'a str>), - Array(&'a StringArray), + Array(StringArrayView<'a>), } impl<'a> ValueSource<'a> { @@ -92,8 +93,7 @@ impl<'a> ValueSource<'a> { fn at(&self, i: usize) -> Option<&str> { match self { ValueSource::Scalar(s) => *s, - ValueSource::Array(arr) if arr.is_null(i) => None, - ValueSource::Array(arr) => Some(arr.value(i)), + ValueSource::Array(view) => view.cell(i), } } } @@ -151,18 +151,10 @@ impl ScalarUDFImpl for ToNumberUdf { _ => None, }; let values: ValueSource = match (&values_arr_ref, value_col) { - (Some(arr), _) => { - let sa = arr.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Internal(format!( - "tonumber: value expected Utf8, got {:?}", - arr.data_type() - )) - })?; - ValueSource::Array(sa) - } - (None, ColumnarValue::Scalar(ScalarValue::Utf8(opt))) => { - ValueSource::Scalar(opt.as_deref()) - } + (Some(arr), _) => ValueSource::Array(StringArrayView::from_array(arr)?), + (None, ColumnarValue::Scalar( + ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) | ScalarValue::Utf8View(opt), + )) => ValueSource::Scalar(opt.as_deref()), (None, other) => { return exec_err!("tonumber: value expected Utf8, got {other:?}"); } @@ -221,7 +213,7 @@ fn parse_with_base(s: Option<&str>, base: Option) -> Option { #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::{AsArray, Int32Array}; + use datafusion::arrow::array::{Array, AsArray, Int32Array, StringArray}; use datafusion::arrow::datatypes::Field; fn invoke_scalar(value: Option<&str>, base: Option) -> Option { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tostring.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tostring.rs index e9593e1a0d79f..b1b14a00a2138 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tostring.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/udf/tostring.rs @@ -46,14 +46,19 @@ pub struct ToStringUdf { impl ToStringUdf { pub fn new() -> Self { + // The format arg accepts any string variant — the body dispatches via + // StringArrayView so the producer's Utf8View doesn't get cast to Utf8 + // at the planner level. + let exacts = [DataType::Int64, DataType::Float64] + .into_iter() + .flat_map(|value_ty| { + [DataType::Utf8, DataType::LargeUtf8, DataType::Utf8View] + .into_iter() + .map(move |fmt_ty| TypeSignature::Exact(vec![value_ty.clone(), fmt_ty])) + }) + .collect(); Self { - signature: Signature::one_of( - vec![ - TypeSignature::Exact(vec![DataType::Int64, DataType::Utf8]), - TypeSignature::Exact(vec![DataType::Float64, DataType::Utf8]), - ], - Volatility::Immutable, - ), + signature: Signature::one_of(exacts, Volatility::Immutable), } } } @@ -182,31 +187,17 @@ fn scalar_to_str( /// format cells. fn format_at(format_col: &ColumnarValue, row: usize) -> Result> { match format_col { - ColumnarValue::Scalar(ScalarValue::Utf8(opt)) | ColumnarValue::Scalar(ScalarValue::LargeUtf8(opt)) => { - Ok(opt.clone()) - } + ColumnarValue::Scalar( + ScalarValue::Utf8(opt) | ScalarValue::LargeUtf8(opt) | ScalarValue::Utf8View(opt), + ) => Ok(opt.clone()), ColumnarValue::Scalar(other) => { exec_err!("tostring: format must be VARCHAR, got {other:?}") } - ColumnarValue::Array(arr) => match arr.data_type() { - DataType::Utf8 => { - let strs = arr.as_string::(); - if strs.is_null(row) { - Ok(None) - } else { - Ok(Some(strs.value(row).to_string())) - } - } - DataType::LargeUtf8 => { - let strs = arr.as_string::(); - if strs.is_null(row) { - Ok(None) - } else { - Ok(Some(strs.value(row).to_string())) - } - } - other => exec_err!("tostring: expected Utf8 format array, got {other:?}"), - }, + ColumnarValue::Array(arr) => { + let view = super::json_common::StringArrayView::from_array(arr) + .map_err(|e| datafusion::common::DataFusionError::Execution(format!("tostring: {e}")))?; + Ok(view.cell(row).map(|s| s.to_string())) + } } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs b/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs index 04dd1d8694028..741c675242ef9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/tests/local_exec_test.rs @@ -29,7 +29,6 @@ use std::sync::{Arc, OnceLock}; use std::thread; use std::time::Duration; -use arrow::ipc::writer::StreamWriter; use arrow_array::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; use arrow_array::{Array, Int64Array, RecordBatch, StructArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; @@ -119,17 +118,36 @@ fn i64_batch(schema: &SchemaRef, values: &[i64]) -> RecordBatch { .expect("batch builds") } -/// Serialize a `Schema` to Arrow IPC stream bytes — the shape -/// `df_register_partition_stream` expects. -fn schema_to_ipc_bytes(schema: &Schema) -> Vec { - let mut buf: Vec = Vec::new(); - { - let mut writer = StreamWriter::try_new(&mut buf, schema).expect("stream writer builds"); - writer.finish().expect("stream writer finishes"); - } +/// Build a substrait `SELECT * FROM ` plan whose lowered output schema +/// matches `schema`. `df_register_partition_stream` derives the registered +/// table's schema by lowering this plan, so the inferred schema is identical +/// to the producer-side schema. +async fn build_seed_substrait(table: &str, schema: SchemaRef) -> Vec { + let ctx = SessionContext::new(); + let empty = MemTable::try_new(Arc::clone(&schema), vec![vec![]]).expect("mem table"); + ctx.register_table(table, Arc::new(empty)) + .expect("register seed table"); + let plan = ctx + .sql(&format!("SELECT * FROM \"{}\"", table)) + .await + .expect("sql parses") + .logical_plan() + .clone(); + let substrait = to_substrait_plan(&plan, &ctx.state()).expect("to_substrait"); + let mut buf = Vec::new(); + substrait.encode(&mut buf).expect("encode"); buf } +/// Decode an IPC schema-only stream back to an arrow Schema. +fn ipc_bytes_to_schema(bytes: &[u8]) -> Schema { + use arrow::ipc::reader::StreamReader; + use std::io::Cursor; + let cursor = Cursor::new(bytes); + let reader = StreamReader::try_new(cursor, None).expect("stream reader builds"); + (*reader.schema()).clone() +} + /// Export a `RecordBatch` as (FFI_ArrowArray*, FFI_ArrowSchema*) and transfer /// ownership to the caller — mirrors what `DatafusionReduceSink.feed` will do /// on the Java side. @@ -184,18 +202,34 @@ async fn build_sum_substrait(schema: SchemaRef) -> Vec { } fn register_input(session_ptr: i64, input_id: &str, schema: &Schema) -> i64 { - let ipc = schema_to_ipc_bytes(schema); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build tokio rt"); + let plan_bytes = rt.block_on(build_seed_substrait(input_id, Arc::new(schema.clone()))); let id_bytes = input_id.as_bytes(); + let mut out_buf = vec![0u8; 64 * 1024]; + let mut out_len: i64 = 0; let rc = unsafe { df_register_partition_stream( session_ptr, id_bytes.as_ptr(), id_bytes.len() as i64, - ipc.as_ptr(), - ipc.len() as i64, + plan_bytes.as_ptr(), + plan_bytes.len() as i64, + out_buf.as_mut_ptr(), + out_buf.len() as i64, + &mut out_len as *mut i64, ) }; assert!(rc > 0, "df_register_partition_stream rc={}", rc); + let derived = ipc_bytes_to_schema(&out_buf[..out_len as usize]); + // Sanity: derived schema matches the caller-provided one (column types). + assert_eq!( + derived.fields().len(), + schema.fields().len(), + "derived schema field count mismatch" + ); rc } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java index 0be42e9690c43..134d3a8f9ddea 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java @@ -43,8 +43,11 @@ * * *

    Multi-input shapes (Union, future Join) are supported at this base by exposing - * {@link #childInputs} (childStageId → schemaIpc) for subclasses to register one - * native partition per child stage. The {@link #INPUT_ID} constant remains as the + * {@link #childInputs} (childStageId → producer-side plan bytes) for subclasses to register + * one native partition per child stage. The native call returns the IPC-encoded schema the + * session settled on after lowering; subclasses populate {@link #childSchemas} from those + * returns so the {@code typesMatch} tripwire validates batches against the same schema the + * native session is registered with. The {@link #INPUT_ID} constant remains as the * conventional name for the single-input case (childStageId=0); the per-child id is * computed via {@link #inputIdFor(int)}. * @@ -75,16 +78,18 @@ abstract class AbstractDatafusionReduceSink implements ExchangeSink { */ protected final DataFusionReduceState preparedState; /** - * Per-child Arrow schema IPC bytes, keyed by childStageId. Iteration order matches + * Per-child producer-side plan bytes, keyed by childStageId. Iteration order matches * the order of {@code ctx.childInputs()} so subclasses get deterministic registration. + * Subclasses pass each entry to the native registration call, which lowers the plan + * and returns the schema the session settled on. */ protected final Map childInputs; /** - * Declared Arrow {@link org.apache.arrow.vector.types.pojo.Schema} per childStageId, - * parallel to {@link #childInputs}. Used by sinks to coerce incoming batches when - * the shard's actual emit type diverges from the declaration (e.g. DataFusion's - * {@code Utf8View} for string group keys vs. declared {@code Utf8}). + * Declared Arrow {@link org.apache.arrow.vector.types.pojo.Schema} per childStageId. + * Populated lazily by subclasses from the IPC bytes the native registration call + * returns — i.e. the schema the native session itself derived from the producer plan. + * Used by sinks to validate incoming batches via the {@code typesMatch} tripwire. */ protected final Map childSchemas; @@ -108,13 +113,11 @@ protected AbstractDatafusionReduceSink( this.preparedState = preparedState; this.session = preparedState != null ? preparedState.session() : new DatafusionLocalSession(runtimeHandle.get()); Map inputs = new LinkedHashMap<>(ctx.childInputs().size()); - Map schemas = new LinkedHashMap<>(ctx.childInputs().size()); for (ExchangeSinkContext.ChildInput child : ctx.childInputs()) { - inputs.put(child.childStageId(), ArrowSchemaIpc.toBytes(child.schema())); - schemas.put(child.childStageId(), child.schema()); + inputs.put(child.childStageId(), child.producerPlanBytes()); } this.childInputs = inputs; - this.childSchemas = schemas; + this.childSchemas = new LinkedHashMap<>(ctx.childInputs().size()); } /** DataFusion table name for an input partition associated with the given child stage id. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrowSchemaIpc.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrowSchemaIpc.java index 1e8cee72d8c4b..ff48d06b4c8a9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrowSchemaIpc.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ArrowSchemaIpc.java @@ -8,10 +8,13 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.ipc.ArrowStreamReader; import org.apache.arrow.vector.ipc.WriteChannel; import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.Schema; +import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.nio.channels.Channels; @@ -39,4 +42,16 @@ public static byte[] toBytes(Schema schema) { } return baos.toByteArray(); } + + /** Inverse of {@link #toBytes(Schema)}: decodes an Arrow IPC schema-only stream. */ + public static Schema fromBytes(byte[] ipcBytes) { + try ( + RootAllocator allocator = new RootAllocator(); + ArrowStreamReader reader = new ArrowStreamReader(new ByteArrayInputStream(ipcBytes), allocator) + ) { + return reader.getVectorSchemaRoot().getSchema(); + } catch (IOException e) { + throw new IllegalStateException("Failed to decode Arrow IPC schema bytes", e); + } + } } 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 ccb972de88bb3..5e7f8dba9c232 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 @@ -19,6 +19,8 @@ import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.EngineCapability; +import org.opensearch.analytics.spi.ExchangeSink; +import org.opensearch.analytics.spi.ExchangeSinkContext; import org.opensearch.analytics.spi.ExchangeSinkProvider; import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.FilterCapability; @@ -623,32 +625,35 @@ public FragmentInstructionHandlerFactory getInstructionHandlerFactory() { @Override public ExchangeSinkProvider getExchangeSinkProvider() { - return (ctx, backendContext) -> { - DataFusionService svc = plugin.getDataFusionService(); - if (svc == null) { - throw new IllegalStateException("DataFusionService not initialized"); - } - // When the FinalAggregateInstructionHandler has already prepared a plan on the - // coordinator, it hands over a DataFusionReduceState carrying the session + - // registered senders. The sink drives executeLocalPreparedPlan against that - // state instead of re-decoding the fragment bytes. - DataFusionReduceState preparedState = backendContext instanceof DataFusionReduceState s ? s : null; - String mode = plugin.getClusterService() != null - ? plugin.getClusterService().getClusterSettings().get(DataFusionPlugin.DATAFUSION_REDUCE_INPUT_MODE) - : "streaming"; - // Memtable mode is single-input only (DatafusionMemtableReduceSink registers - // exactly one MemTable at close time). Multi-input shapes (Union, future Join) - // need per-child input partitions, which only the streaming sink implements via - // MultiInputExchangeSink#sinkForChild. Auto-fall-back to streaming so end users - // don't have to flip the cluster setting per query. Also fall back when a - // prepared state is supplied (memtable sink does not yet support the - // prepared-plan path). - // TODO: lift this fallback once the memtable sink registers one MemTable per - // child stage (see DatafusionMemtableReduceSink class javadoc). - if ("memtable".equals(mode) && ctx.childInputs().size() == 1 && preparedState == null) { - return new DatafusionMemtableReduceSink(ctx, svc.getNativeRuntime()); + return new ExchangeSinkProvider() { + @Override + public ExchangeSink createSink(ExchangeSinkContext ctx, BackendExecutionContext backendContext) { + DataFusionService svc = plugin.getDataFusionService(); + if (svc == null) { + throw new IllegalStateException("DataFusionService not initialized"); + } + // When the FinalAggregateInstructionHandler has already prepared a plan on the + // coordinator, it hands over a DataFusionReduceState carrying the session + + // registered senders. The sink drives executeLocalPreparedPlan against that + // state instead of re-decoding the fragment bytes. + DataFusionReduceState preparedState = backendContext instanceof DataFusionReduceState s ? s : null; + String mode = plugin.getClusterService() != null + ? plugin.getClusterService().getClusterSettings().get(DataFusionPlugin.DATAFUSION_REDUCE_INPUT_MODE) + : "streaming"; + // Memtable mode is single-input only (DatafusionMemtableReduceSink registers + // exactly one MemTable at close time). Multi-input shapes (Union, future Join) + // need per-child input partitions, which only the streaming sink implements via + // MultiInputExchangeSink#sinkForChild. Auto-fall-back to streaming so end users + // don't have to flip the cluster setting per query. Also fall back when a + // prepared state is supplied (memtable sink does not yet support the + // prepared-plan path). + // TODO: lift this fallback once the memtable sink registers one MemTable per + // child stage (see DatafusionMemtableReduceSink class javadoc). + if ("memtable".equals(mode) && ctx.childInputs().size() == 1 && preparedState == null) { + return new DatafusionMemtableReduceSink(ctx, svc.getNativeRuntime()); + } + return new DatafusionReduceSink(ctx, svc.getNativeRuntime(), svc.getDrainExecutor(), preparedState); } - return new DatafusionReduceSink(ctx, svc.getNativeRuntime(), svc.getDrainExecutor(), preparedState); }; } 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 877c30e34929e..efc9512edf899 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 @@ -65,9 +65,9 @@ * *

    Dispatch summary: *

      - *
    • {@link #convertShardScanFragment(String, RelNode)} and - * {@link #convertFinalAggFragment(RelNode)} — full-fragment conversions via - * {@link #convertToSubstrait(RelNode)}.
      • + *
    • {@link #convertFragment(RelNode)} — full-fragment conversion via + * {@link #convertToSubstrait(RelNode)}, with StageInputScan rewriting + * for reduce-stage fragments.
      • *
    • {@link #attachPartialAggOnTop(RelNode, byte[])} and * {@link #attachFragmentOnTop(RelNode, byte[])} — convert the wrapping * operator standalone, then rewire its input to the decoded inner plan's @@ -280,9 +280,14 @@ public DataFusionFragmentConvertor(SimpleExtension.ExtensionCollection extension } @Override - public byte[] convertShardScanFragment(String tableName, RelNode fragment) { - LOGGER.debug("Converting shard scan fragment for table [{}]", tableName); - return convertToSubstrait(fragment); + public byte[] convertFragment(RelNode fragment) { + LOGGER.debug("Converting fragment [{}]", fragment.getClass().getSimpleName()); + // Rewrite any OpenSearchStageInputScan leaves to plain TableScan nodes so the + // isthmus visitor (which only knows about Calcite core / Logical RelNodes) + // emits a ReadRel with the stage-input-id as the named table. No-op when the + // fragment has no StageInputScan leaves (shard-scan and Values cases). + RelNode rewritten = rewriteStageInputScans(fragment); + return convertToSubstrait(rewritten); } @Override @@ -298,16 +303,6 @@ public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByte return serializePlan(SubstraitPlanRewriter.rewrite(rewired)); } - @Override - public byte[] convertFinalAggFragment(RelNode fragment) { - LOGGER.debug("Converting final-aggregate fragment"); - // Rewrite any OpenSearchStageInputScan leaves to plain TableScan nodes so the - // isthmus visitor (which only knows about Calcite core / Logical RelNodes) - // emits a ReadRel with the stage-input-id as the named table. - RelNode rewritten = rewriteStageInputScans(fragment); - return convertToSubstrait(rewritten); - } - @Override public byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { LOGGER.debug("Attaching generic fragment [{}] on top of {} inner bytes", fragment.getClass().getSimpleName(), innerBytes.length); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionReduceState.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionReduceState.java index f43722c6e21b8..1507f2967fd62 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionReduceState.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionReduceState.java @@ -8,6 +8,7 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.analytics.spi.BackendExecutionContext; import java.io.IOException; @@ -16,8 +17,9 @@ /** * Backend-specific execution context for the coordinator-reduce path when a final-aggregate * plan has been prepared. Carries the local session (with the prepared plan stored on the - * Rust side), the runtime handle, and the partition senders used to feed Arrow batches - * into the streaming input partitions. + * Rust side), the runtime handle, the partition senders used to feed Arrow batches into the + * streaming input partitions, and the schemas the native session settled on for each input + * (parallel to {@code senders}; same order as {@code ctx.childInputs()}). * *

      Produced by {@link FinalAggregateInstructionHandler} and consumed by * {@link DatafusionReduceSink} via the {@link org.opensearch.analytics.spi.ExchangeSinkProvider} @@ -26,7 +28,7 @@ * @opensearch.internal */ public record DataFusionReduceState(DatafusionLocalSession session, NativeRuntimeHandle runtimeHandle, List< - DatafusionPartitionSender> senders) implements BackendExecutionContext { + DatafusionPartitionSender> senders, List inputSchemas) implements BackendExecutionContext { @Override public void close() throws IOException { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSink.java index 43665322a64dc..0b705377c43df 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSink.java @@ -59,7 +59,7 @@ public final class DatafusionMemtableReduceSink extends AbstractDatafusionReduce private final List arrays = new ArrayList<>(); private final List schemas = new ArrayList<>(); - private final byte[] schemaIpc; + private final byte[] producerPlanBytes; public DatafusionMemtableReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtimeHandle) { super(ctx, runtimeHandle); @@ -82,7 +82,7 @@ public DatafusionMemtableReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle + " constructed directly." ); } - this.schemaIpc = childInputs.values().iterator().next(); + this.producerPlanBytes = childInputs.values().iterator().next(); } @Override @@ -121,7 +121,14 @@ protected Throwable closeUnderLock() { // distinct "input-" table id and separate buffer accumulation // per child (the constructor enforces single-input today; see class javadoc). int singleChildStageId = childInputs.keySet().iterator().next(); - NativeBridge.registerMemtable(session.getPointer(), inputIdFor(singleChildStageId), schemaIpc, arrayPtrs, schemaPtrs); + NativeBridge.RegisteredInput registered = NativeBridge.registerMemtable( + session.getPointer(), + inputIdFor(singleChildStageId), + producerPlanBytes, + arrayPtrs, + schemaPtrs + ); + childSchemas.put(singleChildStageId, ArrowSchemaIpc.fromBytes(registered.schemaIpc())); streamPtr = NativeBridge.executeLocalPlan(session.getPointer(), ctx.fragmentBytes()); try (StreamHandle outStream = new StreamHandle(streamPtr, runtimeHandle)) { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java index 74f63b1354d68..fa62ac7246a17 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java @@ -12,11 +12,9 @@ import org.apache.arrow.c.ArrowSchema; import org.apache.arrow.c.Data; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ViewVarCharVector; import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -27,6 +25,7 @@ import org.opensearch.be.datafusion.nativelib.StreamHandle; import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; @@ -96,11 +95,13 @@ public DatafusionReduceSink( try { if (preparedState != null) { // Plan was already prepared by FinalAggregateInstructionHandler. The handler - // registered senders in ctx.childInputs() iteration order; we re-index them - // here by childStageId for lookup during feed(). + // registered senders + captured per-input schemas in ctx.childInputs() + // iteration order; re-index them by childStageId here for lookup during feed(). int i = 0; for (Map.Entry child : childInputs.entrySet()) { - senders.put(child.getKey(), preparedState.senders().get(i++)); + senders.put(child.getKey(), preparedState.senders().get(i)); + childSchemas.put(child.getKey(), preparedState.inputSchemas().get(i)); + i++; } streamPtr = NativeBridge.executeLocalPreparedPlan(session.getPointer()); } else { @@ -111,9 +112,14 @@ public DatafusionReduceSink( // (DataFusionFragmentConvertor names them this way during plan conversion). for (Map.Entry child : childInputs.entrySet()) { int childStageId = child.getKey(); - byte[] schemaIpc = child.getValue(); - long senderPtr = NativeBridge.registerPartitionStream(session.getPointer(), inputIdFor(childStageId), schemaIpc); - senders.put(childStageId, new DatafusionPartitionSender(senderPtr)); + byte[] producerPlanBytes = child.getValue(); + NativeBridge.RegisteredInput registered = NativeBridge.registerPartitionStream( + session.getPointer(), + inputIdFor(childStageId), + producerPlanBytes + ); + senders.put(childStageId, new DatafusionPartitionSender(registered.pointer())); + childSchemas.put(childStageId, ArrowSchemaIpc.fromBytes(registered.schemaIpc())); } streamPtr = NativeBridge.executeLocalPlan(session.getPointer(), ctx.fragmentBytes()); } @@ -200,10 +206,17 @@ private void feedToSender(DatafusionPartitionSender sender, VectorSchemaRoot bat return; } BufferAllocator alloc = ctx.allocator(); - // Bridge DataFusion's physical types (e.g. Utf8View for string group keys) to the - // coordinator's declared schema (Utf8) before handing the batch to Rust. Zero-copy - // fast path when schemas already match. See coerceToDeclaredSchema(). - batch = coerceToDeclaredSchema(batch, declaredSchema, alloc); + // Type-only equality check; nullability and Timestamp precision are advisory. + if (!typesMatch(batch.getSchema(), declaredSchema)) { + batch.close(); + throw new IllegalStateException( + "DatafusionReduceSink: batch schema types do not match declared schema. " + + "declared=" + + declaredSchema + + " batch=" + + batch.getSchema() + ); + } ArrowArray array = ArrowArray.allocateNew(alloc); ArrowSchema arrowSchema = ArrowSchema.allocateNew(alloc); try { @@ -242,66 +255,28 @@ private void feedToSender(DatafusionPartitionSender sender, VectorSchemaRoot bat } /** - * Coerces {@code batch} to {@code declaredSchema} at the Java→Rust boundary. - * Bridges the impedance between DataFusion's physical types (e.g. {@code Utf8View} - * for string group keys, a non-configurable HashAggregate optimization) and - * substrait's logical "string" which the coordinator's FINAL plan consumes as - * {@code Utf8}. One place, explicit, grows per-case on observed mismatch. - * - *

      Zero-copy fast path when schemas already match (numeric-only aggregates). - * Closes {@code batch} — caller drops its reference. - * - *

      TODO (revisit): this runtime coercer bridges a logical/physical type - * mismatch between Calcite's declared exchange schema and DataFusion's physical - * output. A cleaner fix would eliminate the mismatch upstream — for example, a Rust - * pass that casts {@code Utf8View} → {@code Utf8} at the PARTIAL plan's root using - * DataFusion's vectorized {@code CastExpr} (one columnar kernel per batch instead of - * per-cell Java copy), or a Substrait extension that carries view-vs-plain type - * information through the serialized plan. Until one of those lands, this Java-side - * coercer is the minimum correct bridge. + * Field-by-field type equality. Ignores nullability; Timestamp precision/timezone + * parameters are tolerated because the data-node parquet reader and physical + * planner pick a precision the Java-side declaration does not predict, and the + * chosen precision round-trips through Arrow C Data — divergence is harmless. */ - private static VectorSchemaRoot coerceToDeclaredSchema(VectorSchemaRoot batch, Schema declaredSchema, BufferAllocator alloc) { - if (batch.getSchema().equals(declaredSchema)) { - return batch; + private static boolean typesMatch(Schema actual, Schema declared) { + List a = actual.getFields(); + List d = declared.getFields(); + if (a.size() != d.size()) { + return false; } - VectorSchemaRoot out = VectorSchemaRoot.create(declaredSchema, alloc); - try { - out.allocateNew(); - int rows = batch.getRowCount(); - for (int col = 0; col < declaredSchema.getFields().size(); col++) { - FieldVector src = batch.getVector(col); - FieldVector dst = out.getVector(col); - if (src.getField().getType().equals(dst.getField().getType())) { - src.makeTransferPair(dst).transfer(); - continue; - } - ArrowType.ArrowTypeID srcId = src.getField().getType().getTypeID(); - ArrowType.ArrowTypeID dstId = dst.getField().getType().getTypeID(); - if (srcId == ArrowType.ArrowTypeID.Utf8View && dstId == ArrowType.ArrowTypeID.Utf8) { - ViewVarCharVector s = (ViewVarCharVector) src; - VarCharVector d = (VarCharVector) dst; - for (int r = 0; r < rows; r++) { - if (s.isNull(r)) { - d.setNull(r); - } else { - d.setSafe(r, s.get(r)); - } - } - d.setValueCount(rows); - continue; - } - throw new IllegalStateException( - "coerceToDeclaredSchema: unsupported " + srcId + " → " + dstId + " for column '" + dst.getField().getName() + "'" - ); + for (int i = 0; i < a.size(); i++) { + ArrowType at = a.get(i).getType(); + ArrowType dt = d.get(i).getType(); + if (at.getTypeID() == ArrowType.ArrowTypeID.Timestamp && dt.getTypeID() == ArrowType.ArrowTypeID.Timestamp) { + continue; + } + if (!at.equals(dt)) { + return false; } - out.setRowCount(rows); - } catch (RuntimeException e) { - out.close(); - throw e; - } finally { - batch.close(); } - return out; + return true; } /** diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java index 1de82997beb1e..41656c952719c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/FinalAggregateInstructionHandler.java @@ -8,6 +8,7 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.CommonExecutionContext; import org.opensearch.analytics.spi.ExchangeSinkContext; @@ -44,12 +45,17 @@ public BackendExecutionContext apply( DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); List senders = new ArrayList<>(ctx.childInputs().size()); + List inputSchemas = new ArrayList<>(ctx.childInputs().size()); try { for (ExchangeSinkContext.ChildInput child : ctx.childInputs()) { String inputId = "input-" + child.childStageId(); - byte[] schemaIpc = ArrowSchemaIpc.toBytes(child.schema()); - long senderPtr = NativeBridge.registerPartitionStream(session.getPointer(), inputId, schemaIpc); - senders.add(new DatafusionPartitionSender(senderPtr)); + NativeBridge.RegisteredInput registered = NativeBridge.registerPartitionStream( + session.getPointer(), + inputId, + child.producerPlanBytes() + ); + senders.add(new DatafusionPartitionSender(registered.pointer())); + inputSchemas.add(ArrowSchemaIpc.fromBytes(registered.schemaIpc())); } NativeBridge.prepareFinalPlan(session.getPointer(), ctx.fragmentBytes()); } catch (RuntimeException e) { @@ -61,6 +67,6 @@ public BackendExecutionContext apply( session.close(); throw e; } - return new DataFusionReduceState(session, runtimeHandle, senders); + return new DataFusionReduceState(session, runtimeHandle, senders, inputSchemas); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index a937904ef9ab2..1a3960830ca37 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -205,7 +205,9 @@ public final class NativeBridge { FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) ); - // i64 df_register_partition_stream(session_ptr, input_id_ptr, input_id_len, schema_ipc_ptr, schema_ipc_len) + // i64 df_register_partition_stream(session_ptr, input_id_ptr, input_id_len, + // partial_plan_ptr, partial_plan_len, + // out_ptr, out_cap, out_len) REGISTER_PARTITION_STREAM = linker.downcallHandle( lib.find("df_register_partition_stream").orElseThrow(), FunctionDescriptor.of( @@ -214,7 +216,10 @@ public final class NativeBridge { ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS ) ); @@ -233,8 +238,10 @@ public final class NativeBridge { // void df_sender_close(sender_ptr) SENDER_CLOSE = linker.downcallHandle(lib.find("df_sender_close").orElseThrow(), FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG)); - // i64 df_register_memtable(session_ptr, input_id_ptr, input_id_len, schema_ipc_ptr, schema_ipc_len, - // array_ptrs, schema_ptrs, n_batches) + // i64 df_register_memtable(session_ptr, input_id_ptr, input_id_len, + // partial_plan_ptr, partial_plan_len, + // array_ptrs, schema_ptrs, n_batches, + // out_ptr, out_cap, out_len) REGISTER_MEMTABLE = linker.downcallHandle( lib.find("df_register_memtable").orElseThrow(), FunctionDescriptor.of( @@ -246,7 +253,10 @@ public final class NativeBridge { ValueLayout.JAVA_LONG, ValueLayout.ADDRESS, ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS ) ); @@ -717,6 +727,16 @@ public static byte[] sqlToSubstrait(long readerPtr, String tableName, String sql // ---- Coordinator-reduce exports ---- + /** + * Pair returned from {@link #registerPartitionStream} / {@link #registerMemtable}: the + * native sender pointer (or 0 for memtable) plus the Arrow IPC-encoded schema the native + * session derived by lowering the producer-side substrait. The Java tripwire + * ({@code typesMatch} in {@code DatafusionReduceSink}) validates fed batches against this + * schema, and downstream callers decode it once into an Arrow {@link org.apache.arrow.vector.types.pojo.Schema}. + */ + public record RegisteredInput(long pointer, byte[] schemaIpc) { + } + /** * Creates a local DataFusion session tied to the given global runtime. Returns an opaque * native pointer freed by {@link #closeLocalSession}. @@ -734,21 +754,28 @@ public static void closeLocalSession(long sessionPtr) { } /** - * Registers an input partition stream on the session under {@code inputId}, with the given - * Arrow IPC-encoded schema. Returns an opaque sender pointer freed by {@link #senderClose}. + * Registers an input partition stream on the session under {@code inputId}, deriving the + * input schema by lowering the producer-side {@code partialPlanBytes}. Returns the native + * sender pointer (freed by {@link #senderClose}) and the Arrow IPC-encoded schema the + * native session settled on after lowering. */ - public static long registerPartitionStream(long sessionPtr, String inputId, byte[] schemaIpc) { + public static RegisteredInput registerPartitionStream(long sessionPtr, String inputId, byte[] partialPlanBytes) { NativeHandle.validatePointer(sessionPtr, "session"); try (var call = new NativeCall()) { var id = call.str(inputId); - return call.invoke( + var out = call.outBuffer(64 * 1024); + long ptr = call.invoke( REGISTER_PARTITION_STREAM, sessionPtr, id.segment(), id.len(), - call.bytes(schemaIpc), - (long) schemaIpc.length + call.bytes(partialPlanBytes), + (long) partialPlanBytes.length, + out.data(), + (long) out.capacity(), + out.lenOut() ); + return new RegisteredInput(ptr, out.toByteArray()); } } @@ -789,10 +816,21 @@ public static void senderClose(long senderPtr) { /** * Memtable variant of {@link #registerPartitionStream}: hands across a list of - * already-exported Arrow C Data batches in two parallel pointer arrays so the native side can - * build a {@code MemTable} in one shot. Native takes ownership of all FFI structs on success. + * already-exported Arrow C Data batches in two parallel pointer arrays so the native side + * can build a {@code MemTable} in one shot. Schema is derived by lowering the producer-side + * {@code partialPlanBytes}; native takes ownership of all FFI structs on success. + * + *

      Returns a {@link RegisteredInput} whose {@code pointer} field is always 0 (memtable + * registration has no sender to return) — the {@code schemaIpc} field carries the schema + * the native session settled on after lowering. */ - public static long registerMemtable(long sessionPtr, String inputId, byte[] schemaIpc, long[] arrayPtrs, long[] schemaPtrs) { + public static RegisteredInput registerMemtable( + long sessionPtr, + String inputId, + byte[] partialPlanBytes, + long[] arrayPtrs, + long[] schemaPtrs + ) { NativeHandle.validatePointer(sessionPtr, "session"); if (arrayPtrs.length != schemaPtrs.length) { throw new IllegalArgumentException( @@ -801,17 +839,22 @@ public static long registerMemtable(long sessionPtr, String inputId, byte[] sche } try (var call = new NativeCall()) { var id = call.str(inputId); - return call.invoke( + var out = call.outBuffer(64 * 1024); + long ptr = call.invoke( REGISTER_MEMTABLE, sessionPtr, id.segment(), id.len(), - call.bytes(schemaIpc), - (long) schemaIpc.length, + call.bytes(partialPlanBytes), + (long) partialPlanBytes.length, call.longs(arrayPtrs), call.longs(schemaPtrs), - (long) arrayPtrs.length + (long) arrayPtrs.length, + out.data(), + (long) out.capacity(), + out.lenOut() ); + return new RegisteredInput(ptr, out.toByteArray()); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java index e1140aa86761d..12a4162f19ba3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java @@ -144,7 +144,7 @@ private LogicalAggregate buildSumAggregate(RelNode input, int columnIndex) { */ public void testConvertShardScanFragment_TableScan() throws Exception { RelNode scan = buildTableScan("test_index", "A", "B"); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", scan); + byte[] bytes = newConvertor().convertFragment(scan); Plan plan = decodeSubstrait(bytes); Rel root = rootRel(plan); @@ -166,7 +166,7 @@ public void testConvertShardScanFragment_FilterOverScan() throws Exception { ); RelNode filter = LogicalFilter.create(scan, predicate); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", filter); + byte[] bytes = newConvertor().convertFragment(filter); Plan plan = decodeSubstrait(bytes); Rel root = rootRel(plan); @@ -187,7 +187,7 @@ public void testAttachPartialAggOnTop_WrapsInner() throws Exception { // Inner bytes from a shard-scan conversion. RelNode scan = buildTableScan("test_index", "A"); - byte[] innerBytes = convertor.convertShardScanFragment("test_index", scan); + byte[] innerBytes = convertor.convertFragment(scan); // Build a bare partial-agg fragment whose input matches the inner's rowType. LogicalAggregate partialAgg = buildSumAggregate(scan, 0); @@ -224,7 +224,7 @@ public void testConvertFinalAggFragment_WithStageInputScanLeaf() throws Exceptio RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, stageRowType, List.of("datafusion")); LogicalAggregate finalAgg = buildSumAggregate(stageInput, 0); - byte[] bytes = newConvertor().convertFinalAggFragment(finalAgg); + byte[] bytes = newConvertor().convertFragment(finalAgg); Plan plan = decodeSubstrait(bytes); Rel root = rootRel(plan); @@ -255,7 +255,7 @@ public void testAttachFragmentOnTop_Sort() throws Exception { int childStageId = 3; RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, stageRowType, List.of("datafusion")); LogicalAggregate finalAgg = buildSumAggregate(stageInput, 0); - byte[] innerBytes = convertor.convertFinalAggFragment(finalAgg); + byte[] innerBytes = convertor.convertFragment(finalAgg); // Contract: attachFragmentOnTop receives a childless operator. Sort requires an // input for row-type validation in the isthmus visitor; give it a bare placeholder @@ -291,7 +291,7 @@ public void testAttachPartialAggOnTop_PlanRootNamesMatchWrapperOutput() throws E // Inner scan has 3 columns; the partial-aggregate emits 1 (sum over col 0). RelNode scan = buildTableScan("test_index", "A", "B", "C"); - byte[] innerBytes = convertor.convertShardScanFragment("test_index", scan); + byte[] innerBytes = convertor.convertFragment(scan); LogicalAggregate partialAgg = buildSumAggregate(scan, 0); byte[] combined = convertor.attachPartialAggOnTop(partialAgg, innerBytes); @@ -321,13 +321,13 @@ public void testAttachFragmentOnTop_AggregateOverMultiColumnInner_PlanRootNamesM RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, wideStageRowType, List.of("datafusion")); // For this regression, the inner doesn't need to be a final-agg — a bare scan-shaped // plan with 3-column rowType is enough to surface the wrapper-vs-inner names mismatch. - // Use convertFinalAggFragment so the inner Plan.Root.names is the 3-column scan list. + // Use convertFragment so the inner Plan.Root.names is the 3-column scan list. RelNode innerStageScan = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, wideStageRowType, List.of("datafusion")); // Wrap it in a no-op aggregate so the convertor accepts it as a final-agg fragment shape. // The inner's Plan.Root.names then carries the agg-output (1 col, "sum_col"), but the // *wrapper* we attach above has its own output rowType. LogicalAggregate innerFinalAgg = buildSumAggregate(innerStageScan, 0); - byte[] innerBytes = convertor.convertFinalAggFragment(innerFinalAgg); + byte[] innerBytes = convertor.convertFragment(innerFinalAgg); // Wrapper: a Project that maps the single inner column to two new aliases — this is // the multisearch-style schema reshape that triggered the bug. We model it as another @@ -352,7 +352,7 @@ public void testAttachFragmentOnTop_AggregateOverMultiColumnInner_PlanRootNamesM /** * Mirror of multisearch's coordinator-stage shape: * {@code Sort(Aggregate(Union(StageInputScan, StageInputScan, StageInputScan)))}. - * After the convertor chain runs (convertFinalAggFragment(Union) → + * After the convertor chain runs (convertFragment(Union) → * attachFragmentOnTop(Aggregate) → attachFragmentOnTop(Sort)), the outermost * {@code Plan.Root.names} must reflect the Sort's output schema (= the * aggregate's 1-column output), not the inner Union's wider row type. @@ -368,7 +368,7 @@ public void testMultisearchShape_SortOverAggregateOverThreeWayUnion_PlanRootName RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion")); RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion")); LogicalUnion union = LogicalUnion.create(List.of(sin1, sin2, sin3), true); - byte[] unionBytes = convertor.convertFinalAggFragment(union); + byte[] unionBytes = convertor.convertFragment(union); // Aggregate over the union: SUM(a) → 1 column output ("sum_col"). // attachFragmentOnTop expects the wrapper to carry its real input so the @@ -408,7 +408,7 @@ public void testMultisearchShape_SystemLimitOverSortOverAggregateOverUnion_Names RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion")); RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion")); LogicalUnion union = LogicalUnion.create(List.of(sin1, sin2, sin3), true); - byte[] unionBytes = convertor.convertFinalAggFragment(union); + byte[] unionBytes = convertor.convertFragment(union); // Aggregate over the union: SUM(a) → 1 column. LogicalAggregate aggregate = buildSumAggregate(union, 0); @@ -444,7 +444,7 @@ public void testConvertShardScanFragment_DelegatedPredicatePlaceholder() throws RexNode placeholder = DelegatedPredicateFunction.makeCall(rexBuilder, 42); RelNode filter = LogicalFilter.create(scan, placeholder); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", filter); + byte[] bytes = newConvertor().convertFragment(filter); Plan plan = decodeSubstrait(bytes); Rel root = rootRel(plan); @@ -474,7 +474,7 @@ public void testConvertShardScanFragment_MixedNativeAndDelegated() throws Except RexNode andCondition = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, delegated); RelNode filter = LogicalFilter.create(scan, andCondition); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", filter); + byte[] bytes = newConvertor().convertFragment(filter); Plan plan = decodeSubstrait(bytes); FilterRel filterRel = rootRel(plan).getFilter(); // Root condition is AND (scalar function with 2 args) @@ -506,7 +506,7 @@ public void testConvertShardScanFragment_ComplexBooleanTreeWithDelegation() thro RexNode andCondition = rexBuilder.makeCall(SqlStdOperatorTable.AND, nativePred, orClause); RelNode filter = LogicalFilter.create(scan, andCondition); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", filter); + byte[] bytes = newConvertor().convertFragment(filter); Plan plan = decodeSubstrait(bytes); logger.info("Substrait plan (complex boolean tree):\n{}", plan); FilterRel filterRel = rootRel(plan).getFilter(); @@ -558,7 +558,7 @@ public void testApproxCountDistinctRenamed() throws Exception { ); LogicalAggregate agg = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(approxCall)); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", agg); + byte[] bytes = newConvertor().convertFragment(agg); Plan plan = decodeSubstrait(bytes); boolean foundApproxDistinct = false; @@ -603,7 +603,7 @@ public void testProjectTimestampOutputCastEmitsToCharExtension() throws Exceptio java.util.Set.of() ); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", project); + byte[] bytes = newConvertor().convertFragment(project); Plan plan = decodeSubstrait(bytes); boolean foundToChar = false; @@ -624,7 +624,7 @@ public void testOtherFunctionsNotRenamed() throws Exception { RelNode scan = buildTableScan("test_index", "A"); LogicalAggregate agg = buildSumAggregate(scan, 0); - byte[] bytes = newConvertor().convertShardScanFragment("test_index", agg); + byte[] bytes = newConvertor().convertFragment(agg); Plan plan = decodeSubstrait(bytes); boolean foundSum = false; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java index 0e6d57134bc41..d8b38a5881a05 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionMemtableReduceSinkTests.java @@ -68,7 +68,7 @@ public void testFeedDrainsSumToDownstream() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstraitBytes(DatafusionMemtableReduceSink.INPUT_ID))), downstream ); @@ -103,7 +103,27 @@ private static byte[] buildSumSubstraitBytes(String inputId) { AggregateCall sumCall = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, bigintNullable, "total"); LogicalAggregate agg = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sumCall)); - return new DataFusionFragmentConvertor(loadExtensions()).convertFinalAggFragment(agg); + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(agg); + } + + /** + * Bare {@code SELECT * FROM "input-0"} substrait — used as the producer-side plan in + * {@link ExchangeSinkContext.ChildInput#producerPlanBytes()}. Its lowered output schema + * is the leaf row type ({@code x: BIGINT}), which the sink registers as the input + * partition's declared schema. + */ + private static byte[] buildPassthroughSubstraitBytes(String inputId) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner hepPlanner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(hepPlanner, rexBuilder); + + RelDataType bigintNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true); + RelDataType rowType = typeFactory.builder().add("x", bigintNullable).build(); + + RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), inputId, rowType); + + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(scan); } private static SimpleExtension.ExtensionCollection loadExtensions() { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java index a1b794bca846a..a7194e3e9ce28 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java @@ -96,7 +96,7 @@ public void testFeedDrainsSumToDownstream() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstraitBytes(DatafusionReduceSink.INPUT_ID))), downstream ); @@ -142,7 +142,7 @@ public void testDrainTaskKeepsUpWithProducer() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstraitBytes(DatafusionReduceSink.INPUT_ID))), downstream ); @@ -167,6 +167,26 @@ public void testDrainTaskKeepsUpWithProducer() throws Exception { // ── Helpers ────────────────────────────────────────────────────────────── + /** + * Builds Substrait bytes for a plain {@code SELECT * FROM "input-0"} — used as + * the producer-side plan in {@link ExchangeSinkContext.ChildInput#producerPlanBytes()}. + * The lowered output schema is the bare leaf row type (single BIGINT column {@code x}) + * which is what the reduce sink registers as the input partition's declared schema. + */ + private static byte[] buildPassthroughSubstraitBytes(String inputId) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner hepPlanner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(hepPlanner, rexBuilder); + + RelDataType bigintNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true); + RelDataType rowType = typeFactory.builder().add("x", bigintNullable).build(); + + RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), inputId, rowType); + + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(scan); + } + /** * Builds Substrait bytes for {@code SELECT SUM(x) FROM "input-0"} using the * production {@link DataFusionFragmentConvertor} path — the same conversion @@ -187,7 +207,7 @@ private static byte[] buildSumSubstraitBytes(String inputId) { AggregateCall sumCall = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, bigintNullable, "total"); LogicalAggregate agg = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sumCall)); - return new DataFusionFragmentConvertor(loadExtensions()).convertFinalAggFragment(agg); + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(agg); } /** diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java index c2b4d8120fdfc..c80e2434c5908 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/NativeBridgeLocalSessionTests.java @@ -14,20 +14,28 @@ import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.ipc.WriteChannel; -import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.pojo.ArrowType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.sql.type.SqlTypeName; import org.opensearch.be.datafusion.nativelib.NativeBridge; import org.opensearch.test.OpenSearchTestCase; -import java.io.ByteArrayOutputStream; -import java.nio.channels.Channels; import java.nio.file.Path; import java.util.List; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; + /** * Smoke test for the coordinator-reduce FFM wrappers added by the datafusion-coordinator-reduce spec. * @@ -50,12 +58,32 @@ private NativeRuntimeHandle createRuntime() { return new NativeRuntimeHandle(runtimePtr); } - private static byte[] schemaIpc(Schema schema) throws Exception { - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - try (WriteChannel channel = new WriteChannel(Channels.newChannel(baos))) { - MessageSerializer.serialize(channel, schema); + /** + * Bare {@code SELECT * FROM "input-0"} substrait whose lowered output schema is a single + * BIGINT column named {@code x} — used as the producer-side plan that + * {@code registerPartitionStream} / {@code registerMemtable} now derive their input + * schema from. + */ + private static byte[] passthroughSubstrait(String inputId) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner hepPlanner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(hepPlanner, rexBuilder); + + RelDataType bigintNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true); + RelDataType rowType = typeFactory.builder().add("x", bigintNullable).build(); + + RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), inputId, rowType); + + Thread t = Thread.currentThread(); + ClassLoader prev = t.getContextClassLoader(); + try { + t.setContextClassLoader(NativeBridgeLocalSessionTests.class.getClassLoader()); + SimpleExtension.ExtensionCollection ext = DefaultExtensionCatalog.DEFAULT_COLLECTION; + return new DataFusionFragmentConvertor(ext).convertFragment(scan); + } finally { + t.setContextClassLoader(prev); } - return baos.toByteArray(); } public void testCreateLocalSessionReturnsNonZeroPtr() { @@ -78,15 +106,20 @@ public void testSenderCloseToleratesZero() { NativeBridge.senderClose(0L); } - public void testRegisterPartitionStreamAndSenderClose() throws Exception { + public void testRegisterPartitionStreamAndSenderClose() { NativeRuntimeHandle runtimeHandle = createRuntime(); try { DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); try { - Schema schema = new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(64, true)), null))); - long senderPtr = NativeBridge.registerPartitionStream(session.getPointer(), "input-0", schemaIpc(schema)); - assertTrue("sender ptr non-zero", senderPtr != 0); - NativeBridge.senderClose(senderPtr); + NativeBridge.RegisteredInput registered = NativeBridge.registerPartitionStream( + session.getPointer(), + "input-0", + passthroughSubstrait("input-0") + ); + assertTrue("sender ptr non-zero", registered.pointer() != 0); + assertNotNull("schema IPC bytes returned", registered.schemaIpc()); + assertTrue("schema IPC non-empty", registered.schemaIpc().length > 0); + NativeBridge.senderClose(registered.pointer()); } finally { session.close(); } @@ -95,13 +128,20 @@ public void testRegisterPartitionStreamAndSenderClose() throws Exception { } } - public void testRegisterMemtableAcceptsZeroBatches() throws Exception { + public void testRegisterMemtableAcceptsZeroBatches() { NativeRuntimeHandle runtimeHandle = createRuntime(); try { DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); try { - Schema schema = new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(64, true)), null))); - NativeBridge.registerMemtable(session.getPointer(), "input-0", schemaIpc(schema), new long[0], new long[0]); + NativeBridge.RegisteredInput registered = NativeBridge.registerMemtable( + session.getPointer(), + "input-0", + passthroughSubstrait("input-0"), + new long[0], + new long[0] + ); + assertNotNull("schema IPC bytes returned", registered.schemaIpc()); + assertTrue("schema IPC non-empty", registered.schemaIpc().length > 0); } finally { session.close(); } @@ -110,7 +150,7 @@ public void testRegisterMemtableAcceptsZeroBatches() throws Exception { } } - public void testRegisterMemtableImportsBatch() throws Exception { + public void testRegisterMemtableImportsBatch() { NativeRuntimeHandle runtimeHandle = createRuntime(); try (RootAllocator alloc = new RootAllocator(Long.MAX_VALUE)) { DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); @@ -128,7 +168,7 @@ public void testRegisterMemtableImportsBatch() throws Exception { NativeBridge.registerMemtable( session.getPointer(), "input-0", - schemaIpc(schema), + passthroughSubstrait("input-0"), new long[] { array.memoryAddress() }, new long[] { arrowSchema.memoryAddress() } ); @@ -143,18 +183,17 @@ public void testRegisterMemtableImportsBatch() throws Exception { } } - public void testRegisterMemtableRejectsLengthMismatch() throws Exception { + public void testRegisterMemtableRejectsLengthMismatch() { NativeRuntimeHandle runtimeHandle = createRuntime(); try { DatafusionLocalSession session = new DatafusionLocalSession(runtimeHandle.get()); try { - Schema schema = new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(64, true)), null))); expectThrows( IllegalArgumentException.class, () -> NativeBridge.registerMemtable( session.getPointer(), "input-0", - schemaIpc(schema), + passthroughSubstrait("input-0"), new long[] { 1L, 2L }, new long[] { 1L } ) diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java index eb0bd161abbd7..a88c4d0e281c4 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java @@ -271,13 +271,8 @@ public ExchangeSinkProvider getExchangeSinkProvider() { public FragmentConvertor getFragmentConvertor() { return new FragmentConvertor() { @Override - public byte[] convertShardScanFragment(String tableName, RelNode fragment) { - return ("shard:" + tableName).getBytes(StandardCharsets.UTF_8); - } - - @Override - public byte[] convertFinalAggFragment(RelNode fragment) { - return "reduce".getBytes(StandardCharsets.UTF_8); + public byte[] convertFragment(RelNode fragment) { + return "fragment".getBytes(StandardCharsets.UTF_8); } @Override 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 deleted file mode 100644 index e1e04e6ab126b..0000000000000 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ArrowSchemaFromCalcite.java +++ /dev/null @@ -1,108 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.analytics.exec.stage; - -import org.apache.arrow.vector.types.DateUnit; -import org.apache.arrow.vector.types.FloatingPointPrecision; -import org.apache.arrow.vector.types.TimeUnit; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.calcite.rel.type.RelDataType; -import org.apache.calcite.rel.type.RelDataTypeField; -import org.apache.calcite.sql.type.SqlTypeName; - -import java.util.ArrayList; -import java.util.List; - -/** - * Translates a Calcite {@link RelDataType} (row type) to an Arrow {@link Schema}. - * Used by distributed stages to declare their exchange-point schema when registering - * {@code StreamingTable} partitions with the native execution engine. - * - *

      All fields are nullable for MVP. - */ -final class ArrowSchemaFromCalcite { - - private ArrowSchemaFromCalcite() {} - - /** - * Convert a Calcite row type to an Arrow schema. All fields are nullable. - * - * @param rowType the Calcite row type from a RelNode fragment - * @return the corresponding Arrow schema - */ - public static Schema arrowSchemaFromRowType(RelDataType rowType) { - List fields = new ArrayList<>(); - for (RelDataTypeField f : rowType.getFieldList()) { - 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: - return new ArrowType.Int(64, true); - case INTEGER: - return new ArrowType.Int(32, true); - case SMALLINT: - return new ArrowType.Int(16, true); - case TINYINT: - return new ArrowType.Int(8, true); - case DOUBLE: - return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); - case FLOAT: - case REAL: - return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); - case BOOLEAN: - return ArrowType.Bool.INSTANCE; - case VARCHAR: - case CHAR: - return ArrowType.Utf8.INSTANCE; - case VARBINARY: - case BINARY: - return ArrowType.Binary.INSTANCE; - case DATE: - return new ArrowType.Date(DateUnit.DAY); - case TIME: - return new ArrowType.Time(TimeUnit.MILLISECOND, 32); - case TIMESTAMP: - case TIMESTAMP_WITH_LOCAL_TIME_ZONE: - return new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); - default: - throw new IllegalArgumentException("Unsupported Calcite SQL type: " + sqlTypeName); - } - } -} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java index c2c44a59dc5f2..367c5f24ea589 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java @@ -27,9 +27,9 @@ * those that run at the coordinator with a backend-provided {@link ExchangeSink}. * Creates the sink via {@link Stage#getExchangeSinkProvider()} using an * {@link ExchangeSinkContext} carrying the plan bytes, allocator, per-child - * input descriptors (one per child stage, each with its stage id + Arrow - * schema), and the downstream sink. Hands the resulting sink to - * {@link LocalStageExecution}. + * input descriptors (one per child stage, each with its stage id + the + * producer-side plan bytes the backend lowers to derive the input schema), + * and the downstream sink. Hands the resulting sink to {@link LocalStageExecution}. * *

      Multi-child stages (Union, future Join) are routed via * {@link LocalStageExecution#inputSink(int)}, which returns a per-child @@ -118,10 +118,10 @@ private static byte[] chosenBytes(Stage stage) { } /** - * Builds one {@link ExchangeSinkContext.ChildInput} per child stage. Each entry - * carries the child's stage id (used by the backend to namespace its registered - * input, e.g. {@code "input-"}) and the Arrow schema derived from the - * child fragment's row type. + * Builds one {@link ExchangeSinkContext.ChildInput} per child stage. Each entry carries + * the child's stage id (used by the backend to namespace its registered input, e.g. + * {@code "input-"}) and the producer-side plan bytes the backend lowers to + * derive the input schema at registration time. */ private static List buildChildInputs(Stage stage) { List children = stage.getChildStages(); @@ -132,12 +132,8 @@ private static List buildChildInputs(Stage stage } List inputs = new ArrayList<>(children.size()); for (Stage child : children) { - inputs.add( - new ExchangeSinkContext.ChildInput( - child.getStageId(), - ArrowSchemaFromCalcite.arrowSchemaFromRowType(child.getFragment().getRowType()) - ) - ); + byte[] producerPlanBytes = child.getPlanAlternatives().getFirst().convertedBytes(); + inputs.add(new ExchangeSinkContext.ChildInput(child.getStageId(), producerPlanBytes)); } return inputs; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java index 4a8386da53321..b4dfee1d05eea 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java @@ -48,12 +48,12 @@ *

      Dispatch logic for PR2 (pure shard-scan path): *

        *
      • Leaf = {@link OpenSearchTableScan}, top = {@link OpenSearchAggregate}(PARTIAL): - * {@code convertShardScanFragment} on everything below partial agg, + * {@code convertFragment} on everything below partial agg, * then {@code attachPartialAggOnTop}
        • *
      • Leaf = {@link OpenSearchTableScan}, top = anything else: - * {@code convertShardScanFragment} on the full fragment
        • + * {@code convertFragment} on the full fragment *
      • Leaf = {@link OpenSearchStageInputScan} (reduce stage): - * {@code convertFinalAggFragment} on the final agg (ExchangeReducer stripped), + * {@code convertFragment} on the final agg (ExchangeReducer stripped), * then {@code attachFragmentOnTop} for any operators above it
        • *
      * @@ -210,21 +210,19 @@ List getResult() { static byte[] convert(RelNode resolvedFragment, FragmentConvertor convertor, IntraOperatorDelegationBytes delegationBytes) { RelNode leaf = findLeaf(resolvedFragment); - if (leaf instanceof OpenSearchTableScan scan) { - String tableName = scan.getTable().getQualifiedName().getLast(); - + if (leaf instanceof OpenSearchTableScan) { // Partial agg at top: convert everything below it, then attach partial agg on top. // strippedInputs passed to stripAnnotations for schema validity (LogicalAggregate needs its inputs). if (resolvedFragment instanceof OpenSearchAggregate agg && agg.getMode() == AggregateMode.PARTIAL) { List strippedInputs = agg.getInputs().stream().map(input -> strip(input, delegationBytes)).toList(); - byte[] innerBytes = convertor.convertShardScanFragment(tableName, strippedInputs.getFirst()); + byte[] innerBytes = convertor.convertFragment(strippedInputs.getFirst()); Function resolver = delegationBytes.resolverFor(agg, agg.getCluster().getRexBuilder()); RelNode strippedAgg = agg.stripAnnotations(strippedInputs, resolver); return convertor.attachPartialAggOnTop(strippedAgg, innerBytes); } RelNode stripped = strip(resolvedFragment, delegationBytes); - return convertor.convertShardScanFragment(tableName, stripped); + return convertor.convertFragment(stripped); } if (leaf instanceof OpenSearchStageInputScan) { @@ -236,7 +234,7 @@ static byte[] convert(RelNode resolvedFragment, FragmentConvertor convertor, Int // path as reduce fragments. isthmus emits ReadRel.VirtualTable for the Values // leaf; DataFusion executes it locally without any input partitions. RelNode stripped = strip(resolvedFragment, delegationBytes); - return convertor.convertFinalAggFragment(stripped); + return convertor.convertFragment(stripped); } throw new IllegalStateException( @@ -250,11 +248,11 @@ static byte[] convert(RelNode resolvedFragment, FragmentConvertor convertor, Int * (Sort, Project, etc.) via attachFragmentOnTop. * *

      Single-input ancestors of a single gathered subtree (Sort/Project/Aggregate over - * a partial agg) reach convertFinalAggFragment as soon as we see a node whose inputs + * a partial agg) reach convertFragment as soon as we see a node whose inputs * are all ExchangeReducers, and attach via attachFragmentOnTop on the way back up. * *

      Multi-input nodes (Join, Union, Intersect, Minus) are converted as a single - * subtree via convertFinalAggFragment: isthmus handles all of them natively, and + * subtree via convertFragment: isthmus handles all of them natively, and * rewriting OpenSearchStageInputScan leaves to plain TableScans (inside the convertor) * lets the whole gathered subtree serialize in one pass. No post-conversion * substrait-level stitching is needed. @@ -271,7 +269,7 @@ private static byte[] convertReduceNode( ) { if (node instanceof OpenSearchExchangeReducer) { // Strip ExchangeReducer — StageInputScan below it is the schema source. - return convertor.convertFinalAggFragment(strip(node.getInputs().getFirst(), delegationBytes)); + return convertor.convertFragment(strip(node.getInputs().getFirst(), delegationBytes)); } if (node instanceof OpenSearchRelNode openSearchNode) { List strippedInputs = node.getInputs().stream().map(input -> strip(input, delegationBytes)).toList(); @@ -295,17 +293,17 @@ private static byte[] convertReduceNode( finalAggInputs.add(strip(input.getInputs().getFirst(), delegationBytes)); } RelNode finalAggFragment = openSearchNode.stripAnnotations(finalAggInputs, resolver); - return convertor.convertFinalAggFragment(finalAggFragment); + return convertor.convertFragment(finalAggFragment); } } // Multi-input node (Join, Union, Intersect, Minus): isthmus handles all of them // natively. The whole subtree — multi-input node + its branches + ERs + - // StageInputScans — serializes in one convertFinalAggFragment pass. The convertor's + // StageInputScans — serializes in one convertFragment pass. The convertor's // StageInputScan → plain TableScan rewrite makes the leaves isthmus-friendly without // any post-conversion substrait-level stitching. if (node.getInputs().size() >= 2) { - return convertor.convertFinalAggFragment(strip(node, delegationBytes)); + return convertor.convertFragment(strip(node, delegationBytes)); } // Single-input operator above the final-fragment boundary — convert child first, then attach. diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java index df66e46b4c77e..5f97f086df1e1 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java @@ -120,7 +120,7 @@ private void assertDoesntContainOperators(RelNode fragment, Set forbidde private void assertShardScanConverted(RecordingConvertor convertor, Stage stage) { assertEquals("expected exactly one alternative", 1, stage.getPlanAlternatives().size()); assertNotNull("convertedBytes must be set", stage.getPlanAlternatives().getFirst().convertedBytes()); - assertTrue("convertShardScanFragment must be called", convertor.shardScanCalled); + assertTrue("convertFragment (shard-scan shape) must be called", convertor.shardScanCalled); assertEquals("test_index", convertor.shardScanTableName); assertDoesntContainOperators(convertor.shardScanFragment, OPENSEARCH_OPERATORS); assertDoesntContainOperators(convertor.shardScanFragment, ANNOTATION_MARKERS); @@ -133,7 +133,7 @@ private void assertShardScanConverted(RecordingConvertor convertor, Stage stage) private void assertReduceStageConverted(RecordingConvertor convertor, Stage stage) { assertEquals("expected exactly one alternative", 1, stage.getPlanAlternatives().size()); assertNotNull("convertedBytes must be set", stage.getPlanAlternatives().getFirst().convertedBytes()); - assertTrue("convertFinalAggFragment must be called", convertor.finalAggCalled); + assertTrue("convertFragment (final-agg shape) must be called", convertor.finalAggCalled); assertDoesntContainOperators(convertor.reduceFragment, OPENSEARCH_OPERATORS); assertDoesntContainOperators(convertor.reduceFragment, ANNOTATION_MARKERS); // Coord-side reduce stages no longer register FinalAggregateInstructionHandler. @@ -149,7 +149,7 @@ private void assertReduceStageConverted(RecordingConvertor convertor, Stage stag /** * Scan, Filter(Scan), Aggregate(Scan), Sort(Filter(Scan)) — single-shard plans now * have a coord stage above the data-node stage (since scans declare RANDOM, the - * coord must gather). The test verifies convertShardScanFragment is called on the + * coord must gather). The test verifies convertFragment (shard-scan shape) is called on the * data-node child stage and the fragment is fully stripped. */ public void testSingleStageQueryShapes() { @@ -225,8 +225,8 @@ public void testSortOnAggregateOnFilteredScan() { // ---- Two-stage shapes ---- /** - * Multi-shard Aggregate(Scan) — child calls convertShardScanFragment, - * root calls convertFinalAggFragment. + * Multi-shard Aggregate(Scan) — child calls convertFragment (shard-scan shape), + * root calls convertFragment (final-agg shape). */ public void testTwoStageAggregateConversion() { RecordingConvertor convertor = new RecordingConvertor(); @@ -264,7 +264,7 @@ public void testTwoStageSortOnAggregateOnFilteredScan() { /** * Coord-side fragment: Aggregate ← Join ← (ER ← ...) | (ER ← ...). * Both branches are gathered subtrees. convertReduceNode must convert the whole Join + - * branches + ERs + StageInputScans subtree in a single {@code convertFinalAggFragment} + * branches + ERs + StageInputScans subtree in a single {@code convertFragment (final-agg shape)} * pass — same path as Union / Intersect / Minus. No substrait-level join stitching. */ public void testJoinDirectlyOverTwoExchanges() { @@ -276,7 +276,7 @@ public void testJoinDirectlyOverTwoExchanges() { Stage joinStage = findStageWithTwoChildren(dag.rootStage()); assertNotNull("expected a stage with 2 child stages (the coord-side Join stage)", joinStage); assertNotNull("join stage alternative must have convertedBytes", joinStage.getPlanAlternatives().getFirst().convertedBytes()); - assertTrue("convertFinalAggFragment must be called for the Join subtree", convertor.finalAggCalled); + assertTrue("convertFragment (final-agg shape) must be called for the Join subtree", convertor.finalAggCalled); } private static Stage findStageWithTwoChildren(Stage stage) { @@ -291,7 +291,7 @@ private static Stage findStageWithTwoChildren(Stage stage) { /** * Coord-side Union with pass-through operators (Sort/Project) between each arm and its * ER. Isthmus's SubstraitRelVisitor handles Union natively; convertReduceNode converts - * the whole Union subtree as one convertFinalAggFragment call — same path as Join. + * the whole Union subtree as one convertFragment (final-agg shape) call — same path as Join. */ public void testUnionOverPassthroughThenExchange() { RecordingConvertor convertor = new RecordingConvertor(); @@ -318,7 +318,7 @@ public org.opensearch.analytics.spi.FragmentConvertor getFragmentConvertor() { Stage root = dag.rootStage(); assertNotNull("root alternative must have convertedBytes", root.getPlanAlternatives().getFirst().convertedBytes()); - assertTrue("convertFinalAggFragment must be called for the Union subtree", convertor.finalAggCalled); + assertTrue("convertFragment (final-agg shape) must be called for the Union subtree", convertor.finalAggCalled); } /** @@ -748,15 +748,20 @@ private static class RecordingConvertor implements FragmentConvertor { RelNode reduceFragment; @Override - public byte[] convertShardScanFragment(String tableName, RelNode fragment) { - this.shardScanCalled = true; - this.shardScanTableName = tableName; - this.shardScanFragment = fragment; - return ("shard:" + tableName).getBytes(StandardCharsets.UTF_8); - } - - @Override - public byte[] convertFinalAggFragment(RelNode fragment) { + public byte[] convertFragment(RelNode fragment) { + // Distinguish shard-scan vs reduce/final by walking down the leftmost spine + // to find a TableScan-shaped leaf (annotations are stripped before this is + // called, so OpenSearchTableScan has been rewritten to LogicalTableScan). + org.apache.calcite.rel.core.TableScan scan = org.opensearch.analytics.planner.RelNodeUtils.findNode( + fragment, + org.apache.calcite.rel.core.TableScan.class + ); + if (scan != null) { + this.shardScanCalled = true; + this.shardScanTableName = scan.getTable().getQualifiedName().getLast(); + this.shardScanFragment = fragment; + return ("shard:" + this.shardScanTableName).getBytes(StandardCharsets.UTF_8); + } this.finalAggCalled = true; this.reduceFragment = fragment; return "reduce".getBytes(StandardCharsets.UTF_8); diff --git a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/sql/ValuesSqlIT.java b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/sql/ValuesSqlIT.java index b964296828e84..8adfe27f2d259 100644 --- a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/sql/ValuesSqlIT.java +++ b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/sql/ValuesSqlIT.java @@ -162,8 +162,7 @@ private void assertInnerJoinScanWithValues() { * {@code (2, total_size_of_GET_rows)}. * *

      1-shard only for now — the multi-shard PARTIAL→FINAL path hits a separate - * Int32/Int64 width drift in {@code DatafusionReduceSink.coerceToDeclaredSchema} - * that's being fixed in another change. + * Int32/Int64 width drift in the reduce sink that's being fixed in another change. */ public void testLiteralWithAggregateAndFilter_1shard() { createAndSeedHttpLogsIndex(1); diff --git a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/CoordinatorReduceStressIT.java b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/CoordinatorReduceStressIT.java index 19d9f05d82796..01b851bbb3f41 100644 --- a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/CoordinatorReduceStressIT.java +++ b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/CoordinatorReduceStressIT.java @@ -156,7 +156,7 @@ public void testReduceHandlesHundredBatches() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); @@ -194,7 +194,7 @@ public void testReduceReleasesOnStubFailure() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); @@ -274,7 +274,7 @@ public void testReduceCancelReleasesNative() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); @@ -375,7 +375,7 @@ public void testReduceCancelDuringParkedSendLeaks() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); @@ -470,7 +470,7 @@ public void testReduceConcurrentSinks() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); DatafusionReduceSink sink = new DatafusionReduceSink(ctx, runtimeHandle, drainExecutor); @@ -563,7 +563,7 @@ private void runOneSink(String queryId, Schema inputSchema, byte[] substrait, in 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); DatafusionReduceSink sink = new DatafusionReduceSink(ctx, runtimeHandle, drainExecutor); @@ -600,7 +600,7 @@ public void testMemtableHandlesBatchedFeed() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); @@ -639,7 +639,7 @@ public void testMemtableAllocatorReleaseUnderLoad() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema)), + List.of(new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait(INPUT_ID))), downstream ); DatafusionMemtableReduceSink sink = new DatafusionMemtableReduceSink(ctx, runtimeHandle); @@ -674,7 +674,10 @@ public void testMemtableRejectsMultiInputAndReleasesSession() { 0, buildSumSubstrait(), alloc, - List.of(new ExchangeSinkContext.ChildInput(0, schemaA), new ExchangeSinkContext.ChildInput(1, schemaB)), + List.of( + new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait("input-0")), + new ExchangeSinkContext.ChildInput(1, buildPassthroughSubstrait("input-1")) + ), new CapturingSink() ); @@ -707,7 +710,10 @@ public void testReduceMultiInputFanIn() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema), new ExchangeSinkContext.ChildInput(1, inputSchema)), + List.of( + new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait("input-0")), + new ExchangeSinkContext.ChildInput(1, buildPassthroughSubstrait("input-1")) + ), downstream ); @@ -763,7 +769,10 @@ public void testReduceMultiInputCancelMidFeed() throws Exception { 0, substrait, alloc, - List.of(new ExchangeSinkContext.ChildInput(0, inputSchema), new ExchangeSinkContext.ChildInput(1, inputSchema)), + List.of( + new ExchangeSinkContext.ChildInput(0, buildPassthroughSubstrait("input-0")), + new ExchangeSinkContext.ChildInput(1, buildPassthroughSubstrait("input-1")) + ), downstream ); @@ -817,6 +826,23 @@ public void testReduceMultiInputCancelMidFeed() throws Exception { // ── Helpers ────────────────────────────────────────────────────────────── + /** + * Bare {@code SELECT * FROM } substrait whose lowered output schema is the + * single BIGINT column {@code x} — used as the producer-side plan in + * {@link ExchangeSinkContext.ChildInput#producerPlanBytes()}. Tests across this file + * all use a single-column BIGINT input shape, so one builder serves every call site. + */ + private static byte[] buildPassthroughSubstrait(String inputId) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner hepPlanner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(hepPlanner, rexBuilder); + RelDataType bigintNullable = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true); + RelDataType rowType = typeFactory.builder().add("x", bigintNullable).build(); + RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), inputId, rowType); + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(scan); + } + /** * Builds Substrait bytes for {@code SELECT SUM(x) FROM "input-0"}. Same shape * as {@code DatafusionReduceSinkTests.buildSumSubstraitBytes}. @@ -831,7 +857,7 @@ private static byte[] buildSumSubstrait() { RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), INPUT_ID, rowType); AggregateCall sumCall = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, bigintNullable, "total"); LogicalAggregate agg = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sumCall)); - return new DataFusionFragmentConvertor(loadExtensions()).convertFinalAggFragment(agg); + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(agg); } /** @@ -851,7 +877,7 @@ private static byte[] buildMultiInputSumSubstrait(int childA, int childB) { LogicalUnion union = LogicalUnion.create(List.of(scanA, scanB), true); AggregateCall sumCall = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, bigintNullable, "total"); LogicalAggregate agg = LogicalAggregate.create(union, List.of(), ImmutableBitSet.of(), null, List.of(sumCall)); - return new DataFusionFragmentConvertor(loadExtensions()).convertFinalAggFragment(agg); + return new DataFusionFragmentConvertor(loadExtensions()).convertFragment(agg); } private static SimpleExtension.ExtensionCollection loadExtensions() { diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/JoinCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/JoinCommandIT.java index f042d8ccb2b68..99ed5f3389a0b 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/JoinCommandIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/JoinCommandIT.java @@ -61,19 +61,11 @@ private void ensureDataProvisioned() throws IOException { } // ── join (direct LogicalJoin) ────────────────────────────────────────────── - // - // NOTE on schema narrowing: the calcs dataset carries date/time/datetime - // fields that map to Calcite TIMESTAMP / DATE types. The analytics-engine - // Arrow schema converter (ArrowSchemaFromCalcite) currently rejects those - // types, so every query below projects down to int/string/boolean columns - // via an explicit {@code fields …} or an aggregation before the join output - // surfaces to Arrow. Removing the projection surfaces - // {@code IllegalArgumentException: Unsupported Calcite SQL type: TIMESTAMP}. /** * Inner equi-join across two indices of the calcs dataset, grouped on * {@code str0}. Both sides are pre-aggregated to a narrow keyword-only - * schema so the join output has no TIMESTAMP/DATE columns. + * schema so the join output is scalar-only. */ public void testInnerJoin() throws IOException { final String ppl = "source=" diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultisearchCommandIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultisearchCommandIT.java index 6434f17f220e4..16cf66ec28182 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultisearchCommandIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultisearchCommandIT.java @@ -60,10 +60,8 @@ public void testMultisearchTwoBranchesByCategory() throws IOException { // 6 null rows excluded by both predicates (5 + 6 + 6 = 17 total). // Verifies: Union over two same-schema projections + Aggregate(count by) on top — // the convertReduceFragment chain attachFragmentOnTop(Sort, - // attachFragmentOnTop(Aggregate, convertFinalAggFragment(Union))). - // Each branch projects to (int0, class) so the union row type is scalar-only — - // calcs has date/time/datetime columns whose TIMESTAMP Calcite SQL type - // ArrowSchemaFromCalcite doesn't yet handle (separate follow-up). + // attachFragmentOnTop(Aggregate, convertFragment(Union))). + // Each branch projects to (int0, class) so the union row type is scalar-only. assertRows( "| multisearch" + " [search source=" + DATASET.indexName + " | where int0 < 5 | eval class = \"low\" | fields int0, class]" @@ -83,8 +81,7 @@ public void testMultisearchThreeBranchesByStr0() throws IOException { // Pre-fix: 500 with "Names list ... 2 uses for {row-type-width} names". Post-fix: the // wrapper aggregate's [count, bucket] names propagate end-to-end, plan deserializes, // DataFusion executes the Union+Aggregate. - // Each branch projects to (str0, bucket) — see testMultisearchTwoBranchesByCategory's - // comment for the reason. + // Each branch projects to (str0, bucket) so the union row type is scalar-only. assertRows( "| multisearch" + " [search source=" + DATASET.indexName + " | where str0 = \"FURNITURE\" | eval bucket = \"F\" | fields str0, bucket]"