From 21b74e0b8948bd1a311a08e50aea72590e21675d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 2 Nov 2026 11:29:48 -0700 Subject: [PATCH 1/2] perf: send a native plan's output schema to the JVM only when it changes executePlan exported an ArrowSchema for every column of every output batch, and the JVM imported each one, although a plan's output schema rarely changes between batches. Native now records the data type and field metadata of each output column it last exported, and while a batch matches them it fills only the ArrowArray structs. It sets bit 32 of the row count it returns when it filled the schemas too. The JVM passes the same empty ArrowSchema structs with every batch, and imports a batch that arrives without a schema with the fields of the last one. A batch with a dictionary-encoded column always carries its schema, because the JVM gives each dictionary it imports a new id. The internal config spark.comet.exec.cacheOutputSchema.enabled, on by default, makes native send the schema with every batch when it is off. --- docs/source/contributor-guide/ffi.md | 51 ++-- .../contributor-guide/memory_management.md | 4 +- native/core/src/execution/jni_api.rs | 287 +++++++++++++++++- native/core/src/execution/spark_config.rs | 2 + native/core/src/execution/utils.rs | 89 ++++-- .../org/apache/arrow/c/ArrowImporter.java | 26 +- .../scala/org/apache/comet/CometConf.scala | 12 + .../org/apache/comet/CometExecIterator.scala | 3 +- .../main/scala/org/apache/comet/Native.scala | 4 +- .../org/apache/comet/vector/NativeUtil.scala | 208 ++++++++++--- .../apache/comet/vector/NativeUtilSuite.scala | 169 ++++++++++- 11 files changed, 740 insertions(+), 115 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 64db891347c..4794031e863 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -27,8 +27,8 @@ Comet transfers Arrow data across the JVM/native boundary in two directions: [Arrow C Stream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html). The JVM exports each per-partition iterator once as an `ArrowArrayStream`, and native pulls every batch through a single C callback. 2. **Native → JVM**: JVM pulls batches from native code using `CometExecIterator`, via the - [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) (one `ArrowArray`/`ArrowSchema` - pair per column of each batch). + [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) (one `ArrowArray` per column of + each batch, and one `ArrowSchema` per column whenever the output schema changes). The following diagram shows an example of the end-to-end flow for a query stage. @@ -196,8 +196,8 @@ When the JVM needs results from native execution: │ produces RecordBatch ▼ ┌─────────────────┐ -│ prepare_output │ ── fills one ArrowArray/ArrowSchema pair per column -│ (Rust/native) │ +│ export_batch │ ── fills one ArrowArray per column, and one ArrowSchema +│ (Rust/native) │ per column when the output schema changes └────────┬────────┘ │ Arrow C Data Interface │ (structs allocated by the JVM, filled by native) @@ -216,25 +216,37 @@ When the JVM needs results from native execution: ### Transfer Process -`CometExecIterator.hasNext` fetches each batch through `NativeUtil.getNextBatch`, with one JNI call per batch: +`CometExecIterator.hasNext` fetches each batch through `NativeUtil.getNextPlanBatch`, with one JNI call per batch: -1. `NativeUtil.getNextBatch` allocates one empty `ArrowArray`/`ArrowSchema` pair per output column from - `CometArrowAllocator` and passes their memory addresses to - `Native.executePlan(stage, partition, plan, arrayAddrs, schemaAddrs)`. +1. `getNextPlanBatch` allocates one empty `ArrowArray` per output column from `CometArrowAllocator`, and passes + their memory addresses to `Native.executePlan(stage, partition, plan, arrayAddrs, schemaAddrs)`, together with + those of one empty `ArrowSchema` per column, which it keeps across batches. 2. `executePlan` (in `jni_api.rs`) polls the native plan for its next `RecordBatch`, and `prepare_output` exports - each column into its pair with `move_to_spark` (in `execution/utils.rs`). `move_to_spark` zeroes the column's - offsets (see [Array Offsets](#array-offsets)), then writes an `FFI_ArrowArray` over the result, and an - `FFI_ArrowSchema` built from its data type and field metadata, into the JVM-allocated structs. `executePlan` - returns the batch's row count, or `-1` at the end of the output. With `spark.comet.debug.enabled` set, + each column into its `ArrowArray` with `move_array_to_spark` (in `execution/utils.rs`), which zeroes the + column's offsets (see [Array Offsets](#array-offsets)) and writes an `FFI_ArrowArray` over the result. When the + batch carries its schema, `move_to_spark` also writes an `FFI_ArrowSchema` built from the column's data type and + field metadata into its `ArrowSchema`. `executePlan` returns the batch's row count, with `OUTPUT_SCHEMA_EXPORTED` + (bit 32) set when the schema went too, or `-1` at the end of the output. With `spark.comet.debug.enabled` set, `prepare_output` first runs `validate_full` on every column. -3. At the end of the output, `NativeUtil` releases the unused structs. Otherwise `NativeUtil.importVector` imports - each column with `ArrowImporter.importVector`, wraps it with `CometVector.getVector`, and returns the vectors as - a `ColumnarBatch`. +3. At the end of the output, `NativeUtil` releases the unused structs. Otherwise `NativeUtil` imports each column + with `ArrowImporter.importVector`, from its schema when the schema came too and with the column's field from the + last schema otherwise, wraps it with `CometVector.getVector`, and returns the vectors as a `ColumnarBatch`. + +`NativeUtil.getNextBatch` does the same for decoded shuffle blocks, which always carry their schema, with new schema +structs for each block. `ArrowImporter` (in `spark/src/main/java/org/apache/arrow/c/`) imports every column through one shared `SchemaImporter`. Arrow Java's own `Data.importField` creates a new `SchemaImporter` for each field, and each one numbers dictionaries from 0, so two dictionary-encoded columns would collide in the shared `CDataDictionaryProvider`. +### Sending the Output Schema Once + +A batch carries its schema only when the schema changed. `OutputSchemaCache` (in `jni_api.rs`) records each output +column's data type and field metadata as it last exported them, and while the next batch matches, only the arrays +go. A batch with a dictionary-encoded column always carries its schema, because the JVM gives each dictionary it +imports a new id, and `NativeUtil` keeps no fields from such a batch. With the internal config +`spark.comet.exec.cacheOutputSchema.enabled` off, every batch carries its schema. + ### Array Offsets Arrow Java's C Data import ignores `ArrowArray.offset` at every level @@ -242,9 +254,8 @@ Arrow Java's C Data import ignores `ArrowArray.offset` at every level folds a slice into the buffers for almost every type, but a sliced `BooleanArray` keeps its bit offset, and a struct exports offset 0 even when its children are sliced. So every array native exports to the JVM first goes through `zero_offsets` (in `native/common/src/ffi_offsets.rs`), which re-slices boolean bitmaps to start at bit 0 at every -level and shares every other buffer. `move_to_spark` applies it to executed batches and decoded shuffle blocks, and -`JvmScalarUdfExpr` applies it to the inputs of the JVM UDF bridge. A new native to JVM export path has to call it too, -or sliced booleans reach the JVM misaligned +level and shares every other buffer. `move_array_to_spark`, which `move_to_spark` also goes through, applies it to +executed batches and decoded shuffle blocks, and `JvmScalarUdfExpr` applies it to the inputs of the JVM UDF bridge. A new native to JVM export path has to call it too, or sliced booleans reach the JVM misaligned ([#6288](https://github.com/apache/datafusion-comet/issues/6288)). `zero_offsets` leaves the offsets of a sliced string, binary, list or map array as they are. arrow-rs slices those @@ -264,8 +275,8 @@ Python runner rebases the same way before it sends a batch to the worker. All th Native allocates the data, and the JVM references it without copying: -- The `FFI_ArrowArray` that `move_to_spark` writes holds a reference to the column's native buffers, and its release - callback drops that reference. +- The `FFI_ArrowArray` that `move_to_spark` or `move_array_to_spark` writes holds a reference to the column's native + buffers, and its release callback drops that reference. - On import, Arrow Java wraps each native buffer in an `ArrowBuf`. Closing the last `ArrowBuf` over an imported column runs its release callback, and native frees the buffers once no Rust reference to them remains. - `CometExecIterator` closes the batch it returned when the consumer next calls `hasNext` or `next`, and on diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 4e7248f2ba3..b109cd0bb21 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -565,8 +565,8 @@ through a unified pool is charged to Spark. The same batch held by an operator t **Native → JVM (`CometExecIterator`).** DataFusion produces the batch in Rust, and whether it is still reserved when it reaches the boundary depends on the producing operator. The sort's output is wrapped in a `ReservationStream` that shrinks the reservation by the batch's size as each batch is -emitted, so by the time `prepare_output` sees a sorted batch its bytes are no longer reserved. -`prepare_output` then calls `move_to_spark`, which writes an `FFI_ArrowArray` whose release callback +emitted, so by the time `executePlan` exports a sorted batch its bytes are no longer reserved. +The export writes an `FFI_ArrowArray` whose release callback drops the Rust `ArrayData`; no reservation travels with it. The JVM wraps the pointers in `ArrowBuf`s and the memory is freed when the JVM calls `close()`. An exported batch is therefore usually _not_ pool-charged while the JVM holds it, but it is resident until the JVM releases it. A diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 50df7a786b0..bade7c7a759 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -111,6 +111,7 @@ use jni::{ use parking_lot::Mutex; use prost::Message; use std::collections::HashMap; +use std::ops::Deref; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; @@ -140,10 +141,11 @@ use crate::execution::tracing::{ use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; use crate::execution::spark_config::{ SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, - COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_SPILL_COMPRESSION_CODEC, - COMET_EXPLAIN_NATIVE_ENABLED, COMET_MAX_TEMP_DIRECTORY_SIZE, - COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, - COMET_STRICT_TIMESTAMP_LABELS, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, + COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED, + COMET_EXEC_SPILL_COMPRESSION_CODEC, COMET_EXPLAIN_NATIVE_ENABLED, + COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_PARQUET_PREFETCH_NEXT_FILE_METADATA_ENABLED, + COMET_PARQUET_ROW_FILTER_PUSHDOWN_ENABLED, COMET_STRICT_TIMESTAMP_LABELS, + COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, }; use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; use crate::parquet::parquet_exec::CometParquetScanOptions; @@ -571,6 +573,8 @@ struct ExecutionContext { pub explain_native: bool, /// Whether to log memory usage on each call to execute_plan pub tracing_enabled: bool, + /// What `executePlan` last exported of the output schema + output_schema: OutputSchemaCache, /// Rust thread ID, used for aggregating tracing metrics per thread pub rust_thread_id: u64, /// Pre-computed metric name for tracing memory usage @@ -727,6 +731,7 @@ struct PlanSettings { strict_timestamp_labels: bool, explain_native: bool, tracing_enabled: bool, + cache_output_schema: bool, max_temp_directory_size: u64, logging_memory_pool: bool, } @@ -748,6 +753,7 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult, ) -> CometResult { let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; - let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; - unsafe { export_batch(&array_addrs, &schema_addrs, output_batch, validate) } + let schema_addrs = || Ok(unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }); + match output_schema { + Some(cache) => unsafe { cache.export(&array_addrs, schema_addrs, output_batch, validate) }, + None => unsafe { + export_batch(&array_addrs, Some(&schema_addrs()?), output_batch, validate) + }, + } +} + +/// The data type and field metadata of an output column, which together determine the +/// `FFI_ArrowSchema` that `move_to_spark` exports for it. +struct ExportedColumn { + data_type: ArrowDataType, + metadata: HashMap, +} + +impl ExportedColumn { + fn of(batch: &RecordBatch) -> Vec { + batch + .columns() + .iter() + .zip(batch.schema_ref().fields()) + .map(|(column, field)| Self { + data_type: column.data_type().clone(), + metadata: field.metadata().clone(), + }) + .collect() + } + + /// Whether every column of `batch` would export the schema that `columns` were exported with. + fn all_match(columns: &[Self], batch: &RecordBatch) -> bool { + columns.len() == batch.num_columns() + && columns + .iter() + .zip(batch.columns()) + .zip(batch.schema_ref().fields()) + .all(|((exported, column), field)| { + exported.data_type == *column.data_type() + && exported.metadata == *field.metadata() + }) + } +} + +/// Set in what `executePlan` returns when it exported the batch's schema, one `FFI_ArrowSchema` +/// per column, for the JVM to import in place of the fields it holds. The row count is in the low +/// 32 bits. Matches `NativeUtil.OUTPUT_SCHEMA_EXPORTED`. +const OUTPUT_SCHEMA_EXPORTED: jlong = 1 << 32; + +/// What `executePlan` last exported of its output schema, so that the batches after it can leave +/// the schema out while it stays the same (`spark.comet.exec.cacheOutputSchema.enabled`). +struct OutputSchemaCache { + enabled: bool, + /// The output columns as last exported. `None` before the first batch, after a batch with a + /// dictionary, and while the cache is disabled. + exported: Option>, +} + +impl OutputSchemaCache { + fn new(enabled: bool) -> Self { + Self { + enabled, + exported: None, + } + } + + /// Exports `batch` like `export_batch`, but writes its schema only when it differs from the one + /// exported last, and records what it writes. `schema_addrs` gives the addresses of the schema + /// structs, and is called only when the schema goes too. The row count it returns has + /// `OUTPUT_SCHEMA_EXPORTED` set when it wrote the schema. + /// + /// A batch with a dictionary carries its schema every time, because the JVM gives each + /// dictionary it imports a new id and keeps its values under that id. + /// + /// # Safety + /// As for `export_batch`. + unsafe fn export>( + &mut self, + array_addrs: &[i64], + schema_addrs: impl FnOnce() -> CometResult, + batch: RecordBatch, + validate: bool, + ) -> CometResult { + // With no addresses, which Spark passes when the results of a query are not used, there is + // no schema to export either. + let unchanged = self + .exported + .as_deref() + .is_some_and(|columns| ExportedColumn::all_match(columns, &batch)); + if array_addrs.is_empty() || unchanged { + return export_batch(array_addrs, None, batch, validate); + } + + let schema_addrs = schema_addrs()?; + let cacheable = self.enabled + && !batch + .columns() + .iter() + .any(|column| has_dictionary(column.data_type())); + let columns = cacheable.then(|| ExportedColumn::of(&batch)); + let num_rows = export_batch(array_addrs, Some(&schema_addrs), batch, validate)?; + self.exported = columns; + Ok(num_rows | OUTPUT_SCHEMA_EXPORTED) + } +} + +/// Whether `data_type` is a dictionary or has one nested in it. +fn has_dictionary(data_type: &ArrowDataType) -> bool { + match data_type { + ArrowDataType::Dictionary(_, _) => true, + ArrowDataType::List(field) + | ArrowDataType::LargeList(field) + | ArrowDataType::ListView(field) + | ArrowDataType::LargeListView(field) + | ArrowDataType::FixedSizeList(field, _) + | ArrowDataType::Map(field, _) => has_dictionary(field.data_type()), + ArrowDataType::Struct(fields) => { + fields.iter().any(|field| has_dictionary(field.data_type())) + } + ArrowDataType::Union(fields, _) => fields + .iter() + .any(|(_, field)| has_dictionary(field.data_type())), + ArrowDataType::RunEndEncoded(_, values) => has_dictionary(values.data_type()), + _ => false, + } } /// Moves the columns of `output_batch` into the Arrow C Data Interface structs at `array_addrs` -/// and `schema_addrs`, one pair per column, and returns the row count. With no addresses, which +/// and `schema_addrs`, one pair per column, and returns the row count. Without `schema_addrs` +/// only the arrays are moved, for a JVM that already holds their schema. With no addresses, which /// Spark passes when the results of a query are not used, only the row count is returned. /// /// # Safety @@ -1219,7 +1351,7 @@ fn prepare_output( /// Whatever the structs hold is overwritten without being released. unsafe fn export_batch( array_addrs: &[i64], - schema_addrs: &[i64], + schema_addrs: Option<&[i64]>, output_batch: RecordBatch, validate: bool, ) -> CometResult { @@ -1252,11 +1384,13 @@ unsafe fn export_batch( let mut i = 0; while i < results.len() { let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?; - let field = output_schema.field(i); - - array_ref - .to_data() - .move_to_spark(field, array_addrs[i], schema_addrs[i])?; + let data = array_ref.to_data(); + match schema_addrs { + Some(schema_addrs) => { + data.move_to_spark(output_schema.field(i), array_addrs[i], schema_addrs[i])? + } + None => data.move_array_to_spark(array_addrs[i])?, + } i += 1; } } @@ -1489,6 +1623,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( schema_addrs, batch, exec_context.debug_native, + Some(&mut exec_context.output_schema), ), None => Ok(-1), } @@ -2276,7 +2411,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_takeShuffleBatch( let batch = unsafe { shuffle_batch_coalescer(handle)? } .take_ready() .ok_or_else(|| CometError::Internal("No decoded shuffle batch is ready".into()))?; - prepare_output(env, array_addrs, schema_addrs, batch, false) + prepare_output(env, array_addrs, schema_addrs, batch, false, None) }) }) } @@ -2435,7 +2570,7 @@ unsafe fn decode_shuffle_block( } else { read_ipc_compressed(block)? }; - export_batch(array_addrs, schema_addrs, batch, false) + export_batch(array_addrs, Some(schema_addrs), batch, false) } #[no_mangle] @@ -3132,6 +3267,126 @@ mod tests { ); } + /// `executePlan` skips the output schema while a batch would export the same one, which + /// depends on each column's type and field metadata but not on its name or nullability. + #[test] + fn exported_columns_match_only_the_same_types_and_metadata() { + let batch = |field: Field, column: ArrayRef| { + RecordBatch::try_new(Arc::new(Schema::new(vec![field])), vec![column]).unwrap() + }; + let ints = || Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef; + let exported = ExportedColumn::of(&batch(Field::new("a", DataType::Int32, true), ints())); + + let renamed = batch(Field::new("b", DataType::Int32, false), ints()); + assert!(ExportedColumn::all_match(&exported, &renamed)); + + let longs = Arc::new(arrow::array::Int64Array::from(vec![1, 2])) as ArrayRef; + let retyped = batch(Field::new("a", DataType::Int64, true), longs); + assert!(!ExportedColumn::all_match(&exported, &retyped)); + + let metadata = HashMap::from([("k".to_string(), "v".to_string())]); + let tagged = batch( + Field::new("a", DataType::Int32, true).with_metadata(metadata), + ints(), + ); + assert!(!ExportedColumn::all_match(&exported, &tagged)); + + let no_columns = RecordBatch::new_empty(Arc::new(Schema::empty())); + assert!(!ExportedColumn::all_match(&exported, &no_columns)); + } + + /// Exports `batch` through `cache` into new structs, as the JVM passes them, and returns what it + /// returned and whether it wrote each column's schema. The structs, and the buffers they hold, + /// are released when they drop at the end. + fn export_plan_output(batch: RecordBatch, cache: &mut OutputSchemaCache) -> (i64, Vec) { + use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; + let columns = batch.num_columns(); + let mut arrays: Vec<_> = (0..columns).map(|_| FFI_ArrowArray::empty()).collect(); + let mut schemas: Vec<_> = (0..columns).map(|_| FFI_ArrowSchema::empty()).collect(); + let array_addrs: Vec = arrays + .iter_mut() + .map(|array| array as *mut FFI_ArrowArray as i64) + .collect(); + let schema_addrs: Vec = schemas + .iter_mut() + .map(|schema| schema as *mut FFI_ArrowSchema as i64) + .collect(); + let schema_addrs = || Ok(schema_addrs.as_slice()); + let result = unsafe { cache.export(&array_addrs, schema_addrs, batch, false) }.unwrap(); + assert!(arrays.iter().all(|array| !array.is_released())); + let written = schemas + .iter() + .map(|schema| !schema.format.is_null()) + .collect(); + (result, written) + } + + #[test] + fn output_schema_cache_writes_the_schema_only_when_it_changes() { + let batch = |column: ArrayRef| RecordBatch::try_from_iter(vec![("c", column)]).unwrap(); + let ints = |values: Vec| batch(Arc::new(Int32Array::from(values))); + let mut cache = OutputSchemaCache::new(true); + + let first = export_plan_output(ints(vec![1, 2]), &mut cache); + assert_eq!(first, (2 | OUTPUT_SCHEMA_EXPORTED, vec![true])); + assert_eq!( + export_plan_output(ints(vec![3]), &mut cache), + (1, vec![false]) + ); + + // A new type is written, and is what the next batch is compared with. + let longs = batch(Arc::new(arrow::array::Int64Array::from(vec![4_i64]))); + let retyped = export_plan_output(longs.clone(), &mut cache); + assert_eq!(retyped, (1 | OUTPUT_SCHEMA_EXPORTED, vec![true])); + assert_eq!(export_plan_output(longs, &mut cache), (1, vec![false])); + + // A dictionary carries its schema every time. + let keys: arrow::array::DictionaryArray = + vec!["x", "y", "x"].into_iter().collect(); + let dictionaries = batch(Arc::new(keys)); + for _ in 0..2 { + let result = export_plan_output(dictionaries.clone(), &mut cache); + assert_eq!(result, (3 | OUTPUT_SCHEMA_EXPORTED, vec![true])); + } + assert!(cache.exported.is_none()); + + // With no addresses there is nothing to write, so nothing is flagged. + let no_schema_addrs = || -> CometResult<&[i64]> { panic!("not needed") }; + let unused = unsafe { cache.export(&[], no_schema_addrs, ints(vec![5, 6, 7]), false) }; + assert_eq!(unused.unwrap(), 3); + } + + #[test] + fn a_disabled_output_schema_cache_writes_the_schema_every_time() { + let ints = RecordBatch::try_from_iter(vec![( + "c", + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + )]) + .unwrap(); + let mut cache = OutputSchemaCache::new(false); + for _ in 0..2 { + let result = export_plan_output(ints.clone(), &mut cache); + assert_eq!(result, (2 | OUTPUT_SCHEMA_EXPORTED, vec![true])); + } + assert!(cache.exported.is_none()); + } + + #[test] + fn has_dictionary_finds_nested_dictionaries() { + let dictionary = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let list_of = |data_type| DataType::List(Arc::new(Field::new("item", data_type, true))); + assert!(has_dictionary(&dictionary)); + assert!(has_dictionary(&list_of(dictionary.clone()))); + assert!(has_dictionary(&DataType::Struct( + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", dictionary, true) + ] + .into() + ))); + assert!(!has_dictionary(&list_of(DataType::Utf8))); + } + #[test] fn length_resolves_to_spark_length_for_string_and_binary() { use datafusion::physical_expr::expressions::{CastExpr, Column}; diff --git a/native/core/src/execution/spark_config.rs b/native/core/src/execution/spark_config.rs index 511b5733388..3521a5d8d84 100644 --- a/native/core/src/execution/spark_config.rs +++ b/native/core/src/execution/spark_config.rs @@ -30,6 +30,8 @@ pub(crate) const COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED: &str = "spark.comet.exec.aggregate.skipPartial.enabled"; pub(crate) const COMET_EXEC_SPILL_COMPRESSION_CODEC: &str = "spark.comet.exec.spill.compression.codec"; +pub(crate) const COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED: &str = + "spark.comet.exec.cacheOutputSchema.enabled"; pub(crate) const COMET_NATIVE_UDF_ENABLED: &str = "spark.comet.nativeUdf.enabled"; pub(crate) const COMET_NATIVE_UDF_ALLOWED_PATHS: &str = "spark.comet.nativeUdf.allowedPaths"; pub(crate) const COMET_STRICT_TIMESTAMP_LABELS: &str = diff --git a/native/core/src/execution/utils.rs b/native/core/src/execution/utils.rs index 527f276b25d..bb56e7bf6d0 100644 --- a/native/core/src/execution/utils.rs +++ b/native/core/src/execution/utils.rs @@ -27,52 +27,47 @@ use datafusion_comet_common::zero_offsets; pub trait SparkArrowConvert { /// Move Arrow Arrays to C data interface. fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError>; + + /// Move the array, but not its schema, to C data interface, for a JVM that already imported + /// the schema `move_to_spark` exported for an array of the same type and field metadata. + fn move_array_to_spark(&self, array: i64) -> Result<(), ExecutionError>; } impl SparkArrowConvert for ArrayData { /// Move this ArrowData to pointers of Arrow C data interface. fn move_to_spark(&self, field: &Field, array: i64, schema: i64) -> Result<(), ExecutionError> { - let array_ptr = array as *mut FFI_ArrowArray; - let schema_ptr = schema as *mut FFI_ArrowSchema; - - let array_align = std::mem::align_of::(); - let schema_align = std::mem::align_of::(); - // Arrow Java ignores `ArrowArray.offset` on import, so every level has to start at 0. - let data = zero_offsets(self)?; - let ffi_array = FFI_ArrowArray::new(&data); // Spark owns the top-level name and nullability. Preserve the existing anonymous schema // shape while carrying logical extension metadata from the RecordBatch field. let ffi_schema = FFI_ArrowSchema::try_from(self.data_type())?.with_metadata(field.metadata())?; + self.move_array_to_spark(array)?; + // SAFETY: the caller passes the address of an `FFI_ArrowSchema` it owns. + unsafe { write_ffi_struct(schema as *mut FFI_ArrowSchema, ffi_schema) }; + Ok(()) + } - // Check if the pointer alignment is correct. - if array_ptr.align_offset(array_align) != 0 || schema_ptr.align_offset(schema_align) != 0 { - unsafe { - std::ptr::write_unaligned(array_ptr, ffi_array); - std::ptr::write_unaligned(schema_ptr, ffi_schema); - } - } else { - // SAFETY: `array_ptr` and `schema_ptr` are aligned correctly. - debug_assert_eq!( - array_ptr.align_offset(array_align), - 0, - "move_to_spark: array_ptr not aligned" - ); - debug_assert_eq!( - schema_ptr.align_offset(schema_align), - 0, - "move_to_spark: schema_ptr not aligned" - ); - unsafe { - std::ptr::write(array_ptr, ffi_array); - std::ptr::write(schema_ptr, ffi_schema); - } - } - + fn move_array_to_spark(&self, array: i64) -> Result<(), ExecutionError> { + // Arrow Java ignores `ArrowArray.offset` on import, so every level has to start at 0. + let data = zero_offsets(self)?; + // SAFETY: the caller passes the address of an `FFI_ArrowArray` it owns. + unsafe { write_ffi_struct(array as *mut FFI_ArrowArray, FFI_ArrowArray::new(&data)) }; Ok(()) } } +/// Writes `value` into the JVM-allocated struct at `ptr`, which Arrow Java may not have aligned +/// for `T`. Whatever the struct held is overwritten without being released. +/// +/// # Safety +/// `ptr` must be valid for writing a `T`. +unsafe fn write_ffi_struct(ptr: *mut T, value: T) { + if ptr.align_offset(std::mem::align_of::()) != 0 { + std::ptr::write_unaligned(ptr, value) + } else { + std::ptr::write(ptr, value) + } +} + #[cfg(test)] mod tests { use super::*; @@ -145,4 +140,34 @@ mod tests { let imported = unsafe { from_ffi(ffi_array, &ffi_schema) }.unwrap(); assert_eq!(imported, data); } + + /// An array moved without its schema reads back through the schema `move_to_spark` exported + /// for an earlier array of the same type, as the JVM reads it. + #[test] + fn test_move_array_to_spark_reads_with_an_earlier_schema() { + let field = Field::new("v", DataType::Int32, true); + let first = Int32Array::from(vec![Some(1), None]).into_data(); + let mut first_array = MaybeUninit::::uninit(); + let mut schema = MaybeUninit::::uninit(); + first + .move_to_spark( + &field, + first_array.as_mut_ptr() as i64, + schema.as_mut_ptr() as i64, + ) + .unwrap(); + drop(unsafe { first_array.assume_init() }); + let schema = unsafe { schema.assume_init() }; + + // Sliced, so the export also has to fold the offset into the buffers. + let second = Int32Array::from(vec![Some(7), Some(8), None, Some(9)]).slice(1, 3); + let mut second_array = MaybeUninit::::uninit(); + second + .to_data() + .move_array_to_spark(second_array.as_mut_ptr() as i64) + .unwrap(); + + let imported = unsafe { from_ffi(second_array.assume_init(), &schema) }.unwrap(); + assert_eq!(Int32Array::from(imported), second); + } } diff --git a/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java b/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java index 94c4916f9f1..116b6735e57 100644 --- a/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java +++ b/spark/src/main/java/org/apache/arrow/c/ArrowImporter.java @@ -53,10 +53,28 @@ Field importField(ArrowSchema schema, CDataDictionaryProvider provider) { public FieldVector importVector( ArrowArray array, ArrowSchema schema, CDataDictionaryProvider provider) { - Field field = null; - FieldVector vector = null; + Field field; try { field = importField(schema, provider); + } catch (RuntimeException | Error failure) { + // No importer has taken the array yet. + if (!array.isClosed()) { + AutoCloseables.close(failure, array::release, array); + } + throw failure; + } + return importVector(array, field, provider); + } + + /** + * Imports {@code array} into a new vector of {@code field}, which this importer imported from the + * schema of {@code array}, or of an earlier array of the same type so that the schema need not + * cross the boundary again. A field with a dictionary serves one array only, because the values + * of each imported dictionary are kept under the id its schema import assigned. + */ + public FieldVector importVector(ArrowArray array, Field field, CDataDictionaryProvider provider) { + FieldVector vector = null; + try { vector = field.createVector(allocator); ArrayImporter importer = new ArrayImporter(allocator, vector, provider); importer.importArray(array); @@ -65,9 +83,7 @@ public FieldVector importVector( if (vector != null) { AutoCloseables.close(failure, vector); } - if (field != null) { - closeDictionaries(field, provider, failure); - } + closeDictionaries(field, provider, failure); // ArrayImporter.importArray moves the array into a copy it owns and closes the source // before doImport can fail, so a closed array means the importer already took it and // releasing it again would be a double release. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..d5162c566d2 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -717,6 +717,18 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + // Used on native side. Check spark_config.rs how the config is used + val COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.cacheOutputSchema.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether a native plan sends the JVM the Arrow schema of its output batches only when " + + "it changes, rather than with every batch. Batches with a dictionary-encoded column " + + "always carry their schema. This is an internal config for benchmarking.") + .internal() + .booleanConf + .createWithDefault(true) + val COMET_EXEC_JOIN_DYNAMIC_FILTER_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.join.dynamicFilter.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index d209e0491d4..2d90baf566c 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -261,7 +261,7 @@ class CometExecIterator( val result = withTrace( s"getNextBatch[JVM] stage=${ctx.stageId()}", tracingEnabled, { - nativeUtil.getNextBatch( + nativeUtil.getNextPlanBatch( numOutputCols, (arrayAddrs, schemaAddrs) => { nativeLib.executePlan(ctx.stageId(), partitionIndex, plan, arrayAddrs, schemaAddrs) @@ -916,6 +916,7 @@ object CometExecIterator extends Logging { CometConf.COMET_DEBUG_ENABLED, CometConf.COMET_DEBUG_MEMORY_ENABLED, CometConf.COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED, + CometConf.COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED, CometConf.COMET_EXEC_SPILL_COMPRESSION_CODEC, CometConf.COMET_EXPLAIN_NATIVE_ENABLED, CometConf.COMET_MAX_TEMP_DIRECTORY_SIZE, diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index c24bd762c3e..14a4c7ed9f3 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -127,7 +127,9 @@ class Native extends NativeBase { * @param schemaAddrs * the addresses of Arrow Schema structures * @return - * the number of rows, if -1, it means end of the output. + * the number of rows, if -1, it means end of the output. Native fills the schema structures + * only when the output schema changes, unless `spark.comet.exec.cacheOutputSchema.enabled` is + * off, and sets `NativeUtil.OUTPUT_SCHEMA_EXPORTED` when it does. */ @native def executePlan( stage: Int, diff --git a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala index 8dd49d8fb1f..aa075512eb9 100644 --- a/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala +++ b/spark/src/main/scala/org/apache/comet/vector/NativeUtil.scala @@ -19,12 +19,14 @@ package org.apache.comet.vector +import java.util.concurrent.atomic.LongAdder + import scala.collection.mutable import scala.jdk.CollectionConverters._ import org.apache.arrow.c.{ArrowArray, ArrowImporter, ArrowSchema, CDataDictionaryProvider, Data} import org.apache.arrow.util.AutoCloseables -import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.{FieldVector, VectorSchemaRoot} import org.apache.arrow.vector.dictionary.DictionaryProvider import org.apache.arrow.vector.types.pojo.{ArrowType, Field} import org.apache.spark.SparkException @@ -32,7 +34,7 @@ import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, ConstantCol import org.apache.spark.sql.comet.util.Utils import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType, TimestampNTZType, TimestampType} -import org.apache.spark.sql.vectorized.ColumnarBatch +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.{CometArrowAllocator, CometArrowImportAllocator, CometRuntimeException} @@ -62,6 +64,23 @@ class NativeUtil extends AutoCloseable { */ private val dictionaryProvider: CDataDictionaryProvider = new CDataDictionaryProvider + /** + * The fields of a native plan's output columns, as [[getNextPlanBatch]] last imported them. + * Native sends the output schema only when it changes, and the batches in between are imported + * with these. `null` until the first batch. + */ + private var planOutputFields: Array[Field] = _ + + /** + * The schema structs [[getNextPlanBatch]] passes native with every batch. Native fills them + * only when the output schema changes, so they stay allocated and empty across the batches in + * between. Importing them consumes them, and the next call allocates new ones. + */ + private var planOutputSchemas: Array[ArrowSchema] = _ + + /** The addresses of [[planOutputSchemas]]. */ + private var planOutputSchemaAddrs: Array[Long] = _ + /** * Allocates Arrow structs for the given number of columns. * @@ -71,21 +90,14 @@ class NativeUtil extends AutoCloseable { * a pair of Arrow arrays and Arrow schemas */ def allocateArrowStructs(numCols: Int): (Array[ArrowArray], Array[ArrowSchema]) = { - val arrays = new Array[ArrowArray](numCols) - val schemas = new Array[ArrowSchema](numCols) - - try { - (0 until numCols).foreach { index => - // Publish each allocation before the next one, so partial allocation failures can free it. - schemas(index) = ArrowSchema.allocateNew(allocator) - arrays(index) = ArrowArray.allocateNew(allocator) + val arrays = allocateArrowArrays(numCols) + val schemas = + try allocateArrowSchemas(numCols) + catch { + case failure: Throwable => + releaseArrowStructs(arrays, Array.empty, failure) + throw failure } - } catch { - case failure: Throwable => - releaseArrowStructs(arrays, schemas, failure) - throw failure - } - (arrays, schemas) } @@ -216,6 +228,86 @@ class NativeUtil extends AutoCloseable { } } + /** + * Gets the next batch of a native plan, which sends its output schema only with the batches + * where it changes, and otherwise works like [[getNextBatch]]. + * + * Native sets [[NativeUtil.OUTPUT_SCHEMA_EXPORTED]] in what `func` returns when it filled the + * schema structs too, as it does for every batch with + * `spark.comet.exec.cacheOutputSchema.enabled` off. The batch is then imported as + * [[getNextBatch]] imports every batch, and its fields are kept for the batches that arrive + * without a schema. + */ + def getNextPlanBatch( + numOutputCols: Int, + func: (Array[Long], Array[Long]) => Long): Option[ColumnarBatch] = { + if (planOutputSchemas == null) { + planOutputSchemas = allocateArrowSchemas(numOutputCols) + planOutputSchemaAddrs = planOutputSchemas.map(_.memoryAddress()) + } + val schemas = planOutputSchemas + val arrays = allocateArrowArrays(numOutputCols) + + val result = + try { + func(arrays.map(_.memoryAddress()), planOutputSchemaAddrs) + } catch { + case failure: Throwable => + // As in getNextBatch, native may have filled some of the structs before failing. + planOutputSchemas = null + releaseArrowStructs(arrays, schemas, failure) + throw failure + } + + result match { + case -1 => + // The schema structs stay allocated for close() to free. + releaseArrowStructs(arrays, Array.empty) + None + case packed => + val numRows = (packed & 0xffffffffL).toInt + val cometVectors = if ((packed & NativeUtil.OUTPUT_SCHEMA_EXPORTED) != 0) { + // The import consumes the schema structs, so the next batch gets new ones. + planOutputSchemas = null + val imported = importVector(arrays, schemas) + val fields = imported.map(_.getValueVector.getField).toArray + // A field with a dictionary is good for its own batch only, and native sends the schema + // of every batch that has one. + planOutputFields = if (fields.exists(NativeUtil.hasDictionary)) null else fields + NativeUtil.planOutputSchemaImports.increment() + imported.toArray + } else { + importVector(arrays, planOutputFields) + } + Some(new ColumnarBatch(cometVectors.toArray[ColumnVector], numRows)) + } + } + + protected def allocateArrowArrays(numCols: Int): Array[ArrowArray] = { + val arrays = new Array[ArrowArray](numCols) + try { + // Publish each allocation before the next one, so partial allocation failures can free it. + (0 until numCols).foreach(index => arrays(index) = ArrowArray.allocateNew(allocator)) + } catch { + case failure: Throwable => + releaseArrowStructs(arrays, Array.empty, failure) + throw failure + } + arrays + } + + protected def allocateArrowSchemas(numCols: Int): Array[ArrowSchema] = { + val schemas = new Array[ArrowSchema](numCols) + try { + (0 until numCols).foreach(index => schemas(index) = ArrowSchema.allocateNew(allocator)) + } catch { + case failure: Throwable => + releaseArrowStructs(Array.empty, schemas, failure) + throw failure + } + schemas + } + private def releaseArrowStructs( arrays: Array[ArrowArray], schemas: Array[ArrowSchema], @@ -259,20 +351,45 @@ class NativeUtil extends AutoCloseable { * @return * a list of Comet vectors */ - def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = { - val arrayVectors = mutable.ArrayBuffer.empty[CometVector] - var firstUnconsumed = 0 + def importVector(arrays: Array[ArrowArray], schemas: Array[ArrowSchema]): Seq[CometVector] = + // importField's finally consumes the schema. ArrayImporter takes the array normally, while + // ArrowImporter's catch releases it if the import fails before that transfer. + importColumns(arrays, schemas)(i => + importer.importVector(arrays(i), schemas(i), dictionaryProvider)).toSeq + /** + * Imports arrays from native execution with the fields of an earlier import of the same schema, + * and returns a list of Comet vectors. On failure it releases everything it was given, as the + * import with schemas does. + */ + private def importVector( + arrays: Array[ArrowArray], + fields: Array[Field]): Array[CometVector] = { + if (arrays.nonEmpty && (fields == null || fields.length != arrays.length)) { + val failure = new IllegalStateException( + s"Native plan sent ${arrays.length} output columns without a schema, but the last " + + s"schema it sent had ${Option(fields).fold(0)(_.length)}") + releaseArrowStructs(arrays, Array.empty, failure) + throw failure + } + importColumns(arrays, Array.empty)(i => + importer.importVector(arrays(i), fields(i), dictionaryProvider)) + } + + /** + * Imports each of `arrays` with `importColumn` and wraps it as a Comet vector. `importColumn` + * consumes its column's array, and schema if any, even when it fails. On failure this releases + * the rest, the arrays and `schemas` not consumed yet and the vectors already imported, so + * callers must not add cleanup of their own. + */ + private def importColumns(arrays: Array[ArrowArray], schemas: Array[ArrowSchema])( + importColumn: Int => FieldVector): Array[CometVector] = { + val vectors = new Array[CometVector](arrays.length) + var i = 0 try { - (0 until arrays.length).foreach { i => - val arrowSchema = schemas(i) - val arrowArray = arrays(i) - - // importField's finally consumes the schema. ArrayImporter takes the array normally, while - // ArrowImporter's catch releases it if the import fails before that transfer. - firstUnconsumed = i + 1 - val arrowVector = importer.importVector(arrowArray, arrowSchema, dictionaryProvider) - val cometVector = + while (i < arrays.length) { + val arrowVector = importColumn(i) + vectors(i) = try CometVector.getVector(arrowVector, dictionaryProvider) catch { case failure: Throwable => @@ -284,13 +401,13 @@ class NativeUtil extends AutoCloseable { ArrowImporter.closeDictionaries(field, dictionaryProvider, failure) throw failure } - arrayVectors += cometVector + i += 1 } - arrayVectors.toSeq + vectors } catch { case failure: Throwable => - AutoCloseables.close(failure, arrayVectors.toSeq: _*) - releaseArrowStructs(arrays.drop(firstUnconsumed), schemas.drop(firstUnconsumed), failure) + AutoCloseables.close(failure, vectors.take(i): _*) + releaseArrowStructs(arrays.drop(i + 1), schemas.drop(i + 1), failure) throw failure } } @@ -326,13 +443,36 @@ class NativeUtil extends AutoCloseable { } override def close(): Unit = { - // closing the dictionary provider also closes the dictionary arrays - dictionaryProvider.close() + val schemas = planOutputSchemas + planOutputSchemas = null + try { + if (schemas != null) releaseArrowStructs(Array.empty, schemas) + } finally { + // closing the dictionary provider also closes the dictionary arrays + dictionaryProvider.close() + } } } object NativeUtil { + /** + * Set in what `Native.executePlan` returns when native filled the schema structs as well as the + * arrays. The row count is in the low 32 bits. Matches `OUTPUT_SCHEMA_EXPORTED` in + * `jni_api.rs`. + */ + val OUTPUT_SCHEMA_EXPORTED: Long = 1L << 32 + + /** + * How many times [[NativeUtil.getNextPlanBatch]] has imported an output schema in this JVM, so + * that tests and benchmarks can check that a plan sends it only when it changes. + */ + private[apache] val planOutputSchemaImports = new LongAdder() + + /** Whether `field` or a field nested in it is dictionary-encoded. */ + private def hasDictionary(field: Field): Boolean = + field.getDictionary != null || field.getChildren.asScala.exists(hasDictionary) + /** * Describes the first timestamp in `field` whose timezone label is not Comet's convention for * the Spark type at the same position in `dataType`, as a path rooted at `path`. That is diff --git a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala index 50eb2959cf1..6d31ba66de4 100644 --- a/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala +++ b/spark/src/test/scala/org/apache/comet/vector/NativeUtilSuite.scala @@ -22,6 +22,7 @@ package org.apache.comet.vector import java.io.IOException import java.nio.charset.StandardCharsets +import scala.collection.mutable import scala.jdk.CollectionConverters._ import scala.util.Using @@ -34,8 +35,9 @@ import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvid import org.apache.arrow.vector.types.TimeUnit import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} import org.apache.spark.sql.CometTestBase -import org.apache.spark.sql.comet.{CapturedSQLConf, CometExec} +import org.apache.spark.sql.comet.{CapturedSQLConf, CometExec, CometNativeExec} import org.apache.spark.sql.comet.util.Utils +import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, IntegerType, MapType, StringType, StructField, StructType, TimestampNTZType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} @@ -50,11 +52,13 @@ class NativeUtilSuite extends CometTestBase { val allocator = new RootAllocator(Long.MaxValue) var allocatedArrays = Array.empty[ArrowArray] val nativeUtil = new NativeUtil { - override def allocateArrowStructs(numCols: Int): (Array[ArrowArray], Array[ArrowSchema]) = { + override protected def allocateArrowArrays(numCols: Int): Array[ArrowArray] = { allocatedArrays = Array.fill(numCols)(ArrowArray.allocateNew(allocator)) - val schemas = Array.fill(numCols)(ArrowSchema.allocateNew(allocator)) - (allocatedArrays, schemas) + allocatedArrays } + + override protected def allocateArrowSchemas(numCols: Int): Array[ArrowSchema] = + Array.fill(numCols)(ArrowSchema.allocateNew(allocator)) } try check(nativeUtil, allocator, () => allocatedArrays) finally { @@ -312,6 +316,163 @@ class NativeUtilSuite extends CometTestBase { } } + test("getNextPlanBatch imports batches sent without a schema with the fields sent last") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + val schemaAddrs = mutable.ArrayBuffer.empty[Seq[Long]] + def nextBatch(values: Seq[Int], withSchema: Boolean): ColumnarBatch = + Using.resource(new IntVector("value", allocator)) { vector => + vector.allocateNew(values.length) + values.indices.foreach(i => vector.setSafe(i, values(i))) + vector.setValueCount(values.length) + nativeUtil + .getNextPlanBatch( + 1, + (arrays, schemas) => { + schemaAddrs += schemas.toSeq + val schema = if (withSchema) ArrowSchema.wrap(schemas(0)) else null + Data.exportVector(allocator, vector, null, ArrowArray.wrap(arrays(0)), schema) + val rows = values.length.toLong + if (withSchema) rows | NativeUtil.OUTPUT_SCHEMA_EXPORTED else rows + }) + .get + } + + val imports = NativeUtil.planOutputSchemaImports.sum() + val batches = Seq( + nextBatch(Seq(1, 2), withSchema = true), + nextBatch(Seq(3), withSchema = false), + nextBatch(Seq(4, 5, 6), withSchema = false)) + try { + val values = batches.map(batch => (0 until batch.numRows()).map(batch.column(0).getInt)) + assert(values == Seq(Seq(1, 2), Seq(3), Seq(4, 5, 6))) + assert(NativeUtil.planOutputSchemaImports.sum() - imports == 1) + // The batches without a schema were passed the same empty schema structs. + assert(schemaAddrs(1) == schemaAddrs(2)) + } finally { + batches.foreach(_.close()) + } + // Only the kept schema structs remain, until close() frees them. + nativeUtil.close() + assert(allocator.getAllocatedMemory == 0) + } + } + + test("getNextPlanBatch releases the structs native filled before failing") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + Using.resource(new IntVector("value", allocator)) { vector => + vector.allocateNew(1) + vector.setSafe(0, 42) + vector.setValueCount(1) + val expected = new IOException("native failed after exporting one column") + val actual = intercept[IOException] { + nativeUtil.getNextPlanBatch( + 2, + (arrays, schemas) => { + Data.exportVector( + allocator, + vector, + null, + ArrowArray.wrap(arrays(0)), + ArrowSchema.wrap(schemas(0))) + throw expected + }) + } + assert(actual eq expected) + } + assert(allocator.getAllocatedMemory == 0) + } + } + + test("getNextPlanBatch keeps its schema structs past the end of the output for close()") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + assert(nativeUtil.getNextPlanBatch(2, (_, _) => -1L).isEmpty) + assert(allocator.getAllocatedMemory > 0) + nativeUtil.close() + assert(allocator.getAllocatedMemory == 0) + } + } + + test("getNextPlanBatch rejects a batch without a schema before any schema was sent") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + Using.resource(new IntVector("value", allocator)) { vector => + vector.allocateNew(1) + vector.setSafe(0, 42) + vector.setValueCount(1) + val failure = intercept[IllegalStateException] { + nativeUtil.getNextPlanBatch( + 1, + (arrays, _) => { + Data.exportVector(allocator, vector, null, ArrowArray.wrap(arrays(0))) + 1L + }) + } + assert(failure.getMessage.contains("without a schema")) + } + nativeUtil.close() + assert(allocator.getAllocatedMemory == 0) + } + } + + test("a native plan sends its output schema with its first batch only") { + withParquetTable((0 until 1000).map(i => (i.toLong, i.toString, (i, s"v$i"))), "t") { + Seq(true, false).foreach { cacheOutputSchema => + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "100", + CometConf.COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED.key -> cacheOutputSchema.toString) { + val df = sql("SELECT _1, _2, _3 FROM t") + val plan = collectFirst(df.queryExecution.executedPlan) { + case native: CometNativeExec => + native + }.get + val imports = NativeUtil.planOutputSchemaImports.sum() + val batchesPerTask = SQLExecution.withNewExecutionId(df.queryExecution) { + plan.executeColumnar().mapPartitions(batches => Iterator(batches.size)).collect() + } + assert(batchesPerTask.sum >= 10) + val expected = + if (cacheOutputSchema) batchesPerTask.count(_ > 0) else batchesPerTask.sum + assert(NativeUtil.planOutputSchemaImports.sum() - imports == expected) + checkSparkAnswerAndOperator(df) + } + } + } + } + + test("getNextPlanBatch keeps no fields from a batch with a dictionary") { + withIsolatedStructAllocator { (nativeUtil, allocator, _) => + val indexType = new ArrowType.Int(32, true) + val encoding = new DictionaryEncoding(0L, false, indexType) + val indexField = new Field("encoded", new FieldType(true, indexType, encoding), null) + Using.resources( + new VarCharVector("dictionary", allocator), + new IntVector(indexField, allocator)) { (values, indices) => + values.allocateNew(1) + values.setSafe(0, "a".getBytes(StandardCharsets.UTF_8)) + values.setValueCount(1) + indices.allocateNew(2) + indices.setSafe(0, 0) + indices.setSafe(1, 0) + indices.setValueCount(2) + val provider = new MapDictionaryProvider(new Dictionary(values, encoding)) + def nextBatch(withSchema: Boolean): Option[ColumnarBatch] = + nativeUtil.getNextPlanBatch( + 1, + (arrays, schemas) => { + val schema = if (withSchema) ArrowSchema.wrap(schemas(0)) else null + Data.exportVector(allocator, indices, provider, ArrowArray.wrap(arrays(0)), schema) + if (withSchema) 2L | NativeUtil.OUTPUT_SCHEMA_EXPORTED else 2L + }) + + nextBatch(withSchema = true).get.close() + // Native sends the schema of every batch with a dictionary, so one without is an error. + val failure = intercept[IllegalStateException](nextBatch(withSchema = false)) + assert(failure.getMessage.contains("without a schema")) + } + nativeUtil.close() + assert(allocator.getAllocatedMemory == 0) + } + } + test("exportBatch round-trips a ConstantColumnVector through Arrow FFI") { // Smoke test for the ConstantColumnVector arm of NativeUtil.exportBatch: a batch carrying // Spark ConstantColumnVectors (partition values / per-batch constants) is exported across the From 566fc872ce6b1446a564757588d64bbfc9973e6c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 2 Nov 2026 11:29:48 -0700 Subject: [PATCH 2/2] bench: add a benchmark of moving native plan output to the JVM CometPlanOutputBenchmark compares spark.comet.exec.cacheOutputSchema.enabled on and off for native plans with 1, 10 and 100 output columns of bigint, string and struct, and for one that reads 100 distinct columns. Besides wall time, which leaves planning out, it reports each arm's process CPU time per batch, since the JVM's import overlaps the native plan when that runs on a Tokio worker. --- .../benchmark/CometPlanOutputBenchmark.scala | 213 ++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometPlanOutputBenchmark.scala diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPlanOutputBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPlanOutputBenchmark.scala new file mode 100644 index 00000000000..5a39655c6f6 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometPlanOutputBenchmark.scala @@ -0,0 +1,213 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.benchmark + +import java.io.File +import java.lang.management.ManagementFactory + +import scala.collection.mutable + +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.comet.CometNativeExec +import org.apache.spark.sql.execution.{QueryExecution, SQLExecution} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf +import org.apache.comet.vector.NativeUtil + +/** + * Benchmark to measure what it costs to move the output batches of a native plan to the JVM, with + * the Arrow schema of the output sent with every batch or only when it changes + * (`spark.comet.exec.cacheOutputSchema.enabled`). + * + * Most cases read one column of a Parquet table and project it `n` times, so the native plan does + * almost nothing for each of its `n` output columns beyond exporting it. The last case reads `n` + * distinct columns, so the work of the scan grows with `n` as well. The JVM only counts the rows + * of the batches it receives, as a JVM operator that passes batches through would, so these + * results bound the effect on a query from above. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make \ + * benchmark-org.apache.spark.sql.benchmark.CometPlanOutputBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometPlanOutputBenchmark-**results.txt". + */ +object CometPlanOutputBenchmark extends CometBenchmarkBase { + + private val narrowRows = 8 * 1024 * 1024 + private val wideRows = 1024 * 1024 + private val wideColumns = 100 + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + // Keep the plan a scan and a projection under the transition to rows. + spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + withTempPath { dir => + withTempTable("narrow", "wide") { + writeTable( + dir, + "narrow", + spark.sql( + "SELECT id AS i, CAST(id AS STRING) AS s, " + + s"named_struct('a', id, 'b', CAST(id AS STRING)) AS st FROM range($narrowRows)")) + writeTable( + dir, + "wide", + spark.sql( + (0 until wideColumns) + .map(c => s"id + $c AS w$c") + .mkString("SELECT ", ", ", s" FROM range($wideRows)"))) + + for { + (column, typeName) <- Seq( + "i" -> "bigint", + "s" -> "string", + "st" -> "struct") + n <- Seq(1, 10, 100) + } { + val query = + (0 until n).map(c => s"$column AS c$c").mkString("SELECT ", ", ", " FROM narrow") + runCases(s"$n $typeName output columns, 1 column read", narrowRows, query) + } + runCases( + s"$wideColumns bigint output columns, $wideColumns columns read", + wideRows, + "SELECT * FROM wide") + } + } + } + + private def writeTable(dir: File, name: String, df: DataFrame): Unit = { + val path = new File(dir, name).getCanonicalPath + df.write.mode("overwrite").option("compression", "snappy").parquet(path) + spark.read.parquet(path).createOrReplaceTempView(name) + } + + private def withComet(cacheOutputSchema: Boolean)(f: => Unit): Unit = + withSQLConf( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_CACHE_OUTPUT_SCHEMA_ENABLED.key -> cacheOutputSchema.toString)(f) + + private def runCases(name: String, rows: Long, query: String): Unit = { + val benchmark = new Benchmark(name, rows, output = output) + // Outside the timed runs, check that both arms return every row from a native plan, and that + // the plan sends its output schema with every batch, or once per task when it caches it. + for (cacheOutputSchema <- Seq(false, true)) withComet(cacheOutputSchema) { + val imports = NativeUtil.planOutputSchemaImports.sum() + val counts = countBatches(query) + val schemas = NativeUtil.planOutputSchemaImports.sum() - imports + assert(counts.rows == rows, s"$query returned ${counts.rows} rows, expected $rows") + val expected = if (cacheOutputSchema) counts.tasks else counts.batches + assert(schemas == expected, s"$query sent its schema $schemas times, expected $expected") + benchmark.out.println( + s"$name, cacheOutputSchema=$cacheOutputSchema: ${counts.batches} batches in " + + s"${counts.tasks} tasks, the schema sent with $schemas of them") + } + + CometBenchmarkBase.onOffModes(onFirst = false).foreach { case (mode, cacheOutputSchema) => + benchmark.addTimerCase(s"cacheOutputSchema $mode") { timer => + withComet(cacheOutputSchema) { + val (queryExecution, plan) = nativePlan(query) + timer.startTiming() + countBatches(queryExecution, plan) + timer.stopTiming() + } + } + } + benchmark.run() + reportCpuPerBatch(benchmark, name, query) + } + + private val processCpu = ManagementFactory.getOperatingSystemMXBean + .asInstanceOf[com.sun.management.OperatingSystemMXBean] + + /** + * Reports each arm's process CPU time per batch, the median of runs that alternate between the + * arms so that drift affects both alike. Planning is left out. The native plan runs on a Tokio + * worker while the task thread imports its batches, so wall time hides what the JVM saves until + * the JVM is the slower of the two, but a busy executor has no idle core to hide it on. + */ + private def reportCpuPerBatch(benchmark: Benchmark, name: String, query: String): Unit = { + val samples = + Map(false -> mutable.ArrayBuffer.empty[Double], true -> mutable.ArrayBuffer.empty[Double]) + for (_ <- 0 until 7; cacheOutputSchema <- Seq(false, true)) { + withComet(cacheOutputSchema) { + val plans = Seq.fill(3)(nativePlan(query)) + val start = processCpu.getProcessCpuTime + val batches = plans.map { case (queryExecution, plan) => + countBatches(queryExecution, plan).batches + }.sum + samples(cacheOutputSchema) += (processCpu.getProcessCpuTime - start).toDouble / batches + } + } + def median(cacheOutputSchema: Boolean): Double = { + val sorted = samples(cacheOutputSchema).sorted + sorted(sorted.length / 2) + } + val (off, on) = (median(false), median(true)) + benchmark.out.println( + f"$name: process CPU per batch ${off / 1000}%.1f us with the schema in every batch, " + + f"${on / 1000}%.1f us with it cached, ${(off - on) / 1000}%.1f us " + + f"(${(off - on) / off * 100}%.1f%%) less") + } + + private case class BatchCounts(rows: Long, batches: Long, tasks: Long) + + /** The planned `query`, and the root of its native plan, below the transition to rows. */ + private def nativePlan(query: String): (QueryExecution, CometNativeExec) = { + val queryExecution = spark.sql(query).queryExecution + val executedPlan = queryExecution.executedPlan + // Below any codegen stage over the transition, too. + val plan = executedPlan.collectFirst { case native: CometNativeExec => native }.getOrElse { + throw new IllegalStateException(s"Expected a native plan for $query, got:\n$executedPlan") + } + (queryExecution, plan) + } + + private def countBatches(query: String): BatchCounts = { + val (queryExecution, plan) = nativePlan(query) + countBatches(queryExecution, plan) + } + + /** + * Counts what the JVM receives from `plan`: rows, batches, and the tasks that received any + * batch. + */ + private def countBatches(queryExecution: QueryExecution, plan: CometNativeExec): BatchCounts = { + // Runs as a SQL execution so that the tasks see the session's configs. + val perTask = SQLExecution.withNewExecutionId(queryExecution) { + plan + .executeColumnar() + .mapPartitions { batches => + var rows = 0L + var count = 0L + batches.foreach { batch => + rows += batch.numRows() + count += 1 + } + Iterator((rows, count)) + } + .collect() + } + BatchCounts(perTask.map(_._1).sum, perTask.map(_._2).sum, perTask.count(_._2 > 0)) + } +}