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 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)) + } +}