[AURON #2419] RowNullChecker may mis-detect null join keys for complex Arrow row types

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::<Result<_>>()?,
-        )?;
+        )?));
         let key_exprs = keys.iter().map(|k| k.expr.clone()).collect::<Vec<_>>();
 
         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::<Result<Vec<_>>>()?;
-                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::<Result<_>>()?,
-        )?;
+        )?));
 
         let projection = projection.to_vec();
         let mut projection_with_keys = projection.clone();
@@ -354,15 +358,19 @@ impl ExecutionContext {
                             .and_then(|r| r.into_array(batch.num_rows()))
                     })
                     .collect::<Result<Vec<_>>>()?;
-                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(),
                         batch.columns()[..projection.len()].to_vec(),
                         &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(

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<Rows>,
+    pub key_row_converter: Arc<Mutex<RowConverter>>,
 }
 
 impl RecordBatchWithKeyRows {
-    pub fn new(batch: RecordBatch, key_rows: Arc<Rows>) -> Self {
-        Self { batch, key_rows }
+    pub fn new(
+        batch: RecordBatch,
+        key_rows: Arc<Rows>,
+        key_row_converter: Arc<Mutex<RowConverter>>,
+    ) -> Self {
+        Self {
+            batch,
+            key_rows,
+            key_row_converter,
+        }
     }
 }
 

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<FieldConfig>,
+    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::<Vec<_>>();
+
         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<dyn Error>> {
+        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::<Vec<_>>();
+        let columns: Vec<ArrayRef> = 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<dyn Error>> {
+        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::<Int32Type, _, _>(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::<DictionaryArray<Int16Type>>();
+        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::<Vec<_>>();
+        let columns: Vec<ArrayRef> = 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(())
+    }
 }