nullable transform

Author: DrakeLin

kernel/src/scan/data_skipping.rs

@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::sync::{Arc, LazyLock};
 
@@ -15,12 +14,14 @@ use crate::expressions::{
 use crate::kernel_predicates::{
     DataSkippingPredicateEvaluator, KernelPredicateEvaluator, KernelPredicateEvaluatorDefaults,
 };
-use crate::schema::{DataType, PrimitiveType, SchemaRef, SchemaTransform, StructField, StructType};
+use crate::schema::{DataType, SchemaRef, SchemaTransform, StructField, StructType};
 use crate::{
     Engine, EngineData, ExpressionEvaluator, JsonHandler, PredicateEvaluator, RowVisitor as _,
 };
 
 pub(crate) mod stats_schema;
+
+use stats_schema::{NullableStatsTransform, NullCountStatsTransform};
 #[cfg(test)]
 mod tests;
 
@@ -77,39 +78,6 @@ impl DataSkippingFilter {
         let (predicate, referenced_schema) = physical_predicate?;
         debug!("Creating a data skipping filter for {:#?}", predicate);
 
-        // Convert all fields into nullable, as stats may not be available for all columns
-        // (and usually aren't for partition columns).
-        struct NullableStatsTransform;
-        impl<'a> SchemaTransform<'a> for NullableStatsTransform {
-            fn transform_struct_field(
-                &mut self,
-                field: &'a StructField,
-            ) -> Option<Cow<'a, StructField>> {
-                use Cow::*;
-                let field = match self.transform(&field.data_type)? {
-                    Borrowed(_) if field.is_nullable() => Borrowed(field),
-                    data_type => Owned(StructField {
-                        name: field.name.clone(),
-                        data_type: data_type.into_owned(),
-                        nullable: true,
-                        metadata: field.metadata.clone(),
-                    }),
-                };
-                Some(field)
-            }
-        }
-
-        // Convert a min/max stats schema into a nullcount schema (all leaf fields are LONG)
-        struct NullCountStatsTransform;
-        impl<'a> SchemaTransform<'a> for NullCountStatsTransform {
-            fn transform_primitive(
-                &mut self,
-                _ptype: &'a PrimitiveType,
-            ) -> Option<Cow<'a, PrimitiveType>> {
-                Some(Cow::Owned(PrimitiveType::Long))
-            }
-        }
-
         let stats_schema = NullableStatsTransform
             .transform_struct(&referenced_schema)?
             .into_owned();

kernel/src/scan/data_skipping/stats_schema.rs

@@ -111,6 +111,25 @@ pub(crate) fn expected_stats_schema(
     StructType::try_new(fields)
 }
 
+/// Transforms a schema to make all fields nullable.
+/// Used for stats schemas where stats may not be available for all columns.
+pub(crate) struct NullableStatsTransform;
+impl<'a> SchemaTransform<'a> for NullableStatsTransform {
+    fn transform_struct_field(&mut self, field: &'a StructField) -> Option<Cow<'a, StructField>> {
+        use Cow::*;
+        let field = match self.transform(&field.data_type)? {
+            Borrowed(_) if field.is_nullable() => Borrowed(field),
+            data_type => Owned(StructField {
+                name: field.name.clone(),
+                data_type: data_type.into_owned(),
+                nullable: true,
+                metadata: field.metadata.clone(),
+            }),
+        };
+        Some(field)
+    }
+}
+
 // Convert a min/max stats schema into a nullcount schema (all leaf fields are LONG)
 #[allow(unused)]
 pub(crate) struct NullCountStatsTransform;
@@ -317,26 +336,6 @@ mod tests {
 
     use super::*;
 
-    pub(crate) struct NullableStatsTransform;
-    impl<'a> SchemaTransform<'a> for NullableStatsTransform {
-        fn transform_struct_field(
-            &mut self,
-            field: &'a StructField,
-        ) -> Option<Cow<'a, StructField>> {
-            use Cow::*;
-            let field = match self.transform(&field.data_type)? {
-                Borrowed(_) if field.is_nullable() => Borrowed(field),
-                data_type => Owned(StructField {
-                    name: field.name.clone(),
-                    data_type: data_type.into_owned(),
-                    nullable: true,
-                    metadata: field.metadata.clone(),
-                }),
-            };
-            Some(field)
-        }
-    }
-
     #[test]
     fn test_should_include_column() {
         let full_name = vec![ColumnName::new(["lvl1", "lvl2", "lvl3", "lvl4"])];

kernel/src/transaction/mod.rs

@@ -1,4 +1,3 @@
-use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};
 use std::iter;
 use std::ops::Deref;
@@ -23,6 +22,7 @@ use crate::expressions::{column_name, ColumnName};
 use crate::expressions::{ArrayData, Scalar, StructData, Transform, UnaryExpressionOp::ToJson};
 use crate::path::{LogRoot, ParsedLogPath};
 use crate::row_tracking::{RowTrackingDomainMetadata, RowTrackingVisitor};
+use crate::scan::data_skipping::stats_schema::NullableStatsTransform;
 use crate::scan::log_replay::{
     get_scan_metadata_transform_expr, BASE_ROW_ID_NAME, DEFAULT_ROW_COMMIT_VERSION_NAME,
     FILE_CONSTANT_VALUES_NAME, TAGS_NAME,
@@ -41,28 +41,6 @@ use crate::{
 };
 use delta_kernel_derive::internal_api;
 
-// This is a workaround due to the fact that expression evaluation happens
-// on the whole EngineData instead of accounting for filtered rows, which can lead to null values in
-// required fields.
-// TODO: Move this to a common place (dedupe from data_skipping.rs) or remove when evaluations work
-// on FilteredEngineData directly.
-struct NullableStatsTransform;
-impl<'a> SchemaTransform<'a> for NullableStatsTransform {
-    fn transform_struct_field(&mut self, field: &'a StructField) -> Option<Cow<'a, StructField>> {
-        use Cow::*;
-        let field = match self.transform(&field.data_type)? {
-            Borrowed(_) if field.is_nullable() => Borrowed(field),
-            data_type => Owned(StructField {
-                name: field.name.clone(),
-                data_type: data_type.into_owned(),
-                nullable: true,
-                metadata: field.metadata.clone(),
-            }),
-        };
-        Some(field)
-    }
-}
-
 /// Type alias for an iterator of [`EngineData`] results.
 pub(crate) type EngineDataResultIterator<'a> =
     Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send + 'a>;