Skip to main content

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, Fields, 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::expressions::Literal;
38use datafusion_physical_expr::projection::{ProjectionExprs, Projector};
39use datafusion_physical_expr::{
40    ScalarFunctionExpr,
41    expressions::{self, CastExpr, Column},
42};
43use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
44use itertools::Itertools;
45
46/// Replace column references in the given physical expression with literal values.
47///
48/// Some use cases for this include:
49/// - Partition column pruning: When scanning partitioned data, partition column references
50///   can be replaced with their literal values for the specific partition being scanned.
51/// - Constant folding: In some cases, columns that can be proven to be constant
52///   from statistical analysis may be replaced with their literal values to optimize expression evaluation.
53/// - Filling in non-null default values: in a custom [`PhysicalExprAdapter`] implementation,
54///   column references can be replaced with default literal values instead of nulls.
55///
56/// # Arguments
57/// - `expr`: The physical expression in which to replace column references.
58/// - `replacements`: A mapping from column names to their corresponding literal `ScalarValue`s.
59///   Accepts various HashMap types including `HashMap<&str, &ScalarValue>`,
60///   `HashMap<String, ScalarValue>`, `HashMap<String, &ScalarValue>`, etc.
61///
62/// # Returns
63/// - `Result<Arc<dyn PhysicalExpr>>`: The rewritten physical expression with columns replaced by literals.
64pub fn replace_columns_with_literals<K, V>(
65    expr: Arc<dyn PhysicalExpr>,
66    replacements: &HashMap<K, V>,
67) -> Result<Arc<dyn PhysicalExpr>>
68where
69    K: Borrow<str> + Eq + Hash,
70    V: Borrow<ScalarValue>,
71{
72    expr.transform_down(|expr| {
73        if let Some(column) = expr.downcast_ref::<Column>()
74            && let Some(replacement_value) = replacements.get(column.name())
75        {
76            return Ok(Transformed::yes(expressions::lit(
77                replacement_value.borrow().clone(),
78            )));
79        }
80        Ok(Transformed::no(expr))
81    })
82    .data()
83}
84
85/// Trait for adapting [`PhysicalExpr`] expressions to match a target schema.
86///
87/// This is used in file scans to rewrite expressions so that they can be
88/// evaluated against the physical schema of the file being scanned. It allows
89/// for handling differences between logical and physical schemas, such as type
90/// mismatches or missing columns common in [Schema evolution] scenarios.
91///
92/// [Schema evolution]: https://www.dremio.com/wiki/schema-evolution/
93///
94/// ## Default Implementations
95///
96/// The default implementation [`DefaultPhysicalExprAdapter`]  handles common
97/// cases.
98///
99/// ## Custom Implementations
100///
101/// You can create a custom implementation of this trait to handle specific rewriting logic.
102/// For example, to fill in missing columns with default values instead of nulls:
103///
104/// ```rust
105/// use datafusion_physical_expr_adapter::{PhysicalExprAdapter, PhysicalExprAdapterFactory};
106/// use arrow::datatypes::{Schema, Field, DataType, FieldRef, SchemaRef};
107/// use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
108/// use datafusion_common::{Result, ScalarValue, tree_node::{Transformed, TransformedResult, TreeNode}};
109/// use datafusion_physical_expr::expressions::{self, Column};
110/// use std::sync::Arc;
111///
112/// #[derive(Debug)]
113/// pub struct CustomPhysicalExprAdapter {
114///     logical_file_schema: SchemaRef,
115///     physical_file_schema: SchemaRef,
116/// }
117///
118/// impl PhysicalExprAdapter for CustomPhysicalExprAdapter {
119///     fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
120///         expr.transform(|expr| {
121///             if let Some(column) = expr.downcast_ref::<Column>() {
122///                 // Check if the column exists in the physical schema
123///                 if self.physical_file_schema.index_of(column.name()).is_err() {
124///                     // If the column is missing, fill it with a default value instead of null
125///                     // The default value could be stored in the table schema's column metadata for example.
126///                     let default_value = ScalarValue::Int32(Some(0));
127///                     return Ok(Transformed::yes(expressions::lit(default_value)));
128///                 }
129///             }
130///             // If the column exists, return it as is
131///             Ok(Transformed::no(expr))
132///         }).data()
133///     }
134/// }
135///
136/// #[derive(Debug)]
137/// pub struct CustomPhysicalExprAdapterFactory;
138///
139/// impl PhysicalExprAdapterFactory for CustomPhysicalExprAdapterFactory {
140///     fn create(
141///         &self,
142///         logical_file_schema: SchemaRef,
143///         physical_file_schema: SchemaRef,
144///     ) -> Result<Arc<dyn PhysicalExprAdapter>> {
145///         Ok(Arc::new(CustomPhysicalExprAdapter {
146///             logical_file_schema,
147///             physical_file_schema,
148///         }))
149///     }
150/// }
151/// ```
152pub trait PhysicalExprAdapter: Send + Sync + std::fmt::Debug {
153    /// Rewrite a physical expression to match the target schema.
154    ///
155    /// This method should return a transformed expression that matches the target schema.
156    ///
157    /// Arguments:
158    /// - `expr`: The physical expression to rewrite.
159    /// - `logical_file_schema`: The logical schema of the table being queried, excluding any partition columns.
160    /// - `physical_file_schema`: The physical schema of the file being scanned.
161    /// - `partition_values`: Optional partition values to use for rewriting partition column references.
162    ///   These are handled as if they were columns appended onto the logical file schema.
163    ///
164    /// Returns:
165    /// - `Arc<dyn PhysicalExpr>`: The rewritten physical expression that can be evaluated against the physical schema.
166    ///
167    /// See Also:
168    /// - [`replace_columns_with_literals`]: for replacing partition column references with their literal values.
169    fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>>;
170}
171
172/// Creates instances of [`PhysicalExprAdapter`] for given logical and physical schemas.
173///
174/// See [`DefaultPhysicalExprAdapterFactory`] for the default implementation.
175pub trait PhysicalExprAdapterFactory: Send + Sync + std::fmt::Debug {
176    /// Create a new instance of the physical expression adapter.
177    fn create(
178        &self,
179        logical_file_schema: SchemaRef,
180        physical_file_schema: SchemaRef,
181    ) -> Result<Arc<dyn PhysicalExprAdapter>>;
182}
183
184#[derive(Debug, Clone)]
185pub struct DefaultPhysicalExprAdapterFactory;
186
187impl PhysicalExprAdapterFactory for DefaultPhysicalExprAdapterFactory {
188    fn create(
189        &self,
190        logical_file_schema: SchemaRef,
191        physical_file_schema: SchemaRef,
192    ) -> Result<Arc<dyn PhysicalExprAdapter>> {
193        Ok(Arc::new(DefaultPhysicalExprAdapter {
194            logical_file_schema,
195            physical_file_schema,
196        }))
197    }
198}
199
200/// Default implementation of [`PhysicalExprAdapter`] for rewriting physical
201/// expressions to match different schemas.
202///
203/// ## Overview
204///
205///  [`DefaultPhysicalExprAdapter`] rewrites physical expressions to match
206///  different schemas, including:
207///
208/// - **Type casting**: When logical and physical schemas have different types, expressions are
209///   automatically wrapped with cast operations. For example, `lit(ScalarValue::Int32(123)) = int64_column`
210///   gets rewritten to `lit(ScalarValue::Int32(123)) = cast(int64_column, 'Int32')`.
211///   Note that this does not attempt to simplify such expressions - that is done by shared simplifiers.
212///
213/// - **Missing columns**: When a column exists in the logical schema but not in the physical schema,
214///   references to it are replaced with null literals.
215///
216/// - **Struct field access**: Expressions like `struct_column.field_that_is_missing_in_schema` are
217///   rewritten to `null` when the field doesn't exist in the physical schema.
218///
219/// - **Default column values**: Partition column references can be replaced with their literal values
220///   when scanning specific partitions. See [`replace_columns_with_literals`] for more details.
221///
222/// # Example
223///
224/// ```rust
225/// # use datafusion_physical_expr_adapter::{DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory};
226/// # use arrow::datatypes::Schema;
227/// # use std::sync::Arc;
228/// #
229/// # fn example(
230/// #     predicate: std::sync::Arc<dyn datafusion_physical_expr_common::physical_expr::PhysicalExpr>,
231/// #     physical_file_schema: &Schema,
232/// #     logical_file_schema: &Schema,
233/// # ) -> datafusion_common::Result<()> {
234/// let factory = DefaultPhysicalExprAdapterFactory;
235/// let adapter =
236///     factory.create(Arc::new(logical_file_schema.clone()), Arc::new(physical_file_schema.clone()))?;
237/// let adapted_predicate = adapter.rewrite(predicate)?;
238/// # Ok(())
239/// # }
240/// ```
241#[derive(Debug, Clone)]
242pub struct DefaultPhysicalExprAdapter {
243    logical_file_schema: SchemaRef,
244    physical_file_schema: SchemaRef,
245}
246
247impl DefaultPhysicalExprAdapter {
248    /// Create a new instance of the default physical expression adapter.
249    ///
250    /// This adapter rewrites expressions to match the physical schema of the file being scanned,
251    /// handling type mismatches and missing columns by filling them with default values.
252    pub fn new(logical_file_schema: SchemaRef, physical_file_schema: SchemaRef) -> Self {
253        Self {
254            logical_file_schema,
255            physical_file_schema,
256        }
257    }
258}
259
260impl PhysicalExprAdapter for DefaultPhysicalExprAdapter {
261    fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
262        let rewriter = DefaultPhysicalExprAdapterRewriter {
263            logical_file_schema: Arc::clone(&self.logical_file_schema),
264            physical_file_schema: Arc::clone(&self.physical_file_schema),
265        };
266        expr.transform(|expr| rewriter.rewrite_expr(Arc::clone(&expr)))
267            .data()
268    }
269}
270
271struct DefaultPhysicalExprAdapterRewriter {
272    logical_file_schema: SchemaRef,
273    physical_file_schema: SchemaRef,
274}
275
276/// Outcome of walking a `get_field` key path through nested struct fields.
277enum FieldPathResolution<'a> {
278    /// The leaf field the path points at.
279    Found(&'a FieldRef),
280    /// Some key along the path does not exist, so the access reads as null.
281    Missing,
282    /// An intermediate field is not a struct, so the path cannot be resolved
283    /// statically.
284    NotAStruct,
285}
286
287/// Follow a `get_field` key path (`['a', 'b']` for `s['a']['b']`) through
288/// nested struct fields.
289///
290/// The first key is taken separately from the rest so that the "at least one
291/// key" invariant is carried by the signature: there is no empty path to
292/// resolve.
293fn resolve_field_path<'a>(
294    fields: &'a Fields,
295    field_name: &str,
296    rest: &[&str],
297) -> FieldPathResolution<'a> {
298    let Some(field) = fields.iter().find(|f| f.name() == field_name) else {
299        return FieldPathResolution::Missing;
300    };
301    let Some((next_field_name, rest)) = rest.split_first() else {
302        return FieldPathResolution::Found(field);
303    };
304    match field.data_type() {
305        DataType::Struct(nested_fields) => {
306            resolve_field_path(nested_fields, next_field_name, rest)
307        }
308        _ => FieldPathResolution::NotAStruct,
309    }
310}
311
312impl DefaultPhysicalExprAdapterRewriter {
313    fn rewrite_expr(
314        &self,
315        expr: Arc<dyn PhysicalExpr>,
316    ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
317        if let Some(transformed) = self.try_rewrite_struct_field_access(&expr)? {
318            return Ok(Transformed::yes(transformed));
319        }
320
321        if let Some(transformed) = self.try_narrow_struct_cast(&expr)? {
322            return Ok(Transformed::yes(transformed));
323        }
324
325        if let Some(column) = expr.downcast_ref::<Column>() {
326            return self.rewrite_column(Arc::clone(&expr), column);
327        }
328
329        Ok(Transformed::no(expr))
330    }
331
332    /// Rewrite `get_field(cast(s AS Struct<..>), 'f')` into
333    /// `cast(get_field(s, 'f') AS <type of f>)`.
334    ///
335    /// Expressions are rewritten bottom-up, so by the time we reach a
336    /// `get_field` node its struct argument has already been wrapped in a cast
337    /// by [`Self::rewrite_column`] whenever the logical and physical struct
338    /// types differ.
339    ///
340    /// Narrowing that cast is worthwhile for two reasons:
341    ///
342    /// 1. Reading one field should not cost a whole struct. The wide form
343    ///    casts every field of the column — including ones the query never
344    ///    reads — to produce a value that is immediately discarded except for
345    ///    one field.
346    /// 2. It keeps the column visible. Consumers throughout the codebase
347    ///    pattern match on `get_field(column, 'f')` to recognise a struct
348    ///    field access; a cast between the `get_field` and its column defeats
349    ///    that match, and each such consumer then falls back to whatever it
350    ///    does for an unrecognised expression.
351    ///
352    /// The Parquet scan is one such consumer, and the reason this is a
353    /// correctness fix rather than only an optimisation: it decides at
354    /// planning time, against the table schema, that a struct-field predicate
355    /// can be evaluated as a row filter, and reports the predicate as fully
356    /// handled. See <https://github.com/apache/datafusion/issues/24109>.
357    ///
358    /// Fixing it here rather than teaching that one consumer to see through
359    /// casts is deliberate: the adapter is where the obscuring cast is
360    /// introduced, so every consumer benefits, and no consumer has to loosen
361    /// its pattern to accept arbitrary casts between a `get_field` and its
362    /// column.
363    ///
364    /// `get_field` also has a flattened multi-key form: `s['a']['b']` is
365    /// simplified to `get_field(s, 'a', 'b')`, so the whole field path is
366    /// resolved here rather than just the first key.
367    ///
368    /// Only struct casts are narrowed. `get_field` on a Map column performs a
369    /// runtime key lookup rather than a schema-level field access, so the map
370    /// value must keep its cast.
371    fn try_narrow_struct_cast(
372        &self,
373        expr: &Arc<dyn PhysicalExpr>,
374    ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
375        let Some(get_field_expr) =
376            ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(expr.as_ref())
377        else {
378            return Ok(None);
379        };
380        let Some((source_expr, field_name_exprs)) = get_field_expr.args().split_first()
381        else {
382            return Ok(None);
383        };
384        let Some(cast) = source_expr.downcast_ref::<CastExpr>() else {
385            return Ok(None);
386        };
387
388        // Every key has to be a string literal, otherwise the leaf field
389        // cannot be resolved statically.
390        let mut field_path = Vec::with_capacity(field_name_exprs.len());
391        for field_name_expr in field_name_exprs {
392            let Some(field_name) = field_name_expr
393                .downcast_ref::<Literal>()
394                .and_then(|lit| lit.value().try_as_str().flatten())
395            else {
396                return Ok(None);
397            };
398            field_path.push(field_name);
399        }
400        // A `get_field` with no keys is not a field access we can narrow.
401        let Some((first_key, rest_keys)) = field_path.split_first() else {
402            return Ok(None);
403        };
404
405        let DataType::Struct(logical_struct_fields) = cast.target_field().data_type()
406        else {
407            return Ok(None);
408        };
409        let FieldPathResolution::Found(logical_struct_field) =
410            resolve_field_path(logical_struct_fields, first_key, rest_keys)
411        else {
412            return Ok(None);
413        };
414
415        let inner = cast.expr();
416        let DataType::Struct(physical_struct_fields) =
417            inner.data_type(&self.physical_file_schema)?
418        else {
419            return Ok(None);
420        };
421        let physical_struct_field =
422            match resolve_field_path(&physical_struct_fields, first_key, rest_keys) {
423                FieldPathResolution::Found(field) => field,
424                FieldPathResolution::Missing => {
425                    // The file does not have this field at all, so reading it
426                    // yields null. Note that the cast would have produced the
427                    // same value: struct casts fill missing target fields with
428                    // nulls.
429                    let null_value =
430                        ScalarValue::Null.cast_to(logical_struct_field.data_type())?;
431                    return Ok(Some(Arc::new(Literal::new_with_metadata(
432                        null_value,
433                        Some(FieldMetadata::from(logical_struct_field.as_ref())),
434                    ))));
435                }
436                FieldPathResolution::NotAStruct => return Ok(None),
437            };
438
439        // Rebuild `get_field` over the uncast struct so its return field is
440        // recomputed from the physical field type.
441        let mut args = Vec::with_capacity(get_field_expr.args().len());
442        args.push(Arc::clone(inner));
443        args.extend(field_name_exprs.iter().map(Arc::clone));
444        let extracted = Arc::new(ScalarFunctionExpr::try_new(
445            Arc::new(get_field_expr.fun().clone()),
446            args,
447            &self.physical_file_schema,
448            Arc::new(get_field_expr.config_options().clone()),
449        )?) as Arc<dyn PhysicalExpr>;
450
451        if physical_struct_field == logical_struct_field {
452            return Ok(Some(extracted));
453        }
454        Ok(Some(Arc::new(CastExpr::new_with_target_field(
455            extracted,
456            Arc::clone(logical_struct_field),
457            Some(cast.cast_options().clone()),
458        ))))
459    }
460
461    /// Attempt to rewrite struct field access expressions to return null if the field does not exist in the physical schema.
462    /// Note that this does *not* handle nested struct fields, only top-level struct field access.
463    /// See <https://github.com/apache/datafusion/issues/17114> for more details.
464    fn try_rewrite_struct_field_access(
465        &self,
466        expr: &Arc<dyn PhysicalExpr>,
467    ) -> Result<Option<Arc<dyn PhysicalExpr>>> {
468        let get_field_expr =
469            match ScalarFunctionExpr::try_downcast_func::<GetFieldFunc>(expr.as_ref()) {
470                Some(expr) => expr,
471                None => return Ok(None),
472            };
473
474        let source_expr = match get_field_expr.args().first() {
475            Some(expr) => expr,
476            None => return Ok(None),
477        };
478
479        let field_name_expr = match get_field_expr.args().get(1) {
480            Some(expr) => expr,
481            None => return Ok(None),
482        };
483
484        let lit = match field_name_expr.downcast_ref::<Literal>() {
485            Some(lit) => lit,
486            None => return Ok(None),
487        };
488
489        let field_name = match lit.value().try_as_str().flatten() {
490            Some(name) => name,
491            None => return Ok(None),
492        };
493
494        let column = match source_expr.downcast_ref::<Column>() {
495            Some(column) => column,
496            None => return Ok(None),
497        };
498
499        let physical_field =
500            match self.physical_file_schema.field_with_name(column.name()) {
501                Ok(field) => field,
502                Err(_) => return Ok(None),
503            };
504
505        let physical_struct_fields = match physical_field.data_type() {
506            DataType::Struct(fields) => fields,
507            _ => return Ok(None),
508        };
509
510        if physical_struct_fields
511            .iter()
512            .any(|f| f.name() == field_name)
513        {
514            return Ok(None);
515        }
516
517        let logical_field = match self.logical_file_schema.field_with_name(column.name())
518        {
519            Ok(field) => field,
520            Err(_) => return Ok(None),
521        };
522
523        let logical_struct_fields = match logical_field.data_type() {
524            DataType::Struct(fields) => fields,
525            _ => return Ok(None),
526        };
527
528        let logical_struct_field = match logical_struct_fields
529            .iter()
530            .find(|f| f.name() == field_name)
531        {
532            Some(field) => field,
533            None => return Ok(None),
534        };
535
536        let null_value = ScalarValue::Null.cast_to(logical_struct_field.data_type())?;
537        Ok(Some(Arc::new(Literal::new_with_metadata(
538            null_value,
539            Some(FieldMetadata::from(logical_struct_field.as_ref())),
540        ))))
541    }
542
543    fn rewrite_column(
544        &self,
545        expr: Arc<dyn PhysicalExpr>,
546        column: &Column,
547    ) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
548        // Get the logical field for this column if it exists in the logical schema
549        let logical_field = match self.logical_file_schema.field_with_name(column.name())
550        {
551            Ok(field) => field,
552            Err(e) => {
553                // This can be hit if a custom rewrite injected a reference to a column that doesn't exist in the logical schema.
554                // For example, a pre-computed column that is kept only in the physical schema.
555                // If the column exists in the physical schema, we can still use it.
556                if let Ok(physical_field) =
557                    self.physical_file_schema.field_with_name(column.name())
558                {
559                    // If the column exists in the physical schema, we can use it in place of the logical column.
560                    // This is nice to users because if they do a rewrite that results in something like `physical_int32_col = 123u64`
561                    // we'll at least handle the casts for them.
562                    physical_field
563                } else {
564                    // A completely unknown column that doesn't exist in either schema!
565                    // This should probably never be hit unless something upstream broke, but nonetheless it's better
566                    // for us to return a handleable error than to panic / do something unexpected.
567                    return Err(e.into());
568                }
569            }
570        };
571
572        let Some((resolved_column, physical_field)) =
573            self.resolve_physical_column(column)?
574        else {
575            if !logical_field.is_nullable() {
576                return exec_err!(
577                    "Non-nullable column '{}' is missing from the physical schema",
578                    column.name()
579                );
580            }
581            // If the column is missing from the physical schema fill it in with nulls.
582            // For a different behavior, provide a custom `PhysicalExprAdapter` implementation.
583            let null_value = ScalarValue::Null.cast_to(logical_field.data_type())?;
584            return Ok(Transformed::yes(Arc::new(Literal::new_with_metadata(
585                null_value,
586                Some(FieldMetadata::from(logical_field)),
587            ))));
588        };
589
590        let fields_match = logical_field == physical_field.as_ref();
591        if fields_match {
592            if resolved_column.index() == column.index() {
593                return Ok(Transformed::no(expr));
594            }
595
596            // If the fields match (including metadata/nullability), we can use the column as is
597            return Ok(Transformed::yes(Arc::new(resolved_column)));
598        }
599
600        // We need a cast expression whenever the logical and physical fields differ,
601        // whether that difference is only metadata/nullability or also data type.
602        // TODO: add optimization to move the cast from the column to literal expressions in the case of `col = 123`
603        // since that's much cheaper to evaluate.
604        // See https://github.com/apache/datafusion/issues/15780#issuecomment-2824716928
605        validate_data_type_compatibility(
606            resolved_column.name(),
607            physical_field.data_type(),
608            logical_field.data_type(),
609        )
610        .map_err(|e| {
611            DataFusionError::Execution(format!(
612                "Cannot cast column '{}' from '{}' (physical data type) to '{}' (logical data type): {e}",
613                resolved_column.name(),
614                physical_field.data_type(),
615                logical_field.data_type()
616            ))
617        })?;
618
619        Ok(Transformed::yes(Arc::new(CastExpr::new_with_target_field(
620            Arc::new(resolved_column),
621            Arc::new(logical_field.clone()),
622            None,
623        ))))
624    }
625
626    /// Resolves a logical column to the corresponding physical column and field.
627    fn resolve_physical_column(
628        &self,
629        column: &Column,
630    ) -> Result<Option<(Column, FieldRef)>> {
631        // The physical schema adaptation step intentionally resolves columns by **name first**
632        // rather than trusting the incoming index. This mirrors what the old refactoring
633        // did before `resolve_physical_column()` was extracted: the planner might hand us a
634        // `Column` whose `index` field is stale (e.g. after projection/rename rewrites), so
635        // resolving by name ensures we match the correct physical slot. Once we know the
636        // proper index we rebuild the `Column` with `new_with_schema` so callers can rely
637        // on `column.index()` later without having to re-query the schema.
638        let Ok(physical_column_index) = self.physical_file_schema.index_of(column.name())
639        else {
640            return Ok(None);
641        };
642
643        let column = if column.index() == physical_column_index {
644            column.clone()
645        } else {
646            Column::new_with_schema(column.name(), self.physical_file_schema.as_ref())?
647        };
648
649        let physical_field = Arc::new(
650            self.physical_file_schema
651                .field(physical_column_index)
652                .clone(),
653        );
654
655        Ok(Some((column, physical_field)))
656    }
657}
658
659/// Factory for creating [`BatchAdapter`] instances to adapt record batches
660/// to a target schema.
661///
662/// This binds a target schema and allows creating adapters for different source schemas.
663/// It handles:
664/// - **Column reordering**: Columns are reordered to match the target schema
665/// - **Type casting**: Automatic type conversion (e.g., Int32 to Int64)
666/// - **Missing columns**: Nullable columns missing from source are filled with nulls
667/// - **Struct field adaptation**: Nested struct fields are recursively adapted
668///
669/// ## Examples
670///
671/// ```rust
672/// use arrow::array::{Int32Array, Int64Array, StringArray, RecordBatch};
673/// use arrow::datatypes::{DataType, Field, Schema};
674/// use datafusion_physical_expr_adapter::BatchAdapterFactory;
675/// use std::sync::Arc;
676///
677/// // Target schema has different column order and types
678/// let target_schema = Arc::new(Schema::new(vec![
679///     Field::new("name", DataType::Utf8, true),
680///     Field::new("id", DataType::Int64, false),    // Int64 in target
681///     Field::new("score", DataType::Float64, true), // Missing from source
682/// ]));
683///
684/// // Source schema has different column order and Int32 for id
685/// let source_schema = Arc::new(Schema::new(vec![
686///     Field::new("id", DataType::Int32, false),    // Int32 in source
687///     Field::new("name", DataType::Utf8, true),
688///     // Note: 'score' column is missing from source
689/// ]));
690///
691/// // Create factory with target schema
692/// let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
693///
694/// // Create adapter for this specific source schema
695/// let adapter = factory.make_adapter(&source_schema).unwrap();
696///
697/// // Create a source batch
698/// let source_batch = RecordBatch::try_new(
699///     source_schema,
700///     vec![
701///         Arc::new(Int32Array::from(vec![1, 2, 3])),
702///         Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol"])),
703///     ],
704/// ).unwrap();
705///
706/// // Adapt the batch to match target schema
707/// let adapted = adapter.adapt_batch(&source_batch).unwrap();
708///
709/// assert_eq!(adapted.num_columns(), 3);
710/// assert_eq!(adapted.column(0).data_type(), &DataType::Utf8);   // name
711/// assert_eq!(adapted.column(1).data_type(), &DataType::Int64);  // id (cast from Int32)
712/// assert_eq!(adapted.column(2).data_type(), &DataType::Float64); // score (filled with nulls)
713/// ```
714#[derive(Debug)]
715pub struct BatchAdapterFactory {
716    target_schema: SchemaRef,
717    expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
718}
719
720impl BatchAdapterFactory {
721    /// Create a new [`BatchAdapterFactory`] with the given target schema.
722    pub fn new(target_schema: SchemaRef) -> Self {
723        let expr_adapter_factory = Arc::new(DefaultPhysicalExprAdapterFactory);
724        Self {
725            target_schema,
726            expr_adapter_factory,
727        }
728    }
729
730    /// Set a custom [`PhysicalExprAdapterFactory`] to use when adapting expressions.
731    ///
732    /// Use this to customize behavior when adapting batches, e.g. to fill in missing values
733    /// with defaults instead of nulls.
734    ///
735    /// See [`PhysicalExprAdapter`] for more details.
736    pub fn with_adapter_factory(
737        self,
738        factory: Arc<dyn PhysicalExprAdapterFactory>,
739    ) -> Self {
740        Self {
741            expr_adapter_factory: factory,
742            ..self
743        }
744    }
745
746    /// Create a new [`BatchAdapter`] for the given source schema.
747    ///
748    /// Batches fed into this [`BatchAdapter`] *must* conform to the source schema,
749    /// no validation is performed at runtime to minimize overheads.
750    pub fn make_adapter(&self, source_schema: &SchemaRef) -> Result<BatchAdapter> {
751        let expr_adapter = self
752            .expr_adapter_factory
753            .create(Arc::clone(&self.target_schema), Arc::clone(source_schema))?;
754
755        let simplifier = PhysicalExprSimplifier::new(&self.target_schema);
756
757        let projection = ProjectionExprs::from_indices(
758            &(0..self.target_schema.fields().len()).collect_vec(),
759            &self.target_schema,
760        );
761
762        let adapted = projection
763            .try_map_exprs(|e| simplifier.simplify(expr_adapter.rewrite(e)?))?;
764        let projector = adapted.make_projector(source_schema)?;
765
766        Ok(BatchAdapter { projector })
767    }
768}
769
770/// Adapter for transforming record batches to match a target schema.
771///
772/// Create instances via [`BatchAdapterFactory`].
773///
774/// ## Performance
775///
776/// The adapter pre-computes the projection expressions during creation,
777/// so the [`adapt_batch`](BatchAdapter::adapt_batch) call is efficient and suitable
778/// for use in hot paths like streaming file scans.
779#[derive(Debug)]
780pub struct BatchAdapter {
781    projector: Projector,
782}
783
784impl BatchAdapter {
785    /// Adapt the given record batch to match the target schema.
786    ///
787    /// The input batch *must* conform to the source schema used when
788    /// creating this adapter.
789    pub fn adapt_batch(&self, batch: &RecordBatch) -> Result<RecordBatch> {
790        self.projector.project_batch(batch)
791    }
792}
793
794#[cfg(test)]
795mod tests {
796    use super::*;
797    use arrow::array::{
798        Array, BooleanArray, GenericListArray, Int32Array, Int64Array, RecordBatch,
799        RecordBatchOptions, StringArray, StringViewArray, StructArray, record_batch,
800    };
801    use arrow::datatypes as arrow_schema;
802    use arrow::datatypes::{Field, Fields, Schema};
803    use datafusion_common::assert_contains;
804    use datafusion_expr::Operator;
805    use datafusion_physical_expr::expressions::{Column, Literal, col};
806
807    fn assert_cast_expr(expr: &Arc<dyn PhysicalExpr>) -> &CastExpr {
808        expr.downcast_ref::<CastExpr>().expect("Expected CastExpr")
809    }
810
811    fn assert_cast_input_column(cast_expr: &CastExpr, name: &str, index: usize) {
812        let inner_col = cast_expr
813            .expr()
814            .downcast_ref::<Column>()
815            .expect("Expected inner Column");
816        assert_eq!(inner_col.name(), name);
817        assert_eq!(inner_col.index(), index);
818    }
819
820    fn stale_index_cast_schemas() -> (SchemaRef, SchemaRef) {
821        let physical_schema = Arc::new(Schema::new(vec![
822            Field::new("b", DataType::Binary, true),
823            Field::new("a", DataType::Int32, false),
824        ]));
825
826        let logical_schema = Arc::new(Schema::new(vec![
827            Field::new("a", DataType::Int64, false),
828            Field::new("b", DataType::Binary, true),
829        ]));
830
831        (logical_schema, physical_schema)
832    }
833
834    fn create_test_schema() -> (Schema, Schema) {
835        let physical_schema = Schema::new(vec![
836            Field::new("a", DataType::Int32, false),
837            Field::new("b", DataType::Utf8, true),
838        ]);
839
840        let logical_schema = Schema::new(vec![
841            Field::new("a", DataType::Int64, false), // Different type
842            Field::new("b", DataType::Utf8, true),
843            Field::new("c", DataType::Float64, true), // Missing from physical
844        ]);
845
846        (physical_schema, logical_schema)
847    }
848
849    #[test]
850    fn test_rewrite_column_with_type_cast() {
851        let (physical_schema, logical_schema) = create_test_schema();
852
853        let factory = DefaultPhysicalExprAdapterFactory;
854        let adapter = factory
855            .create(Arc::new(logical_schema), Arc::new(physical_schema))
856            .unwrap();
857        let column_expr = Arc::new(Column::new("a", 0));
858
859        let result = adapter.rewrite(column_expr).unwrap();
860
861        // Should be wrapped in a cast expression
862        assert!(result.downcast_ref::<CastExpr>().is_some());
863    }
864
865    #[test]
866    fn test_rewrite_column_with_metadata_or_nullability_mismatch() -> Result<()> {
867        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
868        let logical_schema =
869            Schema::new(vec![Field::new("a", DataType::Int64, false).with_metadata(
870                HashMap::from([("logical_meta".to_string(), "1".to_string())]),
871            )]);
872
873        let factory = DefaultPhysicalExprAdapterFactory;
874        let adapter = factory
875            .create(Arc::new(logical_schema), Arc::new(physical_schema.clone()))
876            .unwrap();
877
878        let result = adapter.rewrite(Arc::new(Column::new("a", 0)))?;
879
880        // Ensure the expression preserves the logical field nullability/metadata.
881        let return_field = result.return_field(physical_schema.as_ref())?;
882        assert_eq!(return_field.data_type(), &DataType::Int64);
883        assert!(!return_field.is_nullable());
884        assert_eq!(
885            return_field
886                .metadata()
887                .get("logical_meta")
888                .map(String::as_str),
889            Some("1")
890        );
891
892        Ok(())
893    }
894
895    #[test]
896    fn test_rewrite_multi_column_expr_with_type_cast() {
897        let (physical_schema, logical_schema) = create_test_schema();
898        let factory = DefaultPhysicalExprAdapterFactory;
899        let adapter = factory
900            .create(Arc::new(logical_schema), Arc::new(physical_schema))
901            .unwrap();
902
903        // Create a complex expression: (a + 5) OR (c > 0.0) that tests the recursive case of the rewriter
904        let column_a = Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>;
905        let column_c = Arc::new(Column::new("c", 2)) as Arc<dyn PhysicalExpr>;
906        let expr = expressions::BinaryExpr::new(
907            Arc::clone(&column_a),
908            Operator::Plus,
909            Arc::new(Literal::new(ScalarValue::Int64(Some(5)))),
910        );
911        let expr = expressions::BinaryExpr::new(
912            Arc::new(expr),
913            Operator::Or,
914            Arc::new(expressions::BinaryExpr::new(
915                Arc::clone(&column_c),
916                Operator::Gt,
917                Arc::new(Literal::new(ScalarValue::Float64(Some(0.0)))),
918            )),
919        );
920
921        let result = adapter.rewrite(Arc::new(expr)).unwrap();
922        let outer = result
923            .downcast_ref::<expressions::BinaryExpr>()
924            .expect("Expected outer BinaryExpr");
925        assert_eq!(*outer.op(), Operator::Or);
926
927        let left = outer
928            .left()
929            .downcast_ref::<expressions::BinaryExpr>()
930            .expect("Expected left BinaryExpr");
931        assert_eq!(*left.op(), Operator::Plus);
932
933        let left_cast = assert_cast_expr(left.left());
934        assert_eq!(left_cast.target_field().data_type(), &DataType::Int64);
935        assert_cast_input_column(left_cast, "a", 0);
936
937        let right = outer
938            .right()
939            .downcast_ref::<expressions::BinaryExpr>()
940            .expect("Expected right BinaryExpr");
941        assert_eq!(*right.op(), Operator::Gt);
942        let null_literal = right
943            .left()
944            .downcast_ref::<Literal>()
945            .expect("Expected null literal");
946        assert_eq!(*null_literal.value(), ScalarValue::Float64(None));
947    }
948
949    #[test]
950    fn test_rewrite_struct_column_incompatible() {
951        let physical_schema = Schema::new(vec![Field::new(
952            "data",
953            DataType::Struct(vec![Field::new("field1", DataType::Binary, true)].into()),
954            true,
955        )]);
956
957        let logical_schema = Schema::new(vec![Field::new(
958            "data",
959            DataType::Struct(vec![Field::new("field1", DataType::Int32, true)].into()),
960            true,
961        )]);
962
963        let factory = DefaultPhysicalExprAdapterFactory;
964        let adapter = factory
965            .create(Arc::new(logical_schema), Arc::new(physical_schema))
966            .unwrap();
967        let column_expr = Arc::new(Column::new("data", 0));
968
969        let error_msg = adapter.rewrite(column_expr).unwrap_err().to_string();
970        // validate_struct_compatibility provides more specific error about which field can't be cast
971        assert_contains!(
972            error_msg,
973            "Cannot cast struct field 'field1' from type Binary to type Int32"
974        );
975    }
976
977    #[test]
978    fn test_rewrite_struct_compatible_cast() {
979        let physical_schema = Schema::new(vec![Field::new(
980            "data",
981            DataType::Struct(
982                vec![
983                    Field::new("id", DataType::Int32, false),
984                    Field::new("name", DataType::Utf8, true),
985                ]
986                .into(),
987            ),
988            false,
989        )]);
990
991        let logical_schema = Schema::new(vec![Field::new(
992            "data",
993            DataType::Struct(
994                vec![
995                    Field::new("id", DataType::Int64, false),
996                    Field::new("name", DataType::Utf8View, true),
997                ]
998                .into(),
999            ),
1000            false,
1001        )]);
1002
1003        let factory = DefaultPhysicalExprAdapterFactory;
1004        let adapter = factory
1005            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1006            .unwrap();
1007        let column_expr = Arc::new(Column::new("data", 0));
1008
1009        let result = adapter.rewrite(column_expr).unwrap();
1010
1011        let logical_struct_fields: Fields = vec![
1012            Field::new("id", DataType::Int64, false),
1013            Field::new("name", DataType::Utf8View, true),
1014        ]
1015        .into();
1016        let logical_field = Arc::new(Field::new(
1017            "data",
1018            DataType::Struct(logical_struct_fields),
1019            false,
1020        ));
1021
1022        let expected = Arc::new(CastExpr::new_with_target_field(
1023            Arc::new(Column::new("data", 0)),
1024            logical_field,
1025            None,
1026        )) as Arc<dyn PhysicalExpr>;
1027
1028        assert_eq!(result.to_string(), expected.to_string());
1029    }
1030
1031    #[test]
1032    fn test_rewrite_missing_column() -> Result<()> {
1033        let (physical_schema, logical_schema) = create_test_schema();
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("c", 2));
1040
1041        let result = adapter.rewrite(column_expr)?;
1042
1043        // Should be replaced with a literal null
1044        if let Some(literal) = result.downcast_ref::<Literal>() {
1045            assert_eq!(*literal.value(), ScalarValue::Float64(None));
1046        } else {
1047            panic!("Expected literal expression");
1048        }
1049
1050        Ok(())
1051    }
1052
1053    #[test]
1054    fn test_rewrite_missing_column_propagates_metadata() -> Result<()> {
1055        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1056        let logical_schema = Schema::new(vec![
1057            Field::new("a", DataType::Int32, false),
1058            Field::new("b", DataType::Utf8, true).with_metadata(HashMap::from([(
1059                "logical_meta".to_string(),
1060                "1".to_string(),
1061            )])),
1062        ]);
1063
1064        let factory = DefaultPhysicalExprAdapterFactory;
1065        let adapter = factory
1066            .create(Arc::new(logical_schema), Arc::new(physical_schema.clone()))
1067            .unwrap();
1068
1069        let result = adapter.rewrite(Arc::new(Column::new("b", 1)))?;
1070        let literal = result
1071            .downcast_ref::<Literal>()
1072            .expect("Expected literal expression");
1073
1074        assert_eq!(
1075            literal
1076                .return_field(physical_schema.as_ref())?
1077                .metadata()
1078                .get("logical_meta")
1079                .map(String::as_str),
1080            Some("1")
1081        );
1082        Ok(())
1083    }
1084
1085    #[test]
1086    fn test_rewrite_missing_column_non_nullable_error() {
1087        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1088        let logical_schema = Schema::new(vec![
1089            Field::new("a", DataType::Int64, false),
1090            Field::new("b", DataType::Utf8, false), // Missing and non-nullable
1091        ]);
1092
1093        let factory = DefaultPhysicalExprAdapterFactory;
1094        let adapter = factory
1095            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1096            .unwrap();
1097        let column_expr = Arc::new(Column::new("b", 1));
1098
1099        let error_msg = adapter.rewrite(column_expr).unwrap_err().to_string();
1100        assert_contains!(error_msg, "Non-nullable column 'b' is missing");
1101    }
1102
1103    #[test]
1104    fn test_rewrite_missing_column_nullable() {
1105        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1106        let logical_schema = Schema::new(vec![
1107            Field::new("a", DataType::Int64, false),
1108            Field::new("b", DataType::Utf8, true), // Missing but nullable
1109        ]);
1110
1111        let factory = DefaultPhysicalExprAdapterFactory;
1112        let adapter = factory
1113            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1114            .unwrap();
1115        let column_expr = Arc::new(Column::new("b", 1));
1116
1117        let result = adapter.rewrite(column_expr).unwrap();
1118
1119        let expected =
1120            Arc::new(Literal::new(ScalarValue::Utf8(None))) as Arc<dyn PhysicalExpr>;
1121
1122        assert_eq!(result.to_string(), expected.to_string());
1123    }
1124
1125    #[test]
1126    fn test_replace_columns_with_literals() -> Result<()> {
1127        let partition_value = ScalarValue::Utf8(Some("test_value".to_string()));
1128        let replacements = HashMap::from([("partition_col", &partition_value)]);
1129
1130        let column_expr =
1131            Arc::new(Column::new("partition_col", 0)) as Arc<dyn PhysicalExpr>;
1132        let result = replace_columns_with_literals(column_expr, &replacements)?;
1133
1134        // Should be replaced with the partition value
1135        let literal = result
1136            .downcast_ref::<Literal>()
1137            .expect("Expected literal expression");
1138        assert_eq!(*literal.value(), partition_value);
1139
1140        Ok(())
1141    }
1142
1143    #[test]
1144    fn test_replace_columns_with_literals_no_match() -> Result<()> {
1145        let value = ScalarValue::Utf8(Some("test_value".to_string()));
1146        let replacements = HashMap::from([("other_col", &value)]);
1147
1148        let column_expr =
1149            Arc::new(Column::new("partition_col", 0)) as Arc<dyn PhysicalExpr>;
1150        let result = replace_columns_with_literals(column_expr, &replacements)?;
1151
1152        assert!(result.downcast_ref::<Column>().is_some());
1153        Ok(())
1154    }
1155
1156    #[test]
1157    fn test_replace_columns_with_literals_nested_expr() -> Result<()> {
1158        let value_a = ScalarValue::Int64(Some(10));
1159        let value_b = ScalarValue::Int64(Some(20));
1160        let replacements = HashMap::from([("a", &value_a), ("b", &value_b)]);
1161
1162        let expr = Arc::new(expressions::BinaryExpr::new(
1163            Arc::new(Column::new("a", 0)),
1164            Operator::Plus,
1165            Arc::new(Column::new("b", 1)),
1166        )) as Arc<dyn PhysicalExpr>;
1167
1168        let result = replace_columns_with_literals(expr, &replacements)?;
1169        assert_eq!(result.to_string(), "10 + 20");
1170
1171        Ok(())
1172    }
1173
1174    #[test]
1175    fn test_rewrite_no_change_needed() -> Result<()> {
1176        let (physical_schema, logical_schema) = create_test_schema();
1177
1178        let factory = DefaultPhysicalExprAdapterFactory;
1179        let adapter = factory
1180            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1181            .unwrap();
1182        let column_expr = Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>;
1183
1184        let result = adapter.rewrite(Arc::clone(&column_expr))?;
1185
1186        // Should be the same expression (no transformation needed)
1187        // We compare the underlying pointer through the trait object
1188        assert!(std::ptr::eq(
1189            column_expr.as_ref() as *const dyn PhysicalExpr,
1190            result.as_ref() as *const dyn PhysicalExpr
1191        ));
1192
1193        Ok(())
1194    }
1195
1196    #[test]
1197    fn test_non_nullable_missing_column_error() {
1198        let physical_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
1199        let logical_schema = Schema::new(vec![
1200            Field::new("a", DataType::Int32, false),
1201            Field::new("b", DataType::Utf8, false), // Non-nullable missing column
1202        ]);
1203
1204        let factory = DefaultPhysicalExprAdapterFactory;
1205        let adapter = factory
1206            .create(Arc::new(logical_schema), Arc::new(physical_schema))
1207            .unwrap();
1208        let column_expr = Arc::new(Column::new("b", 1));
1209
1210        let result = adapter.rewrite(column_expr);
1211        assert!(result.is_err());
1212        assert_contains!(
1213            result.unwrap_err().to_string(),
1214            "Non-nullable column 'b' is missing from the physical schema"
1215        );
1216    }
1217
1218    /// Helper function to project expressions onto a RecordBatch
1219    fn batch_project(
1220        expr: Vec<Arc<dyn PhysicalExpr>>,
1221        batch: &RecordBatch,
1222        schema: SchemaRef,
1223    ) -> Result<RecordBatch> {
1224        let arrays = expr
1225            .iter()
1226            .map(|expr| {
1227                expr.evaluate(batch)
1228                    .and_then(|v| v.into_array(batch.num_rows()))
1229            })
1230            .collect::<Result<Vec<_>>>()?;
1231
1232        if arrays.is_empty() {
1233            let options =
1234                RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
1235            RecordBatch::try_new_with_options(Arc::clone(&schema), arrays, &options)
1236                .map_err(Into::into)
1237        } else {
1238            RecordBatch::try_new(Arc::clone(&schema), arrays).map_err(Into::into)
1239        }
1240    }
1241
1242    /// Example showing how we can use the `DefaultPhysicalExprAdapter` to adapt RecordBatches during a scan
1243    /// to apply projections, type conversions and handling of missing columns all at once.
1244    #[test]
1245    fn test_adapt_batches() {
1246        let physical_batch = record_batch!(
1247            ("a", Int32, vec![Some(1), None, Some(3)]),
1248            ("extra", Utf8, vec![Some("x"), Some("y"), None])
1249        )
1250        .unwrap();
1251
1252        let physical_schema = physical_batch.schema();
1253
1254        let logical_schema = Arc::new(Schema::new(vec![
1255            Field::new("a", DataType::Int64, true), // Different type
1256            Field::new("b", DataType::Utf8, true),  // Missing from physical
1257        ]));
1258
1259        let projection = vec![
1260            col("b", &logical_schema).unwrap(),
1261            col("a", &logical_schema).unwrap(),
1262        ];
1263
1264        let factory = DefaultPhysicalExprAdapterFactory;
1265        let adapter = factory
1266            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1267            .unwrap();
1268
1269        let adapted_projection = projection
1270            .into_iter()
1271            .map(|expr| adapter.rewrite(expr).unwrap())
1272            .collect_vec();
1273
1274        let adapted_schema = Arc::new(Schema::new(
1275            adapted_projection
1276                .iter()
1277                .map(|expr| expr.return_field(&physical_schema).unwrap())
1278                .collect_vec(),
1279        ));
1280
1281        let res = batch_project(
1282            adapted_projection,
1283            &physical_batch,
1284            Arc::clone(&adapted_schema),
1285        )
1286        .unwrap();
1287
1288        assert_eq!(res.num_columns(), 2);
1289        assert_eq!(res.column(0).data_type(), &DataType::Utf8);
1290        assert_eq!(res.column(1).data_type(), &DataType::Int64);
1291        assert_eq!(
1292            res.column(0)
1293                .as_any()
1294                .downcast_ref::<StringArray>()
1295                .unwrap()
1296                .iter()
1297                .collect_vec(),
1298            vec![None, None, None]
1299        );
1300        assert_eq!(
1301            res.column(1)
1302                .as_any()
1303                .downcast_ref::<Int64Array>()
1304                .unwrap()
1305                .iter()
1306                .collect_vec(),
1307            vec![Some(1), None, Some(3)]
1308        );
1309    }
1310
1311    /// Test that struct columns are properly adapted including:
1312    /// - Type casting of subfields (Int32 -> Int64, Utf8 -> Utf8View)
1313    /// - Missing fields in logical schema are filled with nulls
1314    #[test]
1315    fn test_adapt_struct_batches() {
1316        // Physical struct: {id: Int32, name: Utf8}
1317        let physical_struct_fields: Fields = vec![
1318            Field::new("id", DataType::Int32, false),
1319            Field::new("name", DataType::Utf8, true),
1320        ]
1321        .into();
1322
1323        let struct_array = StructArray::new(
1324            physical_struct_fields.clone(),
1325            vec![
1326                Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
1327                Arc::new(StringArray::from(vec![
1328                    Some("alice"),
1329                    None,
1330                    Some("charlie"),
1331                ])) as _,
1332            ],
1333            None,
1334        );
1335
1336        let physical_schema = Arc::new(Schema::new(vec![Field::new(
1337            "data",
1338            DataType::Struct(physical_struct_fields),
1339            false,
1340        )]));
1341
1342        let physical_batch = RecordBatch::try_new(
1343            Arc::clone(&physical_schema),
1344            vec![Arc::new(struct_array)],
1345        )
1346        .unwrap();
1347
1348        // Logical struct: {id: Int64, name: Utf8View, extra: Boolean}
1349        // - id: cast from Int32 to Int64
1350        // - name: cast from Utf8 to Utf8View
1351        // - extra: missing from physical, should be filled with nulls
1352        let logical_struct_fields: Fields = vec![
1353            Field::new("id", DataType::Int64, false),
1354            Field::new("name", DataType::Utf8View, true),
1355            Field::new("extra", DataType::Boolean, true), // New field, not in physical
1356        ]
1357        .into();
1358
1359        let logical_schema = Arc::new(Schema::new(vec![Field::new(
1360            "data",
1361            DataType::Struct(logical_struct_fields),
1362            false,
1363        )]));
1364
1365        let projection = vec![col("data", &logical_schema).unwrap()];
1366
1367        let factory = DefaultPhysicalExprAdapterFactory;
1368        let adapter = factory
1369            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1370            .unwrap();
1371
1372        let adapted_projection = projection
1373            .into_iter()
1374            .map(|expr| adapter.rewrite(expr).unwrap())
1375            .collect_vec();
1376
1377        let adapted_schema = Arc::new(Schema::new(
1378            adapted_projection
1379                .iter()
1380                .map(|expr| expr.return_field(&physical_schema).unwrap())
1381                .collect_vec(),
1382        ));
1383
1384        let res = batch_project(
1385            adapted_projection,
1386            &physical_batch,
1387            Arc::clone(&adapted_schema),
1388        )
1389        .unwrap();
1390
1391        assert_eq!(res.num_columns(), 1);
1392
1393        let result_struct = res
1394            .column(0)
1395            .as_any()
1396            .downcast_ref::<StructArray>()
1397            .unwrap();
1398
1399        // Verify id field is cast to Int64
1400        let id_col = result_struct.column_by_name("id").unwrap();
1401        assert_eq!(id_col.data_type(), &DataType::Int64);
1402        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
1403        assert_eq!(
1404            id_values.iter().collect_vec(),
1405            vec![Some(1), Some(2), Some(3)]
1406        );
1407
1408        // Verify name field is cast to Utf8View
1409        let name_col = result_struct.column_by_name("name").unwrap();
1410        assert_eq!(name_col.data_type(), &DataType::Utf8View);
1411        let name_values = name_col.as_any().downcast_ref::<StringViewArray>().unwrap();
1412        assert_eq!(
1413            name_values.iter().collect_vec(),
1414            vec![Some("alice"), None, Some("charlie")]
1415        );
1416
1417        // Verify extra field (missing from physical) is filled with nulls
1418        let extra_col = result_struct.column_by_name("extra").unwrap();
1419        assert_eq!(extra_col.data_type(), &DataType::Boolean);
1420        let extra_values = extra_col.as_any().downcast_ref::<BooleanArray>().unwrap();
1421        assert_eq!(extra_values.iter().collect_vec(), vec![None, None, None]);
1422    }
1423
1424    /// Test that List<Struct> columns are properly adapted with struct evolution.
1425    #[test]
1426    fn test_adapt_list_struct_batches() {
1427        // Physical: List<{id: Int32, name: Utf8}>
1428        let physical_struct_fields: Fields = vec![
1429            Field::new("id", DataType::Int32, false),
1430            Field::new("name", DataType::Utf8, true),
1431        ]
1432        .into();
1433
1434        let struct_array = StructArray::new(
1435            physical_struct_fields.clone(),
1436            vec![
1437                Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
1438                Arc::new(StringArray::from(vec![
1439                    Some("alice"),
1440                    None,
1441                    Some("charlie"),
1442                ])) as _,
1443            ],
1444            None,
1445        );
1446
1447        // One list element per row
1448        let item_field = Arc::new(Field::new(
1449            "item",
1450            DataType::Struct(physical_struct_fields.clone()),
1451            true,
1452        ));
1453        let offsets =
1454            arrow::buffer::OffsetBuffer::from_lengths(vec![1usize; struct_array.len()]);
1455        let list_array = GenericListArray::<i32>::new(
1456            item_field,
1457            offsets,
1458            Arc::new(struct_array),
1459            None,
1460        );
1461
1462        let physical_schema = Arc::new(Schema::new(vec![Field::new(
1463            "data",
1464            DataType::List(Arc::new(Field::new(
1465                "item",
1466                DataType::Struct(physical_struct_fields),
1467                true,
1468            ))),
1469            false,
1470        )]));
1471
1472        let physical_batch = RecordBatch::try_new(
1473            Arc::clone(&physical_schema),
1474            vec![Arc::new(list_array)],
1475        )
1476        .unwrap();
1477
1478        // Logical: List<{id: Int64, name: Utf8View, extra: Boolean}>
1479        let logical_struct_fields: Fields = vec![
1480            Field::new("id", DataType::Int64, false),
1481            Field::new("name", DataType::Utf8View, true),
1482            Field::new("extra", DataType::Boolean, true),
1483        ]
1484        .into();
1485
1486        let logical_schema = Arc::new(Schema::new(vec![Field::new(
1487            "data",
1488            DataType::List(Arc::new(Field::new(
1489                "item",
1490                DataType::Struct(logical_struct_fields.clone()),
1491                true,
1492            ))),
1493            false,
1494        )]));
1495
1496        let projection = vec![col("data", &logical_schema).unwrap()];
1497
1498        let factory = DefaultPhysicalExprAdapterFactory;
1499        let adapter = factory
1500            .create(Arc::clone(&logical_schema), Arc::clone(&physical_schema))
1501            .unwrap();
1502
1503        let adapted_projection = projection
1504            .into_iter()
1505            .map(|expr| adapter.rewrite(expr).unwrap())
1506            .collect_vec();
1507
1508        let adapted_schema = Arc::new(Schema::new(
1509            adapted_projection
1510                .iter()
1511                .map(|expr| expr.return_field(&physical_schema).unwrap())
1512                .collect_vec(),
1513        ));
1514
1515        let res = batch_project(
1516            adapted_projection,
1517            &physical_batch,
1518            Arc::clone(&adapted_schema),
1519        )
1520        .unwrap();
1521
1522        assert_eq!(res.num_columns(), 1);
1523
1524        let result_list = res
1525            .column(0)
1526            .as_any()
1527            .downcast_ref::<GenericListArray<i32>>()
1528            .unwrap();
1529
1530        // Check each list element contains the evolved struct
1531        assert_eq!(result_list.len(), 3);
1532        let flat_structs = result_list
1533            .values()
1534            .as_any()
1535            .downcast_ref::<StructArray>()
1536            .unwrap();
1537
1538        let id_col = flat_structs.column_by_name("id").unwrap();
1539        assert_eq!(id_col.data_type(), &DataType::Int64);
1540        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
1541        assert_eq!(
1542            id_values.iter().collect_vec(),
1543            vec![Some(1), Some(2), Some(3)]
1544        );
1545
1546        let name_col = flat_structs.column_by_name("name").unwrap();
1547        assert_eq!(name_col.data_type(), &DataType::Utf8View);
1548        let name_values = name_col.as_any().downcast_ref::<StringViewArray>().unwrap();
1549        assert_eq!(
1550            name_values.iter().collect_vec(),
1551            vec![Some("alice"), None, Some("charlie")]
1552        );
1553
1554        let extra_col = flat_structs.column_by_name("extra").unwrap();
1555        assert_eq!(extra_col.data_type(), &DataType::Boolean);
1556        let extra_values = extra_col.as_any().downcast_ref::<BooleanArray>().unwrap();
1557        assert_eq!(extra_values.iter().collect_vec(), vec![None, None, None]);
1558    }
1559
1560    #[test]
1561    fn test_try_rewrite_struct_field_access() {
1562        // Test the core logic of try_rewrite_struct_field_access
1563        let physical_schema = Schema::new(vec![Field::new(
1564            "struct_col",
1565            DataType::Struct(
1566                vec![Field::new("existing_field", DataType::Int32, true)].into(),
1567            ),
1568            true,
1569        )]);
1570
1571        let logical_schema = Schema::new(vec![Field::new(
1572            "struct_col",
1573            DataType::Struct(
1574                vec![
1575                    Field::new("existing_field", DataType::Int32, true),
1576                    Field::new("missing_field", DataType::Utf8, true),
1577                ]
1578                .into(),
1579            ),
1580            true,
1581        )]);
1582
1583        let rewriter = DefaultPhysicalExprAdapterRewriter {
1584            logical_file_schema: Arc::new(logical_schema),
1585            physical_file_schema: Arc::new(physical_schema),
1586        };
1587
1588        // Test that when a field exists in physical schema, it returns None
1589        let column = Arc::new(Column::new("struct_col", 0)) as Arc<dyn PhysicalExpr>;
1590        let result = rewriter.try_rewrite_struct_field_access(&column).unwrap();
1591        assert!(result.is_none());
1592
1593        // The actual test for the get_field expression would require creating a proper ScalarFunctionExpr
1594        // with ScalarUDF, which is complex to set up in a unit test. The integration tests in
1595        // datafusion/core/tests/parquet/schema_adapter.rs provide better coverage for this functionality.
1596    }
1597
1598    /// Build `get_field(column, 'field')` against `schema`.
1599    fn get_field_expr(
1600        schema: &Schema,
1601        column: &str,
1602        field: &str,
1603    ) -> Arc<dyn PhysicalExpr> {
1604        let index = schema.index_of(column).unwrap();
1605        Arc::new(
1606            ScalarFunctionExpr::try_new(
1607                Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())),
1608                vec![
1609                    Arc::new(Column::new(column, index)),
1610                    Arc::new(Literal::new(ScalarValue::from(field))),
1611                ],
1612                schema,
1613                Arc::new(datafusion_common::config::ConfigOptions::default()),
1614            )
1615            .unwrap(),
1616        )
1617    }
1618
1619    fn struct_schemas(
1620        physical_fields: Vec<Field>,
1621        logical_fields: Vec<Field>,
1622    ) -> (SchemaRef, SchemaRef) {
1623        let physical = Arc::new(Schema::new(vec![Field::new(
1624            "s",
1625            DataType::Struct(physical_fields.into()),
1626            true,
1627        )]));
1628        let logical = Arc::new(Schema::new(vec![Field::new(
1629            "s",
1630            DataType::Struct(logical_fields.into()),
1631            true,
1632        )]));
1633        (logical, physical)
1634    }
1635
1636    /// `s['x']` where the file stores `x` as `Int32` and the table declares
1637    /// `Int64` must cast the extracted field, not the whole struct, so that
1638    /// the column stays visible under the `get_field`.
1639    ///
1640    /// See <https://github.com/apache/datafusion/issues/24109>.
1641    #[test]
1642    fn test_narrow_struct_cast_to_field_access() {
1643        let (logical_schema, physical_schema) = struct_schemas(
1644            vec![Field::new("x", DataType::Int32, true)],
1645            vec![Field::new("x", DataType::Int64, true)],
1646        );
1647
1648        let adapter = DefaultPhysicalExprAdapterFactory
1649            .create(Arc::clone(&logical_schema), physical_schema)
1650            .unwrap();
1651        let rewritten = adapter
1652            .rewrite(get_field_expr(&logical_schema, "s", "x"))
1653            .unwrap();
1654
1655        let cast = assert_cast_expr(&rewritten);
1656        assert_eq!(cast.cast_type(), &DataType::Int64);
1657        let get_field = cast
1658            .expr()
1659            .downcast_ref::<ScalarFunctionExpr>()
1660            .expect("Expected get_field under the cast");
1661        assert_eq!(get_field.return_type(), &DataType::Int32);
1662        assert!(
1663            get_field.args()[0].downcast_ref::<Column>().is_some(),
1664            "the struct column must not be hidden behind a cast, got: {rewritten}"
1665        );
1666    }
1667
1668    /// A struct field that only differs in a nested leaf type still ends up
1669    /// with a single cast on the extracted field.
1670    #[test]
1671    fn test_narrow_struct_cast_nested_field_access() {
1672        let (logical_schema, physical_schema) = struct_schemas(
1673            vec![Field::new(
1674                "inner",
1675                DataType::Struct(vec![Field::new("x", DataType::Utf8, true)].into()),
1676                true,
1677            )],
1678            vec![Field::new(
1679                "inner",
1680                DataType::Struct(vec![Field::new("x", DataType::Utf8View, true)].into()),
1681                true,
1682            )],
1683        );
1684
1685        let adapter = DefaultPhysicalExprAdapterFactory
1686            .create(Arc::clone(&logical_schema), physical_schema)
1687            .unwrap();
1688        let outer = get_field_expr(&logical_schema, "s", "inner");
1689        let expr = Arc::new(
1690            ScalarFunctionExpr::try_new(
1691                Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())),
1692                vec![outer, Arc::new(Literal::new(ScalarValue::from("x")))],
1693                &logical_schema,
1694                Arc::new(datafusion_common::config::ConfigOptions::default()),
1695            )
1696            .unwrap(),
1697        ) as Arc<dyn PhysicalExpr>;
1698
1699        let rewritten = adapter.rewrite(expr).unwrap();
1700
1701        let cast = assert_cast_expr(&rewritten);
1702        assert_eq!(cast.cast_type(), &DataType::Utf8View);
1703        let outer_get_field = cast
1704            .expr()
1705            .downcast_ref::<ScalarFunctionExpr>()
1706            .expect("Expected get_field under the cast");
1707        let inner_get_field = outer_get_field.args()[0]
1708            .downcast_ref::<ScalarFunctionExpr>()
1709            .expect("Expected a nested get_field");
1710        assert!(
1711            inner_get_field.args()[0].downcast_ref::<Column>().is_some(),
1712            "the struct column must not be hidden behind a cast, got: {rewritten}"
1713        );
1714    }
1715
1716    /// A struct column that needs no adaptation at all is left completely
1717    /// alone — the narrowing must not disturb the common case.
1718    #[test]
1719    fn test_narrow_struct_cast_leaves_matching_schema_alone() {
1720        let (logical_schema, physical_schema) = struct_schemas(
1721            vec![Field::new("x", DataType::Int32, true)],
1722            vec![Field::new("x", DataType::Int32, true)],
1723        );
1724
1725        let adapter = DefaultPhysicalExprAdapterFactory
1726            .create(Arc::clone(&logical_schema), physical_schema)
1727            .unwrap();
1728        let expr = get_field_expr(&logical_schema, "s", "x");
1729
1730        let rewritten = adapter.rewrite(Arc::clone(&expr)).unwrap();
1731
1732        assert_eq!(
1733            rewritten.to_string(),
1734            expr.to_string(),
1735            "an unadapted struct column must pass through untouched"
1736        );
1737    }
1738
1739    /// When the accessed field has the same type in both schemas, the struct
1740    /// cast disappears entirely rather than being replaced by a field cast:
1741    /// only a sibling field forced the column-level cast in the first place.
1742    #[test]
1743    fn test_narrow_struct_cast_drops_cast_when_field_types_match() {
1744        let (logical_schema, physical_schema) = struct_schemas(
1745            vec![
1746                Field::new("x", DataType::Int32, true),
1747                Field::new("y", DataType::Int32, true),
1748            ],
1749            vec![
1750                Field::new("x", DataType::Int32, true),
1751                Field::new("y", DataType::Int64, true),
1752            ],
1753        );
1754
1755        let adapter = DefaultPhysicalExprAdapterFactory
1756            .create(Arc::clone(&logical_schema), physical_schema)
1757            .unwrap();
1758        let expr = Arc::new(
1759            ScalarFunctionExpr::try_new(
1760                Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())),
1761                vec![
1762                    Arc::new(Column::new("s", 0)),
1763                    Arc::new(Literal::new(ScalarValue::from("x"))),
1764                ],
1765                &logical_schema,
1766                Arc::new(datafusion_common::config::ConfigOptions::default()),
1767            )
1768            .unwrap(),
1769        ) as Arc<dyn PhysicalExpr>;
1770
1771        let rewritten = adapter.rewrite(expr).unwrap();
1772
1773        assert!(
1774            rewritten.downcast_ref::<CastExpr>().is_none(),
1775            "`x` has the same type in both schemas, so no cast is needed, got: {rewritten}"
1776        );
1777        let get_field = rewritten
1778            .downcast_ref::<ScalarFunctionExpr>()
1779            .expect("Expected a bare get_field");
1780        assert_eq!(get_field.return_type(), &DataType::Int32);
1781        assert!(
1782            get_field.args()[0].downcast_ref::<Column>().is_some(),
1783            "the struct column must not be hidden behind a cast, got: {rewritten}"
1784        );
1785    }
1786
1787    /// `s['inner']['x']` is simplified to the flattened `get_field(s, 'inner',
1788    /// 'x')`, so the whole key path has to be resolved.
1789    #[test]
1790    fn test_narrow_struct_cast_flattened_field_path() {
1791        let (logical_schema, physical_schema) = struct_schemas(
1792            vec![Field::new(
1793                "inner",
1794                DataType::Struct(vec![Field::new("x", DataType::Utf8, true)].into()),
1795                true,
1796            )],
1797            vec![Field::new(
1798                "inner",
1799                DataType::Struct(vec![Field::new("x", DataType::Utf8View, true)].into()),
1800                true,
1801            )],
1802        );
1803
1804        let adapter = DefaultPhysicalExprAdapterFactory
1805            .create(Arc::clone(&logical_schema), physical_schema)
1806            .unwrap();
1807        let expr = Arc::new(
1808            ScalarFunctionExpr::try_new(
1809                Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())),
1810                vec![
1811                    Arc::new(Column::new("s", 0)),
1812                    Arc::new(Literal::new(ScalarValue::from("inner"))),
1813                    Arc::new(Literal::new(ScalarValue::from("x"))),
1814                ],
1815                &logical_schema,
1816                Arc::new(datafusion_common::config::ConfigOptions::default()),
1817            )
1818            .unwrap(),
1819        ) as Arc<dyn PhysicalExpr>;
1820
1821        let rewritten = adapter.rewrite(expr).unwrap();
1822
1823        let cast = assert_cast_expr(&rewritten);
1824        assert_eq!(cast.cast_type(), &DataType::Utf8View);
1825        let get_field = cast
1826            .expr()
1827            .downcast_ref::<ScalarFunctionExpr>()
1828            .expect("Expected get_field under the cast");
1829        assert_eq!(get_field.return_type(), &DataType::Utf8);
1830        assert_eq!(
1831            get_field.args().len(),
1832            3,
1833            "the full key path must be preserved, got: {rewritten}"
1834        );
1835        assert!(
1836            get_field.args()[0].downcast_ref::<Column>().is_some(),
1837            "the struct column must not be hidden behind a cast, got: {rewritten}"
1838        );
1839    }
1840
1841    /// A key path whose leaf is missing from the file still resolves to a
1842    /// typed null literal.
1843    #[test]
1844    fn test_narrow_struct_cast_flattened_field_path_missing_leaf() {
1845        let (logical_schema, physical_schema) = struct_schemas(
1846            vec![Field::new(
1847                "inner",
1848                DataType::Struct(vec![Field::new("x", DataType::Int32, true)].into()),
1849                true,
1850            )],
1851            vec![Field::new(
1852                "inner",
1853                DataType::Struct(
1854                    vec![
1855                        Field::new("x", DataType::Int32, true),
1856                        Field::new("y", DataType::Utf8, true),
1857                    ]
1858                    .into(),
1859                ),
1860                true,
1861            )],
1862        );
1863
1864        let adapter = DefaultPhysicalExprAdapterFactory
1865            .create(Arc::clone(&logical_schema), physical_schema)
1866            .unwrap();
1867        let expr = Arc::new(
1868            ScalarFunctionExpr::try_new(
1869                Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())),
1870                vec![
1871                    Arc::new(Column::new("s", 0)),
1872                    Arc::new(Literal::new(ScalarValue::from("inner"))),
1873                    Arc::new(Literal::new(ScalarValue::from("y"))),
1874                ],
1875                &logical_schema,
1876                Arc::new(datafusion_common::config::ConfigOptions::default()),
1877            )
1878            .unwrap(),
1879        ) as Arc<dyn PhysicalExpr>;
1880
1881        let rewritten = adapter.rewrite(expr).unwrap();
1882
1883        let literal = rewritten
1884            .downcast_ref::<Literal>()
1885            .expect("Expected a null literal");
1886        assert_eq!(*literal.value(), ScalarValue::Utf8(None));
1887    }
1888
1889    /// Accessing a field the file does not have yields a typed null literal.
1890    #[test]
1891    fn test_narrow_struct_cast_missing_field() {
1892        let (logical_schema, physical_schema) = struct_schemas(
1893            vec![Field::new("x", DataType::Int32, true)],
1894            vec![
1895                Field::new("x", DataType::Int32, true),
1896                Field::new("y", DataType::Utf8, true),
1897            ],
1898        );
1899
1900        let adapter = DefaultPhysicalExprAdapterFactory
1901            .create(Arc::clone(&logical_schema), physical_schema)
1902            .unwrap();
1903        let rewritten = adapter
1904            .rewrite(get_field_expr(&logical_schema, "s", "y"))
1905            .unwrap();
1906
1907        let literal = rewritten
1908            .downcast_ref::<Literal>()
1909            .expect("Expected a null literal");
1910        assert_eq!(*literal.value(), ScalarValue::Utf8(None));
1911    }
1912
1913    /// `get_field` on a Map column is a runtime key lookup, not a schema-level
1914    /// field access, so the map value must keep its cast.
1915    #[test]
1916    fn test_map_field_access_keeps_cast() {
1917        let map_type = |value_type: DataType| {
1918            DataType::Map(
1919                Arc::new(Field::new(
1920                    "entries",
1921                    DataType::Struct(
1922                        vec![
1923                            Field::new("keys", DataType::Utf8, false),
1924                            Field::new("values", value_type, true),
1925                        ]
1926                        .into(),
1927                    ),
1928                    false,
1929                )),
1930                false,
1931            )
1932        };
1933        let physical_schema = Arc::new(Schema::new(vec![Field::new(
1934            "s",
1935            map_type(DataType::Int32),
1936            true,
1937        )]));
1938        let logical_schema = Arc::new(Schema::new(vec![Field::new(
1939            "s",
1940            map_type(DataType::Int64),
1941            true,
1942        )]));
1943
1944        let adapter = DefaultPhysicalExprAdapterFactory
1945            .create(Arc::clone(&logical_schema), physical_schema)
1946            .unwrap();
1947        let rewritten = adapter
1948            .rewrite(get_field_expr(&logical_schema, "s", "k"))
1949            .unwrap();
1950
1951        let get_field = rewritten
1952            .downcast_ref::<ScalarFunctionExpr>()
1953            .expect("Expected the get_field to be preserved");
1954        assert!(
1955            get_field.args()[0].downcast_ref::<CastExpr>().is_some(),
1956            "map columns must keep the whole-column cast, got: {rewritten}"
1957        );
1958    }
1959
1960    // ============================================================================
1961    // BatchAdapterFactory and BatchAdapter tests
1962    // ============================================================================
1963
1964    #[test]
1965    fn test_batch_adapter_factory_basic() {
1966        // Target schema
1967        let target_schema = Arc::new(Schema::new(vec![
1968            Field::new("a", DataType::Int64, false),
1969            Field::new("b", DataType::Utf8, true),
1970        ]));
1971
1972        // Source schema with different column order and type
1973        let source_schema = Arc::new(Schema::new(vec![
1974            Field::new("b", DataType::Utf8, true),
1975            Field::new("a", DataType::Int32, false), // Int32 -> Int64
1976        ]));
1977
1978        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1979        let adapter = factory.make_adapter(&source_schema).unwrap();
1980
1981        // Create source batch
1982        let source_batch = RecordBatch::try_new(
1983            Arc::clone(&source_schema),
1984            vec![
1985                Arc::new(StringArray::from(vec![Some("hello"), None, Some("world")])),
1986                Arc::new(Int32Array::from(vec![1, 2, 3])),
1987            ],
1988        )
1989        .unwrap();
1990
1991        let adapted = adapter.adapt_batch(&source_batch).unwrap();
1992
1993        // Verify schema matches target
1994        assert_eq!(adapted.num_columns(), 2);
1995        assert_eq!(adapted.schema().field(0).name(), "a");
1996        assert_eq!(adapted.schema().field(0).data_type(), &DataType::Int64);
1997        assert_eq!(adapted.schema().field(1).name(), "b");
1998        assert_eq!(adapted.schema().field(1).data_type(), &DataType::Utf8);
1999
2000        // Verify data
2001        let col_a = adapted
2002            .column(0)
2003            .as_any()
2004            .downcast_ref::<Int64Array>()
2005            .unwrap();
2006        assert_eq!(col_a.iter().collect_vec(), vec![Some(1), Some(2), Some(3)]);
2007
2008        let col_b = adapted
2009            .column(1)
2010            .as_any()
2011            .downcast_ref::<StringArray>()
2012            .unwrap();
2013        assert_eq!(
2014            col_b.iter().collect_vec(),
2015            vec![Some("hello"), None, Some("world")]
2016        );
2017    }
2018
2019    #[test]
2020    fn test_batch_adapter_factory_missing_column() {
2021        // Target schema with a column missing from source
2022        let target_schema = Arc::new(Schema::new(vec![
2023            Field::new("a", DataType::Int32, false),
2024            Field::new("b", DataType::Utf8, true), // exists in source
2025            Field::new("c", DataType::Float64, true), // missing from source
2026        ]));
2027
2028        let source_schema = Arc::new(Schema::new(vec![
2029            Field::new("a", DataType::Int32, false),
2030            Field::new("b", DataType::Utf8, true),
2031        ]));
2032
2033        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
2034        let adapter = factory.make_adapter(&source_schema).unwrap();
2035
2036        let source_batch = RecordBatch::try_new(
2037            Arc::clone(&source_schema),
2038            vec![
2039                Arc::new(Int32Array::from(vec![1, 2])),
2040                Arc::new(StringArray::from(vec!["x", "y"])),
2041            ],
2042        )
2043        .unwrap();
2044
2045        let adapted = adapter.adapt_batch(&source_batch).unwrap();
2046
2047        assert_eq!(adapted.num_columns(), 3);
2048
2049        // Missing column should be filled with nulls
2050        let col_c = adapted.column(2);
2051        assert_eq!(col_c.data_type(), &DataType::Float64);
2052        assert_eq!(col_c.null_count(), 2); // All nulls
2053    }
2054
2055    #[test]
2056    fn test_batch_adapter_factory_with_struct() {
2057        // Target has struct with Int64 id
2058        let target_struct_fields: Fields = vec![
2059            Field::new("id", DataType::Int64, false),
2060            Field::new("name", DataType::Utf8, true),
2061        ]
2062        .into();
2063        let target_schema = Arc::new(Schema::new(vec![Field::new(
2064            "data",
2065            DataType::Struct(target_struct_fields),
2066            false,
2067        )]));
2068
2069        // Source has struct with Int32 id
2070        let source_struct_fields: Fields = vec![
2071            Field::new("id", DataType::Int32, false),
2072            Field::new("name", DataType::Utf8, true),
2073        ]
2074        .into();
2075        let source_schema = Arc::new(Schema::new(vec![Field::new(
2076            "data",
2077            DataType::Struct(source_struct_fields.clone()),
2078            false,
2079        )]));
2080
2081        let struct_array = StructArray::new(
2082            source_struct_fields,
2083            vec![
2084                Arc::new(Int32Array::from(vec![10, 20])) as _,
2085                Arc::new(StringArray::from(vec!["a", "b"])) as _,
2086            ],
2087            None,
2088        );
2089
2090        let source_batch = RecordBatch::try_new(
2091            Arc::clone(&source_schema),
2092            vec![Arc::new(struct_array)],
2093        )
2094        .unwrap();
2095
2096        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
2097        let adapter = factory.make_adapter(&source_schema).unwrap();
2098        let adapted = adapter.adapt_batch(&source_batch).unwrap();
2099
2100        let result_struct = adapted
2101            .column(0)
2102            .as_any()
2103            .downcast_ref::<StructArray>()
2104            .unwrap();
2105
2106        // Verify id was cast to Int64
2107        let id_col = result_struct.column_by_name("id").unwrap();
2108        assert_eq!(id_col.data_type(), &DataType::Int64);
2109        let id_values = id_col.as_any().downcast_ref::<Int64Array>().unwrap();
2110        assert_eq!(id_values.iter().collect_vec(), vec![Some(10), Some(20)]);
2111    }
2112
2113    #[test]
2114    fn test_batch_adapter_factory_identity() {
2115        // When source and target schemas are identical, should pass through efficiently
2116        let schema = Arc::new(Schema::new(vec![
2117            Field::new("a", DataType::Int32, false),
2118            Field::new("b", DataType::Utf8, true),
2119        ]));
2120
2121        let factory = BatchAdapterFactory::new(Arc::clone(&schema));
2122        let adapter = factory.make_adapter(&schema).unwrap();
2123
2124        let batch = RecordBatch::try_new(
2125            Arc::clone(&schema),
2126            vec![
2127                Arc::new(Int32Array::from(vec![1, 2, 3])),
2128                Arc::new(StringArray::from(vec!["a", "b", "c"])),
2129            ],
2130        )
2131        .unwrap();
2132
2133        let adapted = adapter.adapt_batch(&batch).unwrap();
2134
2135        assert_eq!(adapted.num_columns(), 2);
2136        assert_eq!(adapted.schema().field(0).data_type(), &DataType::Int32);
2137        assert_eq!(adapted.schema().field(1).data_type(), &DataType::Utf8);
2138    }
2139
2140    #[test]
2141    fn test_batch_adapter_factory_reuse() {
2142        // Factory can create multiple adapters for different source schemas
2143        let target_schema = Arc::new(Schema::new(vec![
2144            Field::new("x", DataType::Int64, false),
2145            Field::new("y", DataType::Utf8, true),
2146        ]));
2147
2148        let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
2149
2150        // First source schema
2151        let source1 = Arc::new(Schema::new(vec![
2152            Field::new("x", DataType::Int32, false),
2153            Field::new("y", DataType::Utf8, true),
2154        ]));
2155        let adapter1 = factory.make_adapter(&source1).unwrap();
2156
2157        // Second source schema (different order)
2158        let source2 = Arc::new(Schema::new(vec![
2159            Field::new("y", DataType::Utf8, true),
2160            Field::new("x", DataType::Int64, false),
2161        ]));
2162        let adapter2 = factory.make_adapter(&source2).unwrap();
2163
2164        // Both should work correctly
2165        assert!(format!("{adapter1:?}").contains("BatchAdapter"));
2166        assert!(format!("{adapter2:?}").contains("BatchAdapter"));
2167    }
2168
2169    #[test]
2170    fn test_rewrite_column_index_and_type_mismatch() {
2171        let physical_schema = Schema::new(vec![
2172            Field::new("b", DataType::Utf8, true),
2173            Field::new("a", DataType::Int32, false), // Index 1
2174        ]);
2175
2176        let logical_schema = Schema::new(vec![
2177            Field::new("a", DataType::Int64, false), // Index 0, Different Type
2178            Field::new("b", DataType::Utf8, true),
2179        ]);
2180
2181        let adapter = DefaultPhysicalExprAdapterFactory
2182            .create(Arc::new(logical_schema), Arc::new(physical_schema))
2183            .unwrap();
2184
2185        // Logical column "a" is at index 0
2186        let column_expr = Arc::new(Column::new("a", 0));
2187
2188        let result = adapter.rewrite(column_expr).unwrap();
2189
2190        // Should be a CastExpr
2191        let cast_expr = assert_cast_expr(&result);
2192
2193        // Verify the inner column points to the correct physical index (1)
2194        assert_cast_input_column(cast_expr, "a", 1);
2195
2196        // Verify cast types
2197        assert_eq!(
2198            cast_expr.data_type(&Schema::empty()).unwrap(),
2199            DataType::Int64
2200        );
2201    }
2202
2203    #[test]
2204    fn test_rewrite_resolves_physical_column_by_name_before_casting() {
2205        let (logical_schema, physical_schema) = stale_index_cast_schemas();
2206        let adapter = DefaultPhysicalExprAdapterFactory
2207            .create(logical_schema, physical_schema)
2208            .unwrap();
2209
2210        // Deliberately provide the wrong index for column `a`.
2211        // Regression: this must still resolve against physical field `a` by name.
2212        let rewritten = adapter.rewrite(Arc::new(Column::new("a", 0))).unwrap();
2213        let cast_expr = assert_cast_expr(&rewritten);
2214        assert_cast_input_column(cast_expr, "a", 1);
2215        assert_eq!(cast_expr.target_field().data_type(), &DataType::Int64);
2216    }
2217}