This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 211e76ec0a Fix check_not_null_constraints null detection (#13033)
211e76ec0a is described below
commit 211e76ec0a2b9f0a62aa526dc059fb19c0fa0486
Author: Piotr Findeisen <[email protected]>
AuthorDate: Wed Oct 23 13:01:12 2024 +0200
Fix check_not_null_constraints null detection (#13033)
* Fix function name typo
* Fix check_not_null_constraints null detection
`check_not_null_constraints` (aka `check_not_null_contraits`) checked
for null using `Array::null_count` which does not return real null
count.
* Drop assertor dependency
---
datafusion/physical-plan/src/execution_plan.rs | 151 +++++++++++++++++++++++--
1 file changed, 144 insertions(+), 7 deletions(-)
diff --git a/datafusion/physical-plan/src/execution_plan.rs
b/datafusion/physical-plan/src/execution_plan.rs
index a89e265ad2..e6484452d4 100644
--- a/datafusion/physical-plan/src/execution_plan.rs
+++ b/datafusion/physical-plan/src/execution_plan.rs
@@ -21,6 +21,7 @@ use std::sync::Arc;
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
+use arrow_array::Array;
use futures::stream::{StreamExt, TryStreamExt};
use tokio::task::JoinSet;
@@ -852,7 +853,7 @@ pub fn execute_input_stream(
Ok(Box::pin(RecordBatchStreamAdapter::new(
sink_schema,
input_stream
- .map(move |batch| check_not_null_contraits(batch?,
&risky_columns)),
+ .map(move |batch| check_not_null_constraints(batch?,
&risky_columns)),
)))
}
}
@@ -872,7 +873,7 @@ pub fn execute_input_stream(
/// This function iterates over the specified column indices and ensures that
none
/// of the columns contain null values. If any column contains null values, an
error
/// is returned.
-pub fn check_not_null_contraits(
+pub fn check_not_null_constraints(
batch: RecordBatch,
column_indices: &Vec<usize>,
) -> Result<RecordBatch> {
@@ -885,7 +886,13 @@ pub fn check_not_null_contraits(
);
}
- if batch.column(index).null_count() > 0 {
+ if batch
+ .column(index)
+ .logical_nulls()
+ .map(|nulls| nulls.null_count())
+ .unwrap_or_default()
+ > 0
+ {
return exec_err!(
"Invalid batch column at '{}' has null but schema specifies
non-nullable",
index
@@ -920,11 +927,11 @@ pub enum CardinalityEffect {
#[cfg(test)]
mod tests {
use super::*;
+ use arrow_array::{DictionaryArray, Int32Array, NullArray, RunArray};
+ use arrow_schema::{DataType, Field, Schema, SchemaRef};
use std::any::Any;
use std::sync::Arc;
- use arrow_schema::{Schema, SchemaRef};
-
use datafusion_common::{Result, Statistics};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
@@ -1068,6 +1075,136 @@ mod tests {
fn use_execution_plan_as_trait_object(plan: &dyn ExecutionPlan) {
let _ = plan.name();
}
-}
-// pub mod test;
+ #[test]
+ fn test_check_not_null_constraints_accept_non_null() -> Result<()> {
+ check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
true)])),
+ vec![Arc::new(Int32Array::from(vec![Some(1), Some(2),
Some(3)]))],
+ )?,
+ &vec![0],
+ )?;
+ Ok(())
+ }
+
+ #[test]
+ fn test_check_not_null_constraints_reject_null() -> Result<()> {
+ let result = check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Int32,
true)])),
+ vec![Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))],
+ )?,
+ &vec![0],
+ );
+ assert!(result.is_err());
+ assert_starts_with(
+ result.err().unwrap().message().as_ref(),
+ "Invalid batch column at '0' has null but schema specifies
non-nullable",
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn test_check_not_null_constraints_with_run_end_array() -> Result<()> {
+ // some null value inside REE array
+ let run_ends = Int32Array::from(vec![1, 2, 3, 4]);
+ let values = Int32Array::from(vec![Some(0), None, Some(1), None]);
+ let run_end_array = RunArray::try_new(&run_ends, &values)?;
+ let result = check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "a",
+ run_end_array.data_type().to_owned(),
+ true,
+ )])),
+ vec![Arc::new(run_end_array)],
+ )?,
+ &vec![0],
+ );
+ assert!(result.is_err());
+ assert_starts_with(
+ result.err().unwrap().message().as_ref(),
+ "Invalid batch column at '0' has null but schema specifies
non-nullable",
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn test_check_not_null_constraints_with_dictionary_array_with_null() ->
Result<()> {
+ let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3),
Some(4)]));
+ let keys = Int32Array::from(vec![0, 1, 2, 3]);
+ let dictionary = DictionaryArray::new(keys, values);
+ let result = check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "a",
+ dictionary.data_type().to_owned(),
+ true,
+ )])),
+ vec![Arc::new(dictionary)],
+ )?,
+ &vec![0],
+ );
+ assert!(result.is_err());
+ assert_starts_with(
+ result.err().unwrap().message().as_ref(),
+ "Invalid batch column at '0' has null but schema specifies
non-nullable",
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn test_check_not_null_constraints_with_dictionary_masking_null() ->
Result<()> {
+ // some null value marked out by dictionary array
+ let values = Arc::new(Int32Array::from(vec![
+ Some(1),
+ None, // this null value is masked by dictionary keys
+ Some(3),
+ Some(4),
+ ]));
+ let keys = Int32Array::from(vec![0, /*1,*/ 2, 3]);
+ let dictionary = DictionaryArray::new(keys, values);
+ check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new(
+ "a",
+ dictionary.data_type().to_owned(),
+ true,
+ )])),
+ vec![Arc::new(dictionary)],
+ )?,
+ &vec![0],
+ )?;
+ Ok(())
+ }
+
+ #[test]
+ fn test_check_not_null_constraints_on_null_type() -> Result<()> {
+ // null value of Null type
+ let result = check_not_null_constraints(
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![Field::new("a", DataType::Null,
true)])),
+ vec![Arc::new(NullArray::new(3))],
+ )?,
+ &vec![0],
+ );
+ assert!(result.is_err());
+ assert_starts_with(
+ result.err().unwrap().message().as_ref(),
+ "Invalid batch column at '0' has null but schema specifies
non-nullable",
+ );
+ Ok(())
+ }
+
+ fn assert_starts_with(actual: impl AsRef<str>, expected_prefix: impl
AsRef<str>) {
+ let actual = actual.as_ref();
+ let expected_prefix = expected_prefix.as_ref();
+ assert!(
+ actual.starts_with(expected_prefix),
+ "Expected '{}' to start with '{}'",
+ actual,
+ expected_prefix
+ );
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
