This is an automated email from the ASF dual-hosted git repository.
slfan1989 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 0f1794ca [AURON #2419] RowNullChecker may mis-detect null join keys
for complex Arrow row types (#2424)
0f1794ca is described below
commit 0f1794cab2eb3db080f31be2cb02a9b9fc43c78d
Author: Ming Wei <[email protected]>
AuthorDate: Fri Jul 31 08:29:23 2026 +0800
[AURON #2419] RowNullChecker may mis-detect null join keys for complex
Arrow row types (#2424)
# Which issue does this PR close?
Closes #2419
**Rationale for this change**
Auron's sort merge join uses encoded Arrow rows for join key comparison
and uses RowNullChecker to skip rows with null join keys.
RowNullChecker previously inspected Arrow row bytes manually. This is
fragile for complex row encodings such as struct, list, and dictionary,
because their layout is controlled by Arrow RowConverter and should not
be interpreted with fixed or estimated field lengths.
This can cause RowNullChecker to mis-detect null join keys for complex
key types.
**What changes are included in this PR?**
Adds an Arrow RowConverter to RowNullChecker.
Updates RowNullChecker's `has_nulls` path to decode Rows through Arrow
RowConverter before building the null-key mask.
Preserves conservative join key null semantics:
top-level null key values are treated as null join keys
nested null fields or elements inside non-null struct/list keys do not
make the key null
Keeps the existing RowNullChecker public method signatures unchanged.
Keeps the existing byte-level helper path unchanged.
Adds unit coverage for complex struct, list, and dictionary key rows.
**Are there any user-facing changes?**
No user-facing API changes.
**How was this patch tested?**
UT.
**Was this patch authored or co-authored using generative AI tooling?**
- [x] Yes
- [ ] No
Generated-by: OpenAI Codex (GPT-5)
---------
Co-authored-by: Shilun Fan <[email protected]>
Signed-off-by: weimingdiit <[email protected]>
---
.../src/common/execution_context.rs | 24 ++-
.../src/common/key_rows_output.rs | 20 ++-
.../src/common/row_null_checker.rs | 193 +++++++++++++++++++--
.../src/joins/stream_cursor.rs | 7 +-
.../datafusion-ext-plans/src/sort_exec.rs | 6 +-
5 files changed, 217 insertions(+), 33 deletions(-)
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 d274e194..69459e76 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::<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,7 +358,7 @@ 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(),
@@ -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 b347bd4d..bf38bcda 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<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,
+ }
}
}
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 4791c9de..78cf62cf 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<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(())
+ }
}
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 adc107af..201baed7 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 8b56b6eb..42836492 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(())
