1use 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
46pub 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
85pub trait PhysicalExprAdapter: Send + Sync + std::fmt::Debug {
153 fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>>;
170}
171
172pub trait PhysicalExprAdapterFactory: Send + Sync + std::fmt::Debug {
176 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#[derive(Debug, Clone)]
242pub struct DefaultPhysicalExprAdapter {
243 logical_file_schema: SchemaRef,
244 physical_file_schema: SchemaRef,
245}
246
247impl DefaultPhysicalExprAdapter {
248 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
276enum FieldPathResolution<'a> {
278 Found(&'a FieldRef),
280 Missing,
282 NotAStruct,
285}
286
287fn 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 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 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 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 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 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 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 let logical_field = match self.logical_file_schema.field_with_name(column.name())
550 {
551 Ok(field) => field,
552 Err(e) => {
553 if let Ok(physical_field) =
557 self.physical_file_schema.field_with_name(column.name())
558 {
559 physical_field
563 } else {
564 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 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 return Ok(Transformed::yes(Arc::new(resolved_column)));
598 }
599
600 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 fn resolve_physical_column(
628 &self,
629 column: &Column,
630 ) -> Result<Option<(Column, FieldRef)>> {
631 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#[derive(Debug)]
715pub struct BatchAdapterFactory {
716 target_schema: SchemaRef,
717 expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
718}
719
720impl BatchAdapterFactory {
721 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 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 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#[derive(Debug)]
780pub struct BatchAdapter {
781 projector: Projector,
782}
783
784impl BatchAdapter {
785 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), Field::new("b", DataType::Utf8, true),
843 Field::new("c", DataType::Float64, true), ]);
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 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 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 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 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 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), ]);
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), ]);
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 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 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), ]);
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 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 #[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), Field::new("b", DataType::Utf8, true), ]));
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]
1315 fn test_adapt_struct_batches() {
1316 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 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), ]
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 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 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 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]
1426 fn test_adapt_list_struct_batches() {
1427 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 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 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 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 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 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 }
1597
1598 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[test]
1965 fn test_batch_adapter_factory_basic() {
1966 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 let source_schema = Arc::new(Schema::new(vec![
1974 Field::new("b", DataType::Utf8, true),
1975 Field::new("a", DataType::Int32, false), ]));
1977
1978 let factory = BatchAdapterFactory::new(Arc::clone(&target_schema));
1979 let adapter = factory.make_adapter(&source_schema).unwrap();
1980
1981 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 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 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 let target_schema = Arc::new(Schema::new(vec![
2023 Field::new("a", DataType::Int32, false),
2024 Field::new("b", DataType::Utf8, true), Field::new("c", DataType::Float64, true), ]));
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 let col_c = adapted.column(2);
2051 assert_eq!(col_c.data_type(), &DataType::Float64);
2052 assert_eq!(col_c.null_count(), 2); }
2054
2055 #[test]
2056 fn test_batch_adapter_factory_with_struct() {
2057 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 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 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 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 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 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 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 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), ]);
2175
2176 let logical_schema = Schema::new(vec![
2177 Field::new("a", DataType::Int64, false), 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 let column_expr = Arc::new(Column::new("a", 0));
2187
2188 let result = adapter.rewrite(column_expr).unwrap();
2189
2190 let cast_expr = assert_cast_expr(&result);
2192
2193 assert_cast_input_column(cast_expr, "a", 1);
2195
2196 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 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}