datafusion_physical_expr_adapter/
schema_rewriter.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Physical expression schema rewriting utilities: [`PhysicalExprAdapter`],
19//! [`PhysicalExprAdapterFactory`], default implementations,
20//! and [`replace_columns_with_literals`].
21
22use std::borrow::Borrow;
23use std::collections::HashMap;
24use std::hash::Hash;
25use std::sync::Arc;
26
27use arrow::array::RecordBatch;
28use arrow::datatypes::{DataType, FieldRef, SchemaRef};
29use datafusion_common::{
30    DataFusionError, Result, ScalarValue, exec_err,
31    metadata::FieldMetadata,
32    nested_struct::validate_data_type_compatibility,
33    tree_node::{Transformed, TransformedResult, TreeNode},
34};
35use datafusion_functions::core::getfield::GetFieldFunc;
36use datafusion_physical_expr::PhysicalExprSimplifier;
37use datafusion_physical_expr::projection::{ProjectionExprs, Projector};
38use datafusion_physical_expr::{
39    ScalarFunctionExpr,
40    expressions::{self, CastExpr, Column},
41};
42use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
43use itertools::Itertools;
44
45/// Replace column references in the given physical expression with literal values.
46///
47/// Some use cases for this include:
48/// - Partition column pruning: When scanning partitioned data, partition column references
49///   can be replaced with their literal values for the specific partition being scanned.
50/// - Constant folding: In some cases, columns that can be proven to be constant
51///   from statistical analysis may be replaced with their literal values to optimize expression evaluation.
52/// - Filling in non-null default values: in a custom [`PhysicalExprAdapter`] implementation,
53///   column references can be replaced with default literal values instead of nulls.
54///
55/// # Arguments
56/// - `expr`: The physical expression in which to replace column references.
57/// - `replacements`: A mapping from column names to their corresponding literal `ScalarValue`s.
58///   Accepts various HashMap types including `HashMap<&str, &ScalarValue>`,
59///   `HashMap<String, ScalarValue>`, `HashMap<String, &ScalarValue>`, etc.
60///
61/// # Returns
62/// - `Result<Arc<dyn PhysicalExpr>>`: The rewritten physical expression with columns replaced by literals.
63pub fn replace_columns_with_literals<K, V>(
64    expr: Arc<dyn PhysicalExpr>,
65    replacements: &HashMap<K, V>,
66) -> Result<Arc<dyn PhysicalExpr>>
67where
68    K: Borrow<str> + Eq + Hash,
69    V: Borrow<ScalarValue>,
70{
71    expr.transform_down(|expr| {
72        if let Some(column) = expr.downcast_ref::<Column>()
73            && let Some(replacement_value) = replacements.get(column.name())
74        {
75            return Ok(Transformed::yes(expressions::lit(
76                replacement_value.borrow().clone(),
77            )));
78        }
79        Ok(Transformed::no(expr))
80    })
81    .data()
82}
83
84/// Trait for adapting [`PhysicalExpr`] expressions to match a target schema.
85///
86/// This is used in file scans to rewrite expressions so that they can be
87/// evaluated against the physical schema of the file being scanned. It allows
88/// for handling differences between logical and physical schemas, such as type
89/// mismatches or missing columns common in [Schema evolution] scenarios.
90///
91/// [Schema evolution]: https://www.dremio.com/wiki/schema-evolution/
92///
93/// ## Default Implementations
94///
95/// The default implementation [`DefaultPhysicalExprAdapter`]  handles common
96/// cases.
97///
98/// ## Custom Implementations
99///
100/// You can create a custom implementation of this trait to handle specific rewriting logic.
101/// For example, to fill in missing columns with default values instead of nulls:
102///
103/// ```rust
104/// use datafusion_physical_expr_adapter::{PhysicalExprAdapter, PhysicalExprAdapterFactory};
105/// use arrow::datatypes::{Schema, Field, DataType, FieldRef, SchemaRef};
106/// use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
107/// use datafusion_common::{Result, ScalarValue, tree_node::{Transformed, TransformedResult, TreeNode}};
108/// use datafusion_physical_expr::expressions::{self, Column};
109/// use std::sync::Arc;
110///
111/// #[derive(Debug)]
112/// pub struct CustomPhysicalExprAdapter {
113///     logical_file_schema: SchemaRef,
114///     physical_file_schema: SchemaRef,
115/// }
116///
117/// impl PhysicalExprAdapter for CustomPhysicalExprAdapter {
118///     fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
119///         expr.transform(|expr| {
120///             if let Some(column) = expr.downcast_ref::<Column>() {
121///                 // Check if the column exists in the physical schema
122///                 if self.physical_file_schema.index_of(column.name()).is_err() {
123///                     // If the column is missing, fill it with a default value instead of null
124///                     // The default value could be stored in the table schema's column metadata for example.
125///                     let default_value = ScalarValue::Int32(Some(0));
126///                     return Ok(Transformed::yes(expressions::lit(default_value)));
127///                 }
128///             }
129///             // If the column exists, return it as is
130///             Ok(Transformed::no(expr))
131///         }).data()
132///     }
133/// }
134///
135/// #[derive(Debug)]
136/// pub struct CustomPhysicalExprAdapterFactory;
137///
138/// impl PhysicalExprAdapterFactory for CustomPhysicalExprAdapterFactory {
139///     fn create(
140///         &self,
141///         logical_file_schema: SchemaRef,
142///         physical_file_schema: SchemaRef,
143///     ) -> Result<Arc<dyn PhysicalExprAdapter>> {
144///         Ok(Arc::new(CustomPhysicalExprAdapter {
145///             logical_file_schema,
146///             physical_file_schema,
147///         }))
148///     }
149/// }
150/// ```
151pub trait PhysicalExprAdapter: Send + Sync + std::fmt::Debug {
152    /// Rewrite a physical expression to match the target schema.
153    ///
154    /// This method should return a transformed expression that matches the target schema.
155    ///
156    /// Arguments:
157    /// - `expr`: The physical expression to rewrite.
158    /// - `logical_file_schema`: The logical schema of the table being queried, excluding any partition columns.
159    /// - `physical_file_schema`: The physical schema of the file being scanned.
160    /// - `partition_values`: Optional partition values to use for rewriting partition column references.
161    ///   These are handled as if they were columns appended onto the logical file schema.
162    ///
163    /// Returns:
164    /// - `Arc<dyn PhysicalExpr>`: The rewritten physical expression that can be evaluated against the physical schema.
165    ///
166    /// See Also:
167    /// - [`replace_columns_with_literals`]: for replacing partition column references with their literal values.
168    fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>>;
169}
170
171/// Creates instances of [`PhysicalExprAdapter`] for given logical and physical schemas.
172///
173/// See [`DefaultPhysicalExprAdapterFactory`] for the default implementation.
174pub trait PhysicalExprAdapterFactory: Send + Sync + std::fmt::Debug {
175    /// Create a new instance of the physical expression adapter.
176    fn create(
177        &self,
178        logical_file_schema: SchemaRef,
179        physical_file_schema: SchemaRef,
180    ) -> Result<Arc<dyn PhysicalExprAdapter>>;
181}
182
183#[derive(Debug, Clone)]
184pub struct DefaultPhysicalExprAdapterFactory;
185
186impl PhysicalExprAdapterFactory for DefaultPhysicalExprAdapterFactory {
187    fn create(
188        &self,
189        logical_file_schema: SchemaRef,
190        physical_file_schema: SchemaRef,
191    ) -> Result<Arc<dyn PhysicalExprAdapter>> {
192        Ok(Arc::new(DefaultPhysicalExprAdapter {
193            logical_file_schema,
194            physical_file_schema,
195        }))
196    }
197}
198
199/// Default implementation of [`PhysicalExprAdapter`] for rewriting physical
200/// expressions to match different schemas.
201///
202/// ## Overview
203///
204///  [`DefaultPhysicalExprAdapter`] rewrites physical expressions to match
205///  different schemas, including:
206///
207/// - **Type casting**: When logical and physical schemas have different types, expressions are
208///   automatically wrapped with cast operations. For example, `lit(ScalarValue::Int32(123)) = int64_column`
209///   gets rewritten to `lit(ScalarValue::Int32(123)) = cast(int64_column, 'Int32')`.
210///   Note that this does not attempt to simplify such expressions - that is done by shared simplifiers.
211///
212/// - **Missing columns**: When a column exists in the logical schema but not in the physical schema,
213///   references to it are replaced with null literals.
214///
215/// - **Struct field access**: Expressions like `struct_column.field_that_is_missing_in_schema` are
216///   rewritten to `null` when the field doesn't exist in the physical schema.
217///
218/// - **Default column values**: Partition column references can be replaced with their literal values
219///   when scanning specific partitions. See [`replace_columns_with_literals`] for more details.
220///
221/// # Example
222///
223/// ```rust
224/// # use datafusion_physical_expr_adapter::{DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory};
225/// # use arrow::datatypes::Schema;
226/// # use std::sync::Arc;
227/// #
228/// # fn example(
229/// #     predicate: std::sync::Arc<dyn datafusion_physical_expr_common::physical_expr::PhysicalExpr>,
230/// #     physical_file_schema: &Schema,
231/// #     logical_file_schema: &Schema,
232/// # ) -> datafusion_common::Result<()> {
233/// let factory = DefaultPhysicalExprAdapterFactory;
234/// let adapter =
235///     factory.create(Arc::new(logical_file_schema.clone()), Arc::new(physical_file_schema.clone()))?;
236/// let adapted_predicate = adapter.rewrite(predicate)?;
237/// # Ok(())
238/// # }
239/// ```
240#[derive(Debug, Clone)]
241pub struct DefaultPhysicalExprAdapter {
242    logical_file_schema: SchemaRef,
243    physical_file_schema: SchemaRef,
244}
245
246impl DefaultPhysicalExprAdapter {
247    /// Create a new instance of the default physical expression adapter.
248    ///
249    /// This adapter rewrites expressions to match the physical schema of the file being scanned,
250    /// handling type mismatches and missing columns by filling them with default values.
251    pub fn new(logical_file_schema: SchemaRef, physical_file_schema: SchemaRef) -> Self {
252        Self {
253            logical_file_schema,
254            physical_file_schema,
255        }
256    }
257}
258
259impl PhysicalExprAdapter for DefaultPhysicalExprAdapter {
260    fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
261        let rewriter = DefaultPhysicalExprAdapterRewriter {
262            logical_file_schema: Arc::clone(&self.logical_file_schema),
263            physical_file_schema: Arc::clone(&self.physical_file_schema),
264        };
265        expr.transform(|expr| rewriter.rewrite_expr(Arc::clone(&expr)))
266            .data()
267    }
268}
269
270struct DefaultPhysicalExprAdapterRewriter {
271    logical_file_schema: SchemaRef,
272    physical_file_schema: SchemaRef,
273}
274
275impl DefaultPhysicalExprAdapterRewriter {
276    fn rewrite_expr(
277        &self,
278        expr: Arc<dyn PhysicalExpr>,
279    ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
280        if let Some(transformed) = self.try_rewrite_struct_field_access(&expr)? {
281            return Ok(Transformed::yes(transformed));
282        }
283
284        if let Some(column) = expr.downcast_ref::<Column>() {
285            return self.rewrite_column(Arc::clone(&expr), column);
286        }
287
288        Ok(Transformed::no(expr))
289    }
290
291    /// Attempt to rewrite struct field access expressions to return null if the field does not exist in the physical schema.
292    /// Note that this does *not* handle nested struct fields, only top-level struct field access.
293    /// See <https://github.com/apache/datafusion/issues/17114> for more details.
294    fn try_rewrite_struct_field_access(
295        &self,
296        expr: &Arc<dyn PhysicalExpr>,
297    ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
298        let get_field_expr =
299            match ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(expr.as_ref()) {
300                Some(expr) => expr,
301                None => return Ok(None),
302            };
303
304        let source_expr = match get_field_expr.args().first() {
305            Some(expr) => expr,
306            None => return Ok(None),
307        };
308
309        let field_name_expr = match get_field_expr.args().get(1) {
310            Some(expr) => expr,
311            None => return Ok(None),
312        };
313
314        let lit = match field_name_expr.downcast_ref::<expressions::Literal>() {
315            Some(lit) => lit,
316            None => return Ok(None),
317        };
318
319        let field_name = match lit.value().try_as_str().flatten() {
320            Some(name) => name,
321            None => return Ok(None),
322        };
323
324        let column = match source_expr.downcast_ref::<Column>() {
325            Some(column) => column,
326            None => return Ok(None),
327        };
328
329        let physical_field =
330            match self.physical_file_schema.field_with_name(column.name()) {
331                Ok(field) => field,
332                Err(_) => return Ok(None),
333            };
334
335        let physical_struct_fields = match physical_field.data_type() {
336            DataType::Struct(fields) => fields,
337            _ => return Ok(None),
338        };
339
340        if physical_struct_fields
341            .iter()
342            .any(|f| f.name() == field_name)
343        {
344            return Ok(None);
345        }
346
347        let logical_field = match self.logical_file_schema.field_with_name(column.name())
348        {
349            Ok(field) => field,
350            Err(_) => return Ok(None),
351        };
352
353        let logical_struct_fields = match logical_field.data_type() {
354            DataType::Struct(fields) => fields,
355            _ => return Ok(None),
356        };
357
358        let logical_struct_field = match logical_struct_fields
359            .iter()
360            .find(|f| f.name() == field_name)
361        {
362            Some(field) => field,
363            None => return Ok(None),
364        };
365
366        let null_value = ScalarValue::Null.cast_to(logical_struct_field.data_type())?;
367        Ok(Some(Arc::new(expressions::Literal::new_with_metadata(
368            null_value,
369            Some(FieldMetadata::from(logical_struct_field.as_ref())),
370        ))))
371    }
372
373    fn rewrite_column(
374        &self,
375        expr: Arc<dyn PhysicalExpr>,
376        column: &Column,
377    ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
378        // Get the logical field for this column if it exists in the logical schema
379        let logical_field = match self.logical_file_schema.field_with_name(column.name())
380        {
381            Ok(field) => field,
382            Err(e) => {
383                // This can be hit if a custom rewrite injected a reference to a column that doesn't exist in the logical schema.
384                // For example, a pre-computed column that is kept only in the physical schema.
385                // If the column exists in the physical schema, we can still use it.
386                if let Ok(physical_field) =
387                    self.physical_file_schema.field_with_name(column.name())
388                {
389                    // If the column exists in the physical schema, we can use it in place of the logical column.
390                    // This is nice to users because if they do a rewrite that results in something like `physical_int32_col = 123u64`
391                    // we'll at least handle the casts for them.
392                    physical_field
393                } else {
394                    // A completely unknown column that doesn't exist in either schema!
395                    // This should probably never be hit unless something upstream broke, but nonetheless it's better
396                    // for us to return a handleable error than to panic / do something unexpected.
397                    return Err(e.into());
398                }
399            }
400        };
401
402        let Some((resolved_column, physical_field)) =
403            self.resolve_physical_column(column)?
404        else {
405            if !logical_field.is_nullable() {
406                return exec_err!(
407                    "Non-nullable column '{}' is missing from the physical schema",
408                    column.name()
409                );
410            }
411            // If the column is missing from the physical schema fill it in with nulls.
412            // For a different behavior, provide a custom `PhysicalExprAdapter` implementation.
413            let null_value = ScalarValue::Null.cast_to(logical_field.data_type())?;
414            return Ok(Transformed::yes(Arc::new(
415                expressions::Literal::new_with_metadata(
416                    null_value,
417                    Some(FieldMetadata::from(logical_field)),
418                ),
419            )));
420        };
421
422        let fields_match = logical_field == physical_field.as_ref();
423        if fields_match {
424            if resolved_column.index() == column.index() {
425                return Ok(Transformed::no(expr));
426            }
427
428            // If the fields match (including metadata/nullability), we can use the column as is
429            return Ok(Transformed::yes(Arc::new(resolved_column)));
430        }
431
432        // We need a cast expression whenever the logical and physical fields differ,
433        // whether that difference is only metadata/nullability or also data type.
434        // TODO: add optimization to move the cast from the column to literal expressions in the case of `col = 123`
435        // since that's much cheaper to evalaute.
436        // See https://github.com/apache/datafusion/issues/15780#issuecomment-2824716928
437        validate_data_type_compatibility(
438            resolved_column.name(),
439            physical_field.data_type(),
440            logical_field.data_type(),
441        )
442        .map_err(|e| {
443            DataFusionError::Execution(format!(
444                "Cannot cast column '{}' from '{}' (physical data type) to '{}' (logical data type): {e}",
445                resolved_column.name(),
446                physical_field.data_type(),
447                logical_field.data_type()
448            ))
449        })?;
450
451        Ok(Transformed::yes(Arc::new(CastExpr::new_with_target_field(
452            Arc::new(resolved_column),
453            Arc::new(logical_field.clone()),
454            None,
455        ))))
456    }
457
458    /// Resolves a logical column to the corresponding physical column and field.
459    fn resolve_physical_column(
460        &self,
461        column: &Column,
462    ) -> Result<Option<(Column, FieldRef)>> {
463        // The physical schema adaptation step intentionally resolves columns by **name first**
464        // rather than trusting the incoming index. This mirrors what the old refactoring
465        // did before `resolve_physical_column()` was extracted: the planner might hand us a
466        // `Column` whose `index` field is stale (e.g. after projection/rename rewrites), so
467        // resolving by name ensures we match the correct physical slot. Once we know the
468        // proper index we rebuild the `Column` with `new_with_schema` so callers can rely
469        // on `column.index()` later without having to re-query the schema.
470        let Ok(physical_column_index) = self.physical_file_schema.index_of(column.name())
471        else {
472            return Ok(None);
473        };
474
475        let column = if column.index() == physical_column_index {
476            column.clone()
477        } else {
478            Column::new_with_schema(column.name(), self.physical_file_schema.as_ref())?
479        };
480
481        let physical_field = Arc::new(
482            self.physical_file_schema
483                .field(physical_column_index)
484                .clone(),
485        );
486
487        Ok(Some((column, physical_field)))
488    }
489}
490
491/// Factory for creating [`BatchAdapter`] instances to adapt record batches
492/// to a target schema.
493///
494/// This binds a target schema and allows creating adapters for different source schemas.
495/// It handles:
496/// - **Column reordering**: Columns are reordered to match the target schema
497/// - **Type casting**: Automatic type conversion (e.g., Int32 to Int64)
498/// - **Missing columns**: Nullable columns missing from source are filled with nulls
499/// - **Struct field adaptation**: Nested struct fields are recursively adapted
500///
501/// ## Examples
502///
503/// ```rust
504/// use arrow::array::{Int32Array, Int64Array, StringArray, RecordBatch};
505/// use arrow::datatypes::{DataType, Field, Schema};
506/// use datafusion_physical_expr_adapter::BatchAdapterFactory;
507/// use std::sync::Arc;
508///
509/// // Target schema has different column order and types
510/// let target_schema = Arc::new(Schema::new(vec![
511///     Field::new("name", DataType::Utf8, true),
512///     Field::new("id", DataType::Int64, false),    // Int64 in target
513///     Field::new("score", DataType::Float64, true), // Missing from source
514/// ]));
515///
516/// // Source schema has different column order and Int32 for id
517/// let source_schema = Arc::new(Schema::new(vec![
518///     Field::new("id", DataType::Int32, false),    // Int32 in source
519///     Field::new("name", DataType::Utf8, true),
520///     // Note: 'score' column is missing from source
521/// ]));
522///
523/// // Create factory with target schema
524/// let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
525///
526/// // Create adapter for this specific source schema
527/// let adapter = factory.make_adapter(&source_schema).unwrap();
528///
529/// // Create a source batch
530/// let source_batch = RecordBatch::try_new(
531///     source_schema,
532///     vec![
533///         Arc::new(Int32Array::from(vec![1, 2, 3])),
534///         Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol"])),
535///     ],
536/// ).unwrap();
537///
538/// // Adapt the batch to match target schema
539/// let adapted = adapter.adapt_batch(&source_batch).unwrap();
540///
541/// assert_eq!(adapted.num_columns(), 3);
542/// assert_eq!(adapted.column(0).data_type(), &DataType::Utf8);   // name
543/// assert_eq!(adapted.column(1).data_type(), &DataType::Int64);  // id (cast from Int32)
544/// assert_eq!(adapted.column(2).data_type(), &DataType::Float64); // score (filled with nulls)
545/// ```
546#[derive(Debug)]
547pub struct BatchAdapterFactory {
548    target_schema: SchemaRef,
549    expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
550}
551
552impl BatchAdapterFactory {
553    /// Create a new [`BatchAdapterFactory`] with the given target schema.
554    pub fn new(target_schema: SchemaRef) -> Self {
555        let expr_adapter_factory = Arc::new(DefaultPhysicalExprAdapterFactory);
556        Self {
557            target_schema,
558            expr_adapter_factory,
559        }
560    }
561
562    /// Set a custom [`PhysicalExprAdapterFactory`] to use when adapting expressions.
563    ///
564    /// Use this to customize behavior when adapting batches, e.g. to fill in missing values
565    /// with defaults instead of nulls.
566    ///
567    /// See [`PhysicalExprAdapter`] for more details.
568    pub fn with_adapter_factory(
569        self,
570        factory: Arc<dyn PhysicalExprAdapterFactory>,
571    ) -> Self {
572        Self {
573            expr_adapter_factory: factory,
574            ..self
575        }
576    }
577
578    /// Create a new [`BatchAdapter`] for the given source schema.
579    ///
580    /// Batches fed into this [`BatchAdapter`] *must* conform to the source schema,
581    /// no validation is performed at runtime to minimize overheads.
582    pub fn make_adapter(&self, source_schema: &SchemaRef) -> Result<BatchAdapter> {
583        let expr_adapter = self
584            .expr_adapter_factory
585            .create(Arc::clone(&self.target_schema), Arc::clone(source_schema))?;
586
587        let simplifier = PhysicalExprSimplifier::new(&self.target_schema);
588
589        let projection = ProjectionExprs::from_indices(
590            &(0..self.target_schema.fields().len()).collect_vec(),
591            &self.target_schema,
592        );
593
594        let adapted = projection
595            .try_map_exprs(|e| simplifier.simplify(expr_adapter.rewrite(e)?))?;
596        let projector = adapted.make_projector(source_schema)?;
597
598        Ok(BatchAdapter { projector })
599    }
600}
601
602/// Adapter for transforming record batches to match a target schema.
603///
604/// Create instances via [`BatchAdapterFactory`].
605///
606/// ## Performance
607///
608/// The adapter pre-computes the projection expressions during creation,
609/// so the [`adapt_batch`](BatchAdapter::adapt_batch) call is efficient and suitable
610/// for use in hot paths like streaming file scans.
611#[derive(Debug)]
612pub struct BatchAdapter {
613    projector: Projector,
614}
615
616impl BatchAdapter {
617    /// Adapt the given record batch to match the target schema.
618    ///
619    /// The input batch *must* conform to the source schema used when
620    /// creating this adapter.
621    pub fn adapt_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
622        self.projector.project_batch(batch)
623    }
624}
625
626#[cfg(test)]
627mod tests {
628    use super::*;
629    use arrow::array::{
630        Array, BooleanArray, GenericListArray, Int32Array, Int64Array, RecordBatch,
631        RecordBatchOptions, StringArray, StringViewArray, StructArray,
632    };
633    use arrow::datatypes::{Field, Fields, Schema};
634    use datafusion_common::{assert_contains, record_batch};
635    use datafusion_expr::Operator;
636    use datafusion_physical_expr::expressions::{Column, Literal, col};
637
638    fn assert_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> &CastExpr {
639        expr.downcast_ref::<CastExpr>().expect("Expected CastExpr")
640    }
641
642    fn assert_cast_input_column(cast_expr: &CastExpr, name: &str, index: usize) {
643        let inner_col = cast_expr
644            .expr()
645            .downcast_ref::<Column>()
646            .expect("Expected inner Column");
647        assert_eq!(inner_col.name(), name);
648        assert_eq!(inner_col.index(), index);
649    }
650
651    fn stale_index_cast_schemas() -> (SchemaRef, SchemaRef) {
652        let physical_schema = Arc::new(Schema::new(vec![
653            Field::new("b", DataType::Binary, true),
654            Field::new("a", DataType::Int32, false),
655        ]));
656
657        let logical_schema = Arc::new(Schema::new(vec![
658            Field::new("a", DataType::Int64, false),
659            Field::new("b", DataType::Binary, true),
660        ]));
661
662        (logical_schema, physical_schema)
663    }
664
665    fn create_test_schema() -> (Schema, Schema) {
666        let physical_schema = Schema::new(vec![
667            Field::new("a", DataType::Int32, false),
668            Field::new("b", DataType::Utf8, true),
669        ]);
670
671        let logical_schema = Schema::new(vec![
672            Field::new("a", DataType::Int64, false), // Different type
673            Field::new("b", DataType::Utf8, true),
674            Field::new("c", DataType::Float64, true), // Missing from physical
675        ]);
676
677        (physical_schema, logical_schema)
678    }
679
680    #[test]
681    fn test_rewrite_column_with_type_cast() {
682        let (physical_schema, logical_schema) = create_test_schema();
683
684        let factory = DefaultPhysicalExprAdapterFactory;
685        let adapter = factory
686            .create(Arc::new(logical_schema), Arc::new(physical_schema))
687            .unwrap();
688        let column_expr = Arc::new(Column::new("a", 0));
689
690        let result = adapter.rewrite(column_expr).unwrap();
691
692        // Should be wrapped in a cast expression
693        assert!(result.downcast_ref::<CastExpr>().is_some());
694    }
695
696    #[test]
697    fn test_rewrite_column_with_metadata_or_nullability_mismatch() -> Result<()> {
698        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
699        let logical_schema =
700            Schema::new(vec![Field::new("a", DataType::Int64, false).with_metadata(
701                HashMap::from([("logical_meta".to_string(), "1".to_string())]),
702            )]);
703
704        let factory = DefaultPhysicalExprAdapterFactory;
705        let adapter = factory
706            .create(Arc::new(logical_schema), Arc::new(physical_schema.clone()))
707            .unwrap();
708
709        let result = adapter.rewrite(Arc::new(Column::new("a", 0)))?;
710
711        // Ensure the expression preserves the logical field nullability/metadata.
712        let return_field = result.return_field(physical_schema.as_ref())?;
713        assert_eq!(return_field.data_type(), &DataType::Int64);
714        assert!(!return_field.is_nullable());
715        assert_eq!(
716            return_field
717                .metadata()
718                .get("logical_meta")
719                .map(String::as_str),
720            Some("1")
721        );
722
723        Ok(())
724    }
725
726    #[test]
727    fn test_rewrite_multi_column_expr_with_type_cast() {
728        let (physical_schema, logical_schema) = create_test_schema();
729        let factory = DefaultPhysicalExprAdapterFactory;
730        let adapter = factory
731            .create(Arc::new(logical_schema), Arc::new(physical_schema))
732            .unwrap();
733
734        // Create a complex expression: (a + 5) OR (c > 0.0) that tests the recursive case of the rewriter
735        let column_a = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
736        let column_c = Arc::new(Column::new("c", 2)) as Arc<dyn PhysicalExpr>;
737        let expr = expressions::BinaryExpr::new(
738            Arc::clone(&column_a),
739            Operator::Plus,
740            Arc::new(Literal::new(ScalarValue::Int64(Some(5)))),
741        );
742        let expr = expressions::BinaryExpr::new(
743            Arc::new(expr),
744            Operator::Or,
745            Arc::new(expressions::BinaryExpr::new(
746                Arc::clone(&column_c),
747                Operator::Gt,
748                Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
749            )),
750        );
751
752        let result = adapter.rewrite(Arc::new(expr)).unwrap();
753        let outer = result
754            .downcast_ref::<expressions::BinaryExpr>()
755            .expect("Expected outer BinaryExpr");
756        assert_eq!(*outer.op(), Operator::Or);
757
758        let left = outer
759            .left()
760            .downcast_ref::<expressions::BinaryExpr>()
761            .expect("Expected left BinaryExpr");
762        assert_eq!(*left.op(), Operator::Plus);
763
764        let left_cast = assert_cast_expr(left.left());
765        assert_eq!(left_cast.target_field().data_type(), &DataType::Int64);
766        assert_cast_input_column(left_cast, "a", 0);
767
768        let right = outer
769            .right()
770            .downcast_ref::<expressions::BinaryExpr>()
771            .expect("Expected right BinaryExpr");
772        assert_eq!(*right.op(), Operator::Gt);
773        let null_literal = right
774            .left()
775            .downcast_ref::<Literal>()
776            .expect("Expected null literal");
777        assert_eq!(*null_literal.value(), ScalarValue::Float64(None));
778    }
779
780    #[test]
781    fn test_rewrite_struct_column_incompatible() {
782        let physical_schema = Schema::new(vec![Field::new(
783            "data",
784            DataType::Struct(vec![Field::new("field1", DataType::Binary, true)].into()),
785            true,
786        )]);
787
788        let logical_schema = Schema::new(vec![Field::new(
789            "data",
790            DataType::Struct(vec![Field::new("field1", DataType::Int32, true)].into()),
791            true,
792        )]);
793
794        let factory = DefaultPhysicalExprAdapterFactory;
795        let adapter = factory
796            .create(Arc::new(logical_schema), Arc::new(physical_schema))
797            .unwrap();
798        let column_expr = Arc::new(Column::new("data", 0));
799
800        let error_msg = adapter.rewrite(column_expr).unwrap_err().to_string();
801        // validate_struct_compatibility provides more specific error about which field can't be cast
802        assert_contains!(
803            error_msg,
804            "Cannot cast struct field 'field1' from type Binary to type Int32"
805        );
806    }
807
808    #[test]
809    fn test_rewrite_struct_compatible_cast() {
810        let physical_schema = Schema::new(vec![Field::new(
811            "data",
812            DataType::Struct(
813                vec![
814                    Field::new("id", DataType::Int32, false),
815                    Field::new("name", DataType::Utf8, true),
816                ]
817                .into(),
818            ),
819            false,
820        )]);
821
822        let logical_schema = Schema::new(vec![Field::new(
823            "data",
824            DataType::Struct(
825                vec![
826                    Field::new("id", DataType::Int64, false),
827                    Field::new("name", DataType::Utf8View, true),
828                ]
829                .into(),
830            ),
831            false,
832        )]);
833
834        let factory = DefaultPhysicalExprAdapterFactory;
835        let adapter = factory
836            .create(Arc::new(logical_schema), Arc::new(physical_schema))
837            .unwrap();
838        let column_expr = Arc::new(Column::new("data", 0));
839
840        let result = adapter.rewrite(column_expr).unwrap();
841
842        let logical_struct_fields: Fields = vec![
843            Field::new("id", DataType::Int64, false),
844            Field::new("name", DataType::Utf8View, true),
845        ]
846        .into();
847        let logical_field = Arc::new(Field::new(
848            "data",
849            DataType::Struct(logical_struct_fields),
850            false,
851        ));
852
853        let expected = Arc::new(CastExpr::new_with_target_field(
854            Arc::new(Column::new("data", 0)),
855            logical_field,
856            None,
857        )) as Arc<dyn PhysicalExpr>;
858
859        assert_eq!(result.to_string(), expected.to_string());
860    }
861
862    #[test]
863    fn test_rewrite_missing_column() -> Result<()> {
864        let (physical_schema, logical_schema) = create_test_schema();
865
866        let factory = DefaultPhysicalExprAdapterFactory;
867        let adapter = factory
868            .create(Arc::new(logical_schema), Arc::new(physical_schema))
869            .unwrap();
870        let column_expr = Arc::new(Column::new("c", 2));
871
872        let result = adapter.rewrite(column_expr)?;
873
874        // Should be replaced with a literal null
875        if let Some(literal) = result.downcast_ref::<Literal>() {
876            assert_eq!(*literal.value(), ScalarValue::Float64(None));
877        } else {
878            panic!("Expected literal expression");
879        }
880
881        Ok(())
882    }
883
884    #[test]
885    fn test_rewrite_missing_column_propagates_metadata() -> Result<()> {
886        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
887        let logical_schema = Schema::new(vec![
888            Field::new("a", DataType::Int32, false),
889            Field::new("b", DataType::Utf8, true).with_metadata(HashMap::from([(
890                "logical_meta".to_string(),
891                "1".to_string(),
892            )])),
893        ]);
894
895        let factory = DefaultPhysicalExprAdapterFactory;
896        let adapter = factory
897            .create(Arc::new(logical_schema), Arc::new(physical_schema.clone()))
898            .unwrap();
899
900        let result = adapter.rewrite(Arc::new(Column::new("b", 1)))?;
901        let literal = result
902            .downcast_ref::<Literal>()
903            .expect("Expected literal expression");
904
905        assert_eq!(
906            literal
907                .return_field(physical_schema.as_ref())?
908                .metadata()
909                .get("logical_meta")
910                .map(String::as_str),
911            Some("1")
912        );
913        Ok(())
914    }
915
916    #[test]
917    fn test_rewrite_missing_column_non_nullable_error() {
918        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
919        let logical_schema = Schema::new(vec![
920            Field::new("a", DataType::Int64, false),
921            Field::new("b", DataType::Utf8, false), // Missing and non-nullable
922        ]);
923
924        let factory = DefaultPhysicalExprAdapterFactory;
925        let adapter = factory
926            .create(Arc::new(logical_schema), Arc::new(physical_schema))
927            .unwrap();
928        let column_expr = Arc::new(Column::new("b", 1));
929
930        let error_msg = adapter.rewrite(column_expr).unwrap_err().to_string();
931        assert_contains!(error_msg, "Non-nullable column 'b' is missing");
932    }
933
934    #[test]
935    fn test_rewrite_missing_column_nullable() {
936        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
937        let logical_schema = Schema::new(vec![
938            Field::new("a", DataType::Int64, false),
939            Field::new("b", DataType::Utf8, true), // Missing but nullable
940        ]);
941
942        let factory = DefaultPhysicalExprAdapterFactory;
943        let adapter = factory
944            .create(Arc::new(logical_schema), Arc::new(physical_schema))
945            .unwrap();
946        let column_expr = Arc::new(Column::new("b", 1));
947
948        let result = adapter.rewrite(column_expr).unwrap();
949
950        let expected =
951            Arc::new(Literal::new(ScalarValue::Utf8(None))) as Arc<dyn PhysicalExpr>;
952
953        assert_eq!(result.to_string(), expected.to_string());
954    }
955
956    #[test]
957    fn test_replace_columns_with_literals() -> Result<()> {
958        let partition_value = ScalarValue::Utf8(Some("test_value".to_string()));
959        let replacements = HashMap::from([("partition_col", &partition_value)]);
960
961        let column_expr =
962            Arc::new(Column::new("partition_col", 0)) as Arc<dyn PhysicalExpr>;
963        let result = replace_columns_with_literals(column_expr, &replacements)?;
964
965        // Should be replaced with the partition value
966        let literal = result
967            .downcast_ref::<Literal>()
968            .expect("Expected literal expression");
969        assert_eq!(*literal.value(), partition_value);
970
971        Ok(())
972    }
973
974    #[test]
975    fn test_replace_columns_with_literals_no_match() -> Result<()> {
976        let value = ScalarValue::Utf8(Some("test_value".to_string()));
977        let replacements = HashMap::from([("other_col", &value)]);
978
979        let column_expr =
980            Arc::new(Column::new("partition_col", 0)) as Arc<dyn PhysicalExpr>;
981        let result = replace_columns_with_literals(column_expr, &replacements)?;
982
983        assert!(result.downcast_ref::<Column>().is_some());
984        Ok(())
985    }
986
987    #[test]
988    fn test_replace_columns_with_literals_nested_expr() -> Result<()> {
989        let value_a = ScalarValue::Int64(Some(10));
990        let value_b = ScalarValue::Int64(Some(20));
991        let replacements = HashMap::from([("a", &value_a), ("b", &value_b)]);
992
993        let expr = Arc::new(expressions::BinaryExpr::new(
994            Arc::new(Column::new("a", 0)),
995            Operator::Plus,
996            Arc::new(Column::new("b", 1)),
997        )) as Arc<dyn PhysicalExpr>;
998
999        let result = replace_columns_with_literals(expr, &replacements)?;
1000        assert_eq!(result.to_string(), "10 + 20");
1001
1002        Ok(())
1003    }
1004
1005    #[test]
1006    fn test_rewrite_no_change_needed() -> Result<()> {
1007        let (physical_schema, logical_schema) = create_test_schema();
1008
1009        let factory = DefaultPhysicalExprAdapterFactory;
1010        let adapter = factory
1011            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1012            .unwrap();
1013        let column_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
1014
1015        let result = adapter.rewrite(Arc::clone(&column_expr))?;
1016
1017        // Should be the same expression (no transformation needed)
1018        // We compare the underlying pointer through the trait object
1019        assert!(std::ptr::eq(
1020            column_expr.as_ref() as *const dyn PhysicalExpr,
1021            result.as_ref() as *const dyn PhysicalExpr
1022        ));
1023
1024        Ok(())
1025    }
1026
1027    #[test]
1028    fn test_non_nullable_missing_column_error() {
1029        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1030        let logical_schema = Schema::new(vec![
1031            Field::new("a", DataType::Int32, false),
1032            Field::new("b", DataType::Utf8, false), // Non-nullable missing column
1033        ]);
1034
1035        let factory = DefaultPhysicalExprAdapterFactory;
1036        let adapter = factory
1037            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1038            .unwrap();
1039        let column_expr = Arc::new(Column::new("b", 1));
1040
1041        let result = adapter.rewrite(column_expr);
1042        assert!(result.is_err());
1043        assert_contains!(
1044            result.unwrap_err().to_string(),
1045            "Non-nullable column 'b' is missing from the physical schema"
1046        );
1047    }
1048
1049    /// Helper function to project expressions onto a RecordBatch
1050    fn batch_project(
1051        expr: Vec<Arc<dyn PhysicalExpr>>,
1052        batch: &RecordBatch,
1053        schema: SchemaRef,
1054    ) -> Result<RecordBatch> {
1055        let arrays = expr
1056            .iter()
1057            .map(|expr| {
1058                expr.evaluate(batch)
1059                    .and_then(|v| v.into_array(batch.num_rows()))
1060            })
1061            .collect::<Result<Vec<_>>>()?;
1062
1063        if arrays.is_empty() {
1064            let options =
1065                RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1066            RecordBatch::try_new_with_options(Arc::clone(&schema), arrays, &options)
1067                .map_err(Into::into)
1068        } else {
1069            RecordBatch::try_new(Arc::clone(&schema), arrays).map_err(Into::into)
1070        }
1071    }
1072
1073    /// Example showing how we can use the `DefaultPhysicalExprAdapter` to adapt RecordBatches during a scan
1074    /// to apply projections, type conversions and handling of missing columns all at once.
1075    #[test]
1076    fn test_adapt_batches() {
1077        let physical_batch = record_batch!(
1078            ("a", Int32, vec![Some(1), None, Some(3)]),
1079            ("extra", Utf8, vec![Some("x"), Some("y"), None])
1080        )
1081        .unwrap();
1082
1083        let physical_schema = physical_batch.schema();
1084
1085        let logical_schema = Arc::new(Schema::new(vec![
1086            Field::new("a", DataType::Int64, true), // Different type
1087            Field::new("b", DataType::Utf8, true),  // Missing from physical
1088        ]));
1089
1090        let projection = vec![
1091            col("b", &logical_schema).unwrap(),
1092            col("a", &logical_schema).unwrap(),
1093        ];
1094
1095        let factory = DefaultPhysicalExprAdapterFactory;
1096        let adapter = factory
1097            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1098            .unwrap();
1099
1100        let adapted_projection = projection
1101            .into_iter()
1102            .map(|expr| adapter.rewrite(expr).unwrap())
1103            .collect_vec();
1104
1105        let adapted_schema = Arc::new(Schema::new(
1106            adapted_projection
1107                .iter()
1108                .map(|expr| expr.return_field(&physical_schema).unwrap())
1109                .collect_vec(),
1110        ));
1111
1112        let res = batch_project(
1113            adapted_projection,
1114            &physical_batch,
1115            Arc::clone(&adapted_schema),
1116        )
1117        .unwrap();
1118
1119        assert_eq!(res.num_columns(), 2);
1120        assert_eq!(res.column(0).data_type(), &DataType::Utf8);
1121        assert_eq!(res.column(1).data_type(), &DataType::Int64);
1122        assert_eq!(
1123            res.column(0)
1124                .as_any()
1125                .downcast_ref::<StringArray>()
1126                .unwrap()
1127                .iter()
1128                .collect_vec(),
1129            vec![None, None, None]
1130        );
1131        assert_eq!(
1132            res.column(1)
1133                .as_any()
1134                .downcast_ref::<Int64Array>()
1135                .unwrap()
1136                .iter()
1137                .collect_vec(),
1138            vec![Some(1), None, Some(3)]
1139        );
1140    }
1141
1142    /// Test that struct columns are properly adapted including:
1143    /// - Type casting of subfields (Int32 -> Int64, Utf8 -> Utf8View)
1144    /// - Missing fields in logical schema are filled with nulls
1145    #[test]
1146    fn test_adapt_struct_batches() {
1147        // Physical struct: {id: Int32, name: Utf8}
1148        let physical_struct_fields: Fields = vec![
1149            Field::new("id", DataType::Int32, false),
1150            Field::new("name", DataType::Utf8, true),
1151        ]
1152        .into();
1153
1154        let struct_array = StructArray::new(
1155            physical_struct_fields.clone(),
1156            vec![
1157                Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
1158                Arc::new(StringArray::from(vec![
1159                    Some("alice"),
1160                    None,
1161                    Some("charlie"),
1162                ])) as _,
1163            ],
1164            None,
1165        );
1166
1167        let physical_schema = Arc::new(Schema::new(vec![Field::new(
1168            "data",
1169            DataType::Struct(physical_struct_fields),
1170            false,
1171        )]));
1172
1173        let physical_batch = RecordBatch::try_new(
1174            Arc::clone(&physical_schema),
1175            vec![Arc::new(struct_array)],
1176        )
1177        .unwrap();
1178
1179        // Logical struct: {id: Int64, name: Utf8View, extra: Boolean}
1180        // - id: cast from Int32 to Int64
1181        // - name: cast from Utf8 to Utf8View
1182        // - extra: missing from physical, should be filled with nulls
1183        let logical_struct_fields: Fields = vec![
1184            Field::new("id", DataType::Int64, false),
1185            Field::new("name", DataType::Utf8View, true),
1186            Field::new("extra", DataType::Boolean, true), // New field, not in physical
1187        ]
1188        .into();
1189
1190        let logical_schema = Arc::new(Schema::new(vec![Field::new(
1191            "data",
1192            DataType::Struct(logical_struct_fields),
1193            false,
1194        )]));
1195
1196        let projection = vec![col("data", &logical_schema).unwrap()];
1197
1198        let factory = DefaultPhysicalExprAdapterFactory;
1199        let adapter = factory
1200            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1201            .unwrap();
1202
1203        let adapted_projection = projection
1204            .into_iter()
1205            .map(|expr| adapter.rewrite(expr).unwrap())
1206            .collect_vec();
1207
1208        let adapted_schema = Arc::new(Schema::new(
1209            adapted_projection
1210                .iter()
1211                .map(|expr| expr.return_field(&physical_schema).unwrap())
1212                .collect_vec(),
1213        ));
1214
1215        let res = batch_project(
1216            adapted_projection,
1217            &physical_batch,
1218            Arc::clone(&adapted_schema),
1219        )
1220        .unwrap();
1221
1222        assert_eq!(res.num_columns(), 1);
1223
1224        let result_struct = res
1225            .column(0)
1226            .as_any()
1227            .downcast_ref::<StructArray>()
1228            .unwrap();
1229
1230        // Verify id field is cast to Int64
1231        let id_col = result_struct.column_by_name("id").unwrap();
1232        assert_eq!(id_col.data_type(), &DataType::Int64);
1233        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
1234        assert_eq!(
1235            id_values.iter().collect_vec(),
1236            vec![Some(1), Some(2), Some(3)]
1237        );
1238
1239        // Verify name field is cast to Utf8View
1240        let name_col = result_struct.column_by_name("name").unwrap();
1241        assert_eq!(name_col.data_type(), &DataType::Utf8View);
1242        let name_values = name_col.as_any().downcast_ref::<StringViewArray>().unwrap();
1243        assert_eq!(
1244            name_values.iter().collect_vec(),
1245            vec![Some("alice"), None, Some("charlie")]
1246        );
1247
1248        // Verify extra field (missing from physical) is filled with nulls
1249        let extra_col = result_struct.column_by_name("extra").unwrap();
1250        assert_eq!(extra_col.data_type(), &DataType::Boolean);
1251        let extra_values = extra_col.as_any().downcast_ref::<BooleanArray>().unwrap();
1252        assert_eq!(extra_values.iter().collect_vec(), vec![None, None, None]);
1253    }
1254
1255    /// Test that List<Struct> columns are properly adapted with struct evolution.
1256    #[test]
1257    fn test_adapt_list_struct_batches() {
1258        // Physical: List<{id: Int32, name: Utf8}>
1259        let physical_struct_fields: Fields = vec![
1260            Field::new("id", DataType::Int32, false),
1261            Field::new("name", DataType::Utf8, true),
1262        ]
1263        .into();
1264
1265        let struct_array = StructArray::new(
1266            physical_struct_fields.clone(),
1267            vec![
1268                Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
1269                Arc::new(StringArray::from(vec![
1270                    Some("alice"),
1271                    None,
1272                    Some("charlie"),
1273                ])) as _,
1274            ],
1275            None,
1276        );
1277
1278        // One list element per row
1279        let item_field = Arc::new(Field::new(
1280            "item",
1281            DataType::Struct(physical_struct_fields.clone()),
1282            true,
1283        ));
1284        let offsets =
1285            arrow::buffer::OffsetBuffer::from_lengths(vec![1usize; struct_array.len()]);
1286        let list_array = GenericListArray::<i32>::new(
1287            item_field,
1288            offsets,
1289            Arc::new(struct_array),
1290            None,
1291        );
1292
1293        let physical_schema = Arc::new(Schema::new(vec![Field::new(
1294            "data",
1295            DataType::List(Arc::new(Field::new(
1296                "item",
1297                DataType::Struct(physical_struct_fields),
1298                true,
1299            ))),
1300            false,
1301        )]));
1302
1303        let physical_batch = RecordBatch::try_new(
1304            Arc::clone(&physical_schema),
1305            vec![Arc::new(list_array)],
1306        )
1307        .unwrap();
1308
1309        // Logical: List<{id: Int64, name: Utf8View, extra: Boolean}>
1310        let logical_struct_fields: Fields = vec![
1311            Field::new("id", DataType::Int64, false),
1312            Field::new("name", DataType::Utf8View, true),
1313            Field::new("extra", DataType::Boolean, true),
1314        ]
1315        .into();
1316
1317        let logical_schema = Arc::new(Schema::new(vec![Field::new(
1318            "data",
1319            DataType::List(Arc::new(Field::new(
1320                "item",
1321                DataType::Struct(logical_struct_fields.clone()),
1322                true,
1323            ))),
1324            false,
1325        )]));
1326
1327        let projection = vec![col("data", &logical_schema).unwrap()];
1328
1329        let factory = DefaultPhysicalExprAdapterFactory;
1330        let adapter = factory
1331            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1332            .unwrap();
1333
1334        let adapted_projection = projection
1335            .into_iter()
1336            .map(|expr| adapter.rewrite(expr).unwrap())
1337            .collect_vec();
1338
1339        let adapted_schema = Arc::new(Schema::new(
1340            adapted_projection
1341                .iter()
1342                .map(|expr| expr.return_field(&physical_schema).unwrap())
1343                .collect_vec(),
1344        ));
1345
1346        let res = batch_project(
1347            adapted_projection,
1348            &physical_batch,
1349            Arc::clone(&adapted_schema),
1350        )
1351        .unwrap();
1352
1353        assert_eq!(res.num_columns(), 1);
1354
1355        let result_list = res
1356            .column(0)
1357            .as_any()
1358            .downcast_ref::<GenericListArray<i32>>()
1359            .unwrap();
1360
1361        // Check each list element contains the evolved struct
1362        assert_eq!(result_list.len(), 3);
1363        let flat_structs = result_list
1364            .values()
1365            .as_any()
1366            .downcast_ref::<StructArray>()
1367            .unwrap();
1368
1369        let id_col = flat_structs.column_by_name("id").unwrap();
1370        assert_eq!(id_col.data_type(), &DataType::Int64);
1371        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
1372        assert_eq!(
1373            id_values.iter().collect_vec(),
1374            vec![Some(1), Some(2), Some(3)]
1375        );
1376
1377        let name_col = flat_structs.column_by_name("name").unwrap();
1378        assert_eq!(name_col.data_type(), &DataType::Utf8View);
1379        let name_values = name_col.as_any().downcast_ref::<StringViewArray>().unwrap();
1380        assert_eq!(
1381            name_values.iter().collect_vec(),
1382            vec![Some("alice"), None, Some("charlie")]
1383        );
1384
1385        let extra_col = flat_structs.column_by_name("extra").unwrap();
1386        assert_eq!(extra_col.data_type(), &DataType::Boolean);
1387        let extra_values = extra_col.as_any().downcast_ref::<BooleanArray>().unwrap();
1388        assert_eq!(extra_values.iter().collect_vec(), vec![None, None, None]);
1389    }
1390
1391    #[test]
1392    fn test_try_rewrite_struct_field_access() {
1393        // Test the core logic of try_rewrite_struct_field_access
1394        let physical_schema = Schema::new(vec![Field::new(
1395            "struct_col",
1396            DataType::Struct(
1397                vec![Field::new("existing_field", DataType::Int32, true)].into(),
1398            ),
1399            true,
1400        )]);
1401
1402        let logical_schema = Schema::new(vec![Field::new(
1403            "struct_col",
1404            DataType::Struct(
1405                vec![
1406                    Field::new("existing_field", DataType::Int32, true),
1407                    Field::new("missing_field", DataType::Utf8, true),
1408                ]
1409                .into(),
1410            ),
1411            true,
1412        )]);
1413
1414        let rewriter = DefaultPhysicalExprAdapterRewriter {
1415            logical_file_schema: Arc::new(logical_schema),
1416            physical_file_schema: Arc::new(physical_schema),
1417        };
1418
1419        // Test that when a field exists in physical schema, it returns None
1420        let column = Arc::new(Column::new("struct_col", 0)) as Arc<dyn PhysicalExpr>;
1421        let result = rewriter.try_rewrite_struct_field_access(&column).unwrap();
1422        assert!(result.is_none());
1423
1424        // The actual test for the get_field expression would require creating a proper ScalarFunctionExpr
1425        // with ScalarUDF, which is complex to set up in a unit test. The integration tests in
1426        // datafusion/core/tests/parquet/schema_adapter.rs provide better coverage for this functionality.
1427    }
1428
1429    // ============================================================================
1430    // BatchAdapterFactory and BatchAdapter tests
1431    // ============================================================================
1432
1433    #[test]
1434    fn test_batch_adapter_factory_basic() {
1435        // Target schema
1436        let target_schema = Arc::new(Schema::new(vec![
1437            Field::new("a", DataType::Int64, false),
1438            Field::new("b", DataType::Utf8, true),
1439        ]));
1440
1441        // Source schema with different column order and type
1442        let source_schema = Arc::new(Schema::new(vec![
1443            Field::new("b", DataType::Utf8, true),
1444            Field::new("a", DataType::Int32, false), // Int32 -> Int64
1445        ]));
1446
1447        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1448        let adapter = factory.make_adapter(&source_schema).unwrap();
1449
1450        // Create source batch
1451        let source_batch = RecordBatch::try_new(
1452            Arc::clone(&source_schema),
1453            vec![
1454                Arc::new(StringArray::from(vec![Some("hello"), None, Some("world")])),
1455                Arc::new(Int32Array::from(vec![1, 2, 3])),
1456            ],
1457        )
1458        .unwrap();
1459
1460        let adapted = adapter.adapt_batch(&source_batch).unwrap();
1461
1462        // Verify schema matches target
1463        assert_eq!(adapted.num_columns(), 2);
1464        assert_eq!(adapted.schema().field(0).name(), "a");
1465        assert_eq!(adapted.schema().field(0).data_type(), &DataType::Int64);
1466        assert_eq!(adapted.schema().field(1).name(), "b");
1467        assert_eq!(adapted.schema().field(1).data_type(), &DataType::Utf8);
1468
1469        // Verify data
1470        let col_a = adapted
1471            .column(0)
1472            .as_any()
1473            .downcast_ref::<Int64Array>()
1474            .unwrap();
1475        assert_eq!(col_a.iter().collect_vec(), vec![Some(1), Some(2), Some(3)]);
1476
1477        let col_b = adapted
1478            .column(1)
1479            .as_any()
1480            .downcast_ref::<StringArray>()
1481            .unwrap();
1482        assert_eq!(
1483            col_b.iter().collect_vec(),
1484            vec![Some("hello"), None, Some("world")]
1485        );
1486    }
1487
1488    #[test]
1489    fn test_batch_adapter_factory_missing_column() {
1490        // Target schema with a column missing from source
1491        let target_schema = Arc::new(Schema::new(vec![
1492            Field::new("a", DataType::Int32, false),
1493            Field::new("b", DataType::Utf8, true), // exists in source
1494            Field::new("c", DataType::Float64, true), // missing from source
1495        ]));
1496
1497        let source_schema = Arc::new(Schema::new(vec![
1498            Field::new("a", DataType::Int32, false),
1499            Field::new("b", DataType::Utf8, true),
1500        ]));
1501
1502        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1503        let adapter = factory.make_adapter(&source_schema).unwrap();
1504
1505        let source_batch = RecordBatch::try_new(
1506            Arc::clone(&source_schema),
1507            vec![
1508                Arc::new(Int32Array::from(vec![1, 2])),
1509                Arc::new(StringArray::from(vec!["x", "y"])),
1510            ],
1511        )
1512        .unwrap();
1513
1514        let adapted = adapter.adapt_batch(&source_batch).unwrap();
1515
1516        assert_eq!(adapted.num_columns(), 3);
1517
1518        // Missing column should be filled with nulls
1519        let col_c = adapted.column(2);
1520        assert_eq!(col_c.data_type(), &DataType::Float64);
1521        assert_eq!(col_c.null_count(), 2); // All nulls
1522    }
1523
1524    #[test]
1525    fn test_batch_adapter_factory_with_struct() {
1526        // Target has struct with Int64 id
1527        let target_struct_fields: Fields = vec![
1528            Field::new("id", DataType::Int64, false),
1529            Field::new("name", DataType::Utf8, true),
1530        ]
1531        .into();
1532        let target_schema = Arc::new(Schema::new(vec![Field::new(
1533            "data",
1534            DataType::Struct(target_struct_fields),
1535            false,
1536        )]));
1537
1538        // Source has struct with Int32 id
1539        let source_struct_fields: Fields = vec![
1540            Field::new("id", DataType::Int32, false),
1541            Field::new("name", DataType::Utf8, true),
1542        ]
1543        .into();
1544        let source_schema = Arc::new(Schema::new(vec![Field::new(
1545            "data",
1546            DataType::Struct(source_struct_fields.clone()),
1547            false,
1548        )]));
1549
1550        let struct_array = StructArray::new(
1551            source_struct_fields,
1552            vec![
1553                Arc::new(Int32Array::from(vec![10, 20])) as _,
1554                Arc::new(StringArray::from(vec!["a", "b"])) as _,
1555            ],
1556            None,
1557        );
1558
1559        let source_batch = RecordBatch::try_new(
1560            Arc::clone(&source_schema),
1561            vec![Arc::new(struct_array)],
1562        )
1563        .unwrap();
1564
1565        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1566        let adapter = factory.make_adapter(&source_schema).unwrap();
1567        let adapted = adapter.adapt_batch(&source_batch).unwrap();
1568
1569        let result_struct = adapted
1570            .column(0)
1571            .as_any()
1572            .downcast_ref::<StructArray>()
1573            .unwrap();
1574
1575        // Verify id was cast to Int64
1576        let id_col = result_struct.column_by_name("id").unwrap();
1577        assert_eq!(id_col.data_type(), &DataType::Int64);
1578        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
1579        assert_eq!(id_values.iter().collect_vec(), vec![Some(10), Some(20)]);
1580    }
1581
1582    #[test]
1583    fn test_batch_adapter_factory_identity() {
1584        // When source and target schemas are identical, should pass through efficiently
1585        let schema = Arc::new(Schema::new(vec![
1586            Field::new("a", DataType::Int32, false),
1587            Field::new("b", DataType::Utf8, true),
1588        ]));
1589
1590        let factory = BatchAdapterFactory::new(Arc::clone(&schema));
1591        let adapter = factory.make_adapter(&schema).unwrap();
1592
1593        let batch = RecordBatch::try_new(
1594            Arc::clone(&schema),
1595            vec![
1596                Arc::new(Int32Array::from(vec![1, 2, 3])),
1597                Arc::new(StringArray::from(vec!["a", "b", "c"])),
1598            ],
1599        )
1600        .unwrap();
1601
1602        let adapted = adapter.adapt_batch(&batch).unwrap();
1603
1604        assert_eq!(adapted.num_columns(), 2);
1605        assert_eq!(adapted.schema().field(0).data_type(), &DataType::Int32);
1606        assert_eq!(adapted.schema().field(1).data_type(), &DataType::Utf8);
1607    }
1608
1609    #[test]
1610    fn test_batch_adapter_factory_reuse() {
1611        // Factory can create multiple adapters for different source schemas
1612        let target_schema = Arc::new(Schema::new(vec![
1613            Field::new("x", DataType::Int64, false),
1614            Field::new("y", DataType::Utf8, true),
1615        ]));
1616
1617        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1618
1619        // First source schema
1620        let source1 = Arc::new(Schema::new(vec![
1621            Field::new("x", DataType::Int32, false),
1622            Field::new("y", DataType::Utf8, true),
1623        ]));
1624        let adapter1 = factory.make_adapter(&source1).unwrap();
1625
1626        // Second source schema (different order)
1627        let source2 = Arc::new(Schema::new(vec![
1628            Field::new("y", DataType::Utf8, true),
1629            Field::new("x", DataType::Int64, false),
1630        ]));
1631        let adapter2 = factory.make_adapter(&source2).unwrap();
1632
1633        // Both should work correctly
1634        assert!(format!("{adapter1:?}").contains("BatchAdapter"));
1635        assert!(format!("{adapter2:?}").contains("BatchAdapter"));
1636    }
1637
1638    #[test]
1639    fn test_rewrite_column_index_and_type_mismatch() {
1640        let physical_schema = Schema::new(vec![
1641            Field::new("b", DataType::Utf8, true),
1642            Field::new("a", DataType::Int32, false), // Index 1
1643        ]);
1644
1645        let logical_schema = Schema::new(vec![
1646            Field::new("a", DataType::Int64, false), // Index 0, Different Type
1647            Field::new("b", DataType::Utf8, true),
1648        ]);
1649
1650        let adapter = DefaultPhysicalExprAdapterFactory
1651            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1652            .unwrap();
1653
1654        // Logical column "a" is at index 0
1655        let column_expr = Arc::new(Column::new("a", 0));
1656
1657        let result = adapter.rewrite(column_expr).unwrap();
1658
1659        // Should be a CastExpr
1660        let cast_expr = assert_cast_expr(&result);
1661
1662        // Verify the inner column points to the correct physical index (1)
1663        assert_cast_input_column(cast_expr, "a", 1);
1664
1665        // Verify cast types
1666        assert_eq!(
1667            cast_expr.data_type(&Schema::empty()).unwrap(),
1668            DataType::Int64
1669        );
1670    }
1671
1672    #[test]
1673    fn test_rewrite_resolves_physical_column_by_name_before_casting() {
1674        let (logical_schema, physical_schema) = stale_index_cast_schemas();
1675        let adapter = DefaultPhysicalExprAdapterFactory
1676            .create(logical_schema, physical_schema)
1677            .unwrap();
1678
1679        // Deliberately provide the wrong index for column `a`.
1680        // Regression: this must still resolve against physical field `a` by name.
1681        let rewritten = adapter.rewrite(Arc::new(Column::new("a", 0))).unwrap();
1682        let cast_expr = assert_cast_expr(&rewritten);
1683        assert_cast_input_column(cast_expr, "a", 1);
1684        assert_eq!(cast_expr.target_field().data_type(), &DataType::Int64);
1685    }
1686}
datafusion_physical_expr_adapter/schema_rewriter.rs

datafusion_physical_expr_adapter/
schema_rewriter.rs
