fix: reject NULL partition keys during partitioned writes

datafusion/datasource/src/write/demux.rs

@@ -38,7 +38,9 @@ use datafusion_common::cast::{
     as_int64_array, as_large_string_array, as_string_array, as_string_view_array,
     as_uint8_array, as_uint16_array, as_uint32_array, as_uint64_array,
 };
-use datafusion_common::{exec_datafusion_err, internal_datafusion_err, not_impl_err};
+use datafusion_common::{
+    exec_datafusion_err, exec_err, internal_datafusion_err, not_impl_err,
+};
 use datafusion_common_runtime::SpawnedTask;
 
 use chrono::NaiveDate;
@@ -392,6 +394,11 @@ fn compute_partition_keys_by_row<'a>(
             "PartitionBy Column {} does not exist in source data! Got schema {schema}.",
             col
         ))?;
+        if col_array.logical_null_count() > 0 {
+            return exec_err!(
+                "NULL values are not supported for partition column '{col}'"
+            );
+        }
 
         match dtype {
             DataType::Utf8 => {
@@ -595,3 +602,69 @@ fn compute_hive_style_file_path(
 
     file_path.join(format!("{write_id}.{file_extension}"))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow::array::{Array, ArrayRef, DictionaryArray, Int32Array};
+    use arrow::datatypes::Int32Type;
+
+    fn partition_batch(column: ArrayRef) -> RecordBatch {
+        RecordBatch::try_from_iter_with_nullable([("p", column, true)]).unwrap()
+    }
+
+    #[test]
+    fn partition_keys_reject_null_values() {
+        let columns: Vec<ArrayRef> = vec![
+            Arc::new(Int32Array::from(vec![Some(0), None])),
+            Arc::new(StringArray::from(vec![Some(""), None])),
+        ];
+        for column in columns {
+            let partition_by = vec![("p".to_string(), column.data_type().clone())];
+            let batch = partition_batch(column);
+            let err = compute_partition_keys_by_row(&batch, &partition_by).unwrap_err();
+            assert_eq!(
+                err.to_string(),
+                "Execution error: NULL values are not supported for partition column 'p'"
+            );
+        }
+    }
+
+    #[test]
+    fn partition_keys_accept_slice_without_nulls() {
+        let column = Int32Array::from(vec![None, Some(0), Some(1)]).slice(1, 2);
+        let batch = partition_batch(Arc::new(column));
+        let partition_by = vec![("p".to_string(), DataType::Int32)];
+        let keys = compute_partition_keys_by_row(&batch, &partition_by).unwrap();
+        assert_eq!(keys, vec![vec![Cow::Borrowed("0"), Cow::Borrowed("1")]]);
+    }
+
+    fn dictionary_with_null_value() -> DictionaryArray<Int32Type> {
+        DictionaryArray::try_new(
+            Int32Array::from(vec![0, 1]),
+            Arc::new(StringArray::from(vec![Some("a"), None])),
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn partition_keys_reject_dictionary_null_values() {
+        let column = dictionary_with_null_value();
+        let partition_by = vec![("p".to_string(), column.data_type().clone())];
+        let batch = partition_batch(Arc::new(column));
+        let err = compute_partition_keys_by_row(&batch, &partition_by).unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Execution error: NULL values are not supported for partition column 'p'"
+        );
+    }
+
+    #[test]
+    fn partition_keys_accept_unused_dictionary_null_values() {
+        let column = dictionary_with_null_value().slice(0, 1);
+        let partition_by = vec![("p".to_string(), column.data_type().clone())];
+        let batch = partition_batch(Arc::new(column));
+        let keys = compute_partition_keys_by_row(&batch, &partition_by).unwrap();
+        assert_eq!(keys, vec![vec![Cow::Borrowed("a")]]);
+    }
+}

datafusion/sqllogictest/test_files/copy.slt

@@ -42,6 +42,28 @@ select * from validate_partitioned_parquet order by col1, col2;
 1 Foo
 2 Bar
 
+# Reject NULL partition values instead of writing them into the zero partition.
+statement ok
+CREATE TABLE nullable_partition_source (id INT, p INT) AS VALUES (1, NULL), (2, 0);
+
+statement error DataFusion error: Execution error: NULL values are not supported for partition column 'p'
+COPY nullable_partition_source TO 'test_files/scratch/copy/null_partition/' STORED AS PARQUET PARTITIONED BY (p);
+
+# Nullable partition columns still work when the copied rows contain no NULLs.
+query I
+COPY (SELECT * FROM nullable_partition_source WHERE p IS NOT NULL) TO 'test_files/scratch/copy/non_null_partition/' STORED AS PARQUET PARTITIONED BY (p);
+----
+1
+
+statement ok
+CREATE EXTERNAL TABLE validate_non_null_partition (id INT, p INT)
+STORED AS PARQUET LOCATION 'test_files/scratch/copy/non_null_partition/' PARTITIONED BY (p);
+
+query II
+SELECT id, p FROM validate_non_null_partition;
+----
+2 0
+
 # validate partition paths were actually generated
 statement ok
 CREATE EXTERNAL TABLE validate_partitioned_parquet_bar STORED AS PARQUET