diff --git a/native-engine/datafusion-ext-plans/src/common/execution_context.rs b/native-engine/datafusion-ext-plans/src/common/execution_context.rs index d274e194d..69459e76d 100644 --- a/native-engine/datafusion-ext-plans/src/common/execution_context.rs +++ b/native-engine/datafusion-ext-plans/src/common/execution_context.rs @@ -267,7 +267,7 @@ impl ExecutionContext { } let output_schema = self.output_schema(); - let key_converter = RowConverter::new( + let key_converter = Arc::new(Mutex::new(RowConverter::new( keys.iter() .map(|k| { Ok(SortField::new_with_options( @@ -276,7 +276,7 @@ impl ExecutionContext { )) }) .collect::>()?, - )?; + )?)); let key_exprs = keys.iter().map(|k| k.expr.clone()).collect::>(); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); @@ -291,8 +291,12 @@ impl ExecutionContext { .and_then(|r| r.into_array(batch.num_rows())) }) .collect::>>()?; - let key_rows = key_converter.convert_columns(&keys)?; - Ok(RecordBatchWithKeyRows::new(batch, Arc::new(key_rows))) + let key_rows = key_converter.lock().convert_columns(&keys)?; + Ok(RecordBatchWithKeyRows::new( + batch, + Arc::new(key_rows), + key_converter.clone(), + )) }) }); Ok(Box::pin(RecordBatchWithKeyRowsStreamAdapter::new( @@ -319,7 +323,7 @@ impl ExecutionContext { } let input_schema = input.schema(); - let key_converter = RowConverter::new( + let key_converter = Arc::new(Mutex::new(RowConverter::new( keys.iter() .map(|k| { Ok(SortField::new_with_options( @@ -328,7 +332,7 @@ impl ExecutionContext { )) }) .collect::>()?, - )?; + )?)); let projection = projection.to_vec(); let mut projection_with_keys = projection.clone(); @@ -354,7 +358,7 @@ impl ExecutionContext { .and_then(|r| r.into_array(batch.num_rows())) }) .collect::>>()?; - let key_rows = key_converter.convert_columns(&keys)?; + let key_rows = key_converter.lock().convert_columns(&keys)?; if projection.len() < projection_with_keys.len() { batch = RecordBatch::try_new_with_options( projected_schema_cloned.clone(), @@ -362,7 +366,11 @@ impl ExecutionContext { &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), )?; } - Ok(RecordBatchWithKeyRows::new(batch, Arc::new(key_rows))) + Ok(RecordBatchWithKeyRows::new( + batch, + Arc::new(key_rows), + key_converter.clone(), + )) }) }); Ok(Box::pin(RecordBatchWithKeyRowsStreamAdapter::new( diff --git a/native-engine/datafusion-ext-plans/src/common/key_rows_output.rs b/native-engine/datafusion-ext-plans/src/common/key_rows_output.rs index b347bd4d2..bf38bcda4 100644 --- a/native-engine/datafusion-ext-plans/src/common/key_rows_output.rs +++ b/native-engine/datafusion-ext-plans/src/common/key_rows_output.rs @@ -19,21 +19,35 @@ use std::{ task::{Context, Poll}, }; -use arrow::{datatypes::SchemaRef, record_batch::RecordBatch, row::Rows}; +use arrow::{ + datatypes::SchemaRef, + record_batch::RecordBatch, + row::{RowConverter, Rows}, +}; use datafusion::{common::Result, physical_expr::PhysicalSortExpr}; use futures::Stream; use futures_util::StreamExt; +use parking_lot::Mutex; /// A batch with associated key rows for optimization #[derive(Debug)] pub struct RecordBatchWithKeyRows { pub batch: RecordBatch, pub key_rows: Arc, + pub key_row_converter: Arc>, } impl RecordBatchWithKeyRows { - pub fn new(batch: RecordBatch, key_rows: Arc) -> Self { - Self { batch, key_rows } + pub fn new( + batch: RecordBatch, + key_rows: Arc, + key_row_converter: Arc>, + ) -> Self { + Self { + batch, + key_rows, + key_row_converter, + } } } diff --git a/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs b/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs index 4791c9dec..78cf62cf5 100644 --- a/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs +++ b/native-engine/datafusion-ext-plans/src/common/row_null_checker.rs @@ -15,12 +15,17 @@ // specific language governing permissions and limitations // under the License. -use arrow::{buffer::NullBuffer, row::Rows}; +use arrow::{ + array::Array, + buffer::NullBuffer, + row::{RowConverter, Rows}, +}; use arrow_schema::{DataType, Field, SchemaRef, SortOptions}; /// Utility class for checking if a Row contains null values pub struct RowNullChecker { field_configs: Vec, + needs_decoded_nulls: bool, } impl RowNullChecker { @@ -39,30 +44,54 @@ impl RowNullChecker { field_configs.push(field_config); } - Self { field_configs } + let needs_decoded_nulls = Self::needs_decoded_nulls(&field_configs); + Self { + field_configs, + needs_decoded_nulls, + } } /// Check if multiple rows contain null values and return a NullBuffer /// /// # Parameters /// * `rows` - Collection of row data to check + /// * `row_converter` - The `RowConverter` that produced `rows`; using a + /// different converter violates the precondition of + /// `RowConverter::convert_rows` and may panic /// /// # Returns /// * `NullBuffer` - Buffer indicating which rows contain null values - /// - `false` bits indicate rows that contain at least one null value - /// - `true` bits indicate rows where all fields are non-null - pub fn has_nulls(&self, rows: &Rows) -> NullBuffer { - // Create NullBuffer from the collected bits + /// - `false` bits indicate rows that contain at least one null top-level + /// key field + /// - `true` bits indicate rows where all top-level key fields are + /// non-null + pub fn has_nulls(&self, rows: &Rows, row_converter: &RowConverter) -> NullBuffer { + if !self.needs_decoded_nulls { + return NullBuffer::from_iter(rows.iter().map(|row| !self.has_null(row.data()))); + } + + let key_columns = row_converter + .convert_rows(rows.iter()) + .expect("failed to decode row data in RowNullChecker"); + let logical_nulls = key_columns + .iter() + .map(|column| column.logical_nulls()) + .collect::>(); + NullBuffer::from_iter((0..rows.num_rows()).map(|row_index| { - let row_data = rows.row(row_index); - // Check if this row has any null values - let has_null = self.has_null(row_data.as_ref()); - // NullBuffer uses true for valid (non-null) and false for null - // So we need to invert the result since has_null returns true for "has nulls" - !has_null + logical_nulls.iter().all(|nulls| match nulls { + Some(nulls) => nulls.is_valid(row_index), + None => true, + }) })) } + fn needs_decoded_nulls(field_configs: &[FieldConfig]) -> bool { + field_configs + .iter() + .any(|field_config| field_config.data_type.needs_decoded_nulls()) + } + /// Create a FieldConfig from DataType and SortOptions #[allow(clippy::panic)] // Temporarily allow panic to refactor to Result later fn create_field_config_from_data_type( @@ -143,7 +172,11 @@ impl RowNullChecker { field_configs.push(field_config); } - Self { field_configs } + let needs_decoded_nulls = Self::needs_decoded_nulls(&field_configs); + Self { + field_configs, + needs_decoded_nulls, + } } /// Create a FieldConfig from an Arrow Field (for backward compatibility) @@ -382,6 +415,12 @@ enum DataTypeInfo { Dictionary, // Dictionary type } +impl DataTypeInfo { + fn needs_decoded_nulls(&self) -> bool { + matches!(self, Self::Struct | Self::List | Self::Dictionary) + } +} + impl FieldConfig { /// Create configuration for primitive numeric types fn new_primitive(sort_options: SortOptions, encoded_length: usize) -> Self { @@ -452,10 +491,14 @@ mod tests { use std::{error::Error, sync::Arc}; use arrow::{ - array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, StringArray}, + array::{ + Array, ArrayRef, BooleanArray, DictionaryArray, Int32Array, ListArray, NullArray, + RecordBatch, StringArray, StructArray, + }, + datatypes::{Int16Type, Int32Type}, row::{RowConverter, SortField}, }; - use arrow_schema::{DataType, Field, Schema}; + use arrow_schema::{DataType, Field, Fields, Schema}; use super::*; @@ -738,7 +781,7 @@ mod tests { let checker = RowNullChecker::new(&field_configs); // Test has_nulls method - let null_buffer = checker.has_nulls(&rows); + let null_buffer = checker.has_nulls(&rows, &converter); // Verify results // Row 0: (1, "Alice") - no nulls -> should be true (valid) @@ -778,7 +821,7 @@ mod tests { let converter = RowConverter::new(sort_fields.clone())?; let rows = converter.convert_columns(&batch.columns())?; - let null_buffer = checker.has_nulls(&rows); + let null_buffer = checker.has_nulls(&rows, &converter); assert_eq!(null_buffer.len(), 0); Ok(()) } @@ -810,7 +853,7 @@ mod tests { .collect(); let checker = RowNullChecker::new(&field_configs); - let null_buffer = checker.has_nulls(&rows); + let null_buffer = checker.has_nulls(&rows, &converter); // All rows should be invalid (false) since they all contain nulls assert_eq!(null_buffer.len(), 3); @@ -851,7 +894,7 @@ mod tests { .collect(); let checker = RowNullChecker::new(&field_configs); - let null_buffer = checker.has_nulls(&rows); + let null_buffer = checker.has_nulls(&rows, &converter); // All rows should be valid (true) since none contain nulls assert_eq!(null_buffer.len(), 3); @@ -860,4 +903,116 @@ mod tests { } Ok(()) } + + fn assert_validity(null_buffer: &NullBuffer, expected: &[bool]) { + assert_eq!(null_buffer.len(), expected.len()); + for (idx, expected_valid) in expected.iter().enumerate() { + assert_eq!( + null_buffer.is_valid(idx), + *expected_valid, + "unexpected validity at row {idx}" + ); + } + } + + #[test] + fn test_has_nulls_with_null_type_on_decoded_path() -> Result<(), Box> { + let child_fields = Fields::from(vec![Field::new("id", DataType::Int32, true)]); + let struct_type = DataType::Struct(child_fields.clone()); + let struct_array = StructArray::new( + child_fields, + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2)])) as ArrayRef], + None, + ); + let field_configs = vec![ + (DataType::Null, SortOptions::default()), + (struct_type, SortOptions::default()), + ]; + let sort_fields = field_configs + .iter() + .map(|(data_type, sort_options)| { + SortField::new_with_options(data_type.clone(), *sort_options) + }) + .collect::>(); + let columns: Vec = vec![Arc::new(NullArray::new(2)), Arc::new(struct_array)]; + let converter = RowConverter::new(sort_fields)?; + let rows = converter.convert_columns(&columns)?; + + let checker = RowNullChecker::new(&field_configs); + let null_buffer = checker.has_nulls(&rows, &converter); + + assert_validity(&null_buffer, &[false, false]); + Ok(()) + } + + #[test] + fn test_has_nulls_with_complex_rows() -> Result<(), Box> { + let child_fields = Fields::from(vec![ + Field::new("id", DataType::Int32, true), + Field::new("name", DataType::Utf8, true), + ]); + let struct_type = DataType::Struct(child_fields.clone()); + let struct_array = StructArray::new( + child_fields, + vec![ + Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + None, + ])) as ArrayRef, + Arc::new(StringArray::from(vec![ + None, + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])) as ArrayRef, + ], + Some(NullBuffer::from_iter([true, false, true, true, true])), + ); + let list_array = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), None]), + Some(vec![Some(2)]), + None, + Some(vec![]), + Some(vec![None]), + ]); + let dictionary = vec![ + Some("alpha"), + Some("beta"), + Some("gamma"), + None, + Some("alpha"), + ] + .into_iter() + .collect::>(); + let field_configs = vec![ + (struct_type, SortOptions::default()), + (list_array.data_type().clone(), SortOptions::default()), + (dictionary.data_type().clone(), SortOptions::default()), + ]; + let sort_fields = field_configs + .iter() + .map(|(data_type, sort_options)| { + SortField::new_with_options(data_type.clone(), *sort_options) + }) + .collect::>(); + let columns: Vec = vec![ + Arc::new(struct_array), + Arc::new(list_array), + Arc::new(dictionary), + ]; + let converter = RowConverter::new(sort_fields)?; + let rows = converter.convert_columns(&columns)?; + + let checker = RowNullChecker::new(&field_configs); + let null_buffer = checker.has_nulls(&rows, &converter); + + // Child/null elements in rows 0 and 4 should not make top-level keys + // null. Rows 1, 2, and 3 have a null struct/list/dictionary key. + assert_validity(&null_buffer, &[true, false, false, false, true]); + Ok(()) + } } diff --git a/native-engine/datafusion-ext-plans/src/joins/stream_cursor.rs b/native-engine/datafusion-ext-plans/src/joins/stream_cursor.rs index adc107afc..201baed7a 100644 --- a/native-engine/datafusion-ext-plans/src/joins/stream_cursor.rs +++ b/native-engine/datafusion-ext-plans/src/joins/stream_cursor.rs @@ -114,9 +114,12 @@ impl StreamCursor { } let keys = batch_with_keys.key_rows.clone(); + let key_has_nulls = Some( + self.key_null_checker + .has_nulls(&keys, &batch_with_keys.key_row_converter.lock()), + ) + .filter(|nb| nb.null_count() > 0); let batch = batch_with_keys.batch; - let key_has_nulls = Some(self.key_null_checker.has_nulls(&keys)) - .filter(|nb| nb.null_count() > 0); self.batches.push(batch); self.key_has_nulls.push(key_has_nulls); diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs index 8b56b6ebe..42836492e 100644 --- a/native-engine/datafusion-ext-plans/src/sort_exec.rs +++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs @@ -805,7 +805,11 @@ async fn send_output_batch( let batch = prune_sort_keys_from_batch.restore_from_existed_key_rows(pruned_batch, &key_rows)?; sender - .send(RecordBatchWithKeyRows { batch, key_rows }) + .send(RecordBatchWithKeyRows::new( + batch, + key_rows, + prune_sort_keys_from_batch.sort_row_converter.clone(), + )) .await; } Ok(())