better test

Author: adriangb

datafusion/core/src/datasource/listing/table.rs

@@ -938,6 +938,8 @@ pub struct ListingTable {
     column_defaults: HashMap<String, Expr>,
     /// Optional [`SchemaAdapterFactory`] for creating schema adapters
     schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
+    /// Optional [`PhysicalExprAdapterFactory`] for creating physical expression adapters
+    expr_adapter_factory: Option<Arc<dyn PhysicalExprAdapterFactory>>,
 }
 
 impl ListingTable {
@@ -979,6 +981,7 @@ impl ListingTable {
             constraints: Constraints::default(),
             column_defaults: HashMap::new(),
             schema_adapter_factory: config.schema_adapter_factory,
+            expr_adapter_factory: config.physical_expr_adapter_factory,
         };
 
         Ok(table)
@@ -1223,6 +1226,7 @@ impl TableProvider for ListingTable {
                 .with_limit(limit)
                 .with_output_ordering(output_ordering)
                 .with_table_partition_cols(table_partition_cols)
+                .with_expr_adapter(self.expr_adapter_factory.clone())
                 .build(),
             )
             .await

datafusion/core/tests/parquet/schema_adapter.rs

@@ -17,16 +17,27 @@
 
 use std::sync::Arc;
 
-use arrow::array::{record_batch, RecordBatch};
-use arrow_schema::{DataType, Field, Schema};
+use arrow::array::{record_batch, RecordBatch, RecordBatchOptions};
+use arrow::compute::{cast_with_options, CastOptions};
+use arrow_schema::{DataType, Field, FieldRef, Schema, SchemaRef};
 use bytes::{BufMut, BytesMut};
 use datafusion::assert_batches_eq;
+use datafusion::common::Result;
 use datafusion::datasource::listing::{ListingTable, ListingTableConfig};
-use datafusion::prelude::SessionContext;
-use datafusion_datasource::schema_adapter::DefaultSchemaAdapterFactory;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::{ColumnStatistics, ScalarValue};
+use datafusion_datasource::schema_adapter::{
+    DefaultSchemaAdapterFactory, SchemaAdapter, SchemaAdapterFactory, SchemaMapper,
+};
 use datafusion_datasource::ListingTableUrl;
 use datafusion_execution::object_store::ObjectStoreUrl;
-use datafusion_physical_expr::schema_rewriter::DefaultPhysicalExprAdapterFactory;
+use datafusion_physical_expr::expressions::{self, Column};
+use datafusion_physical_expr::schema_rewriter::{
+    DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
+};
+use datafusion_physical_expr::{DefaultPhysicalExprAdapter, PhysicalExpr};
+use itertools::Itertools;
 use object_store::{memory::InMemory, path::Path, ObjectStore};
 use parquet::arrow::ArrowWriter;
 
@@ -41,6 +52,180 @@ async fn write_parquet(batch: RecordBatch, store: Arc<dyn ObjectStore>, path: &s
     store.put(&Path::from(path), data.into()).await.unwrap();
 }
 
+#[derive(Debug)]
+struct CustomSchemaAdapterFactory;
+
+impl SchemaAdapterFactory for CustomSchemaAdapterFactory {
+    fn create(
+        &self,
+        projected_table_schema: SchemaRef,
+        _table_schema: SchemaRef,
+    ) -> Box<dyn SchemaAdapter> {
+        Box::new(CustomSchemaAdapter {
+            logical_file_schema: projected_table_schema,
+        })
+    }
+}
+
+#[derive(Debug)]
+struct CustomSchemaAdapter {
+    logical_file_schema: SchemaRef,
+}
+
+impl SchemaAdapter for CustomSchemaAdapter {
+    fn map_column_index(&self, index: usize, file_schema: &Schema) -> Option<usize> {
+        for (idx, field) in file_schema.fields().iter().enumerate() {
+            if field.name() == self.logical_file_schema.field(index).name() {
+                return Some(idx);
+            }
+        }
+        None
+    }
+
+    fn map_schema(
+        &self,
+        file_schema: &Schema,
+    ) -> Result<(Arc<dyn SchemaMapper>, Vec<usize>)> {
+        let projection = (0..file_schema.fields().len()).collect_vec();
+        Ok((
+            Arc::new(CustomSchemaMapper {
+                logical_file_schema: Arc::clone(&self.logical_file_schema),
+            }),
+            projection,
+        ))
+    }
+}
+
+#[derive(Debug)]
+struct CustomSchemaMapper {
+    logical_file_schema: SchemaRef,
+}
+
+impl SchemaMapper for CustomSchemaMapper {
+    fn map_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
+        let mut output_columns =
+            Vec::with_capacity(self.logical_file_schema.fields().len());
+        for field in self.logical_file_schema.fields() {
+            if let Some(array) = batch.column_by_name(field.name()) {
+                output_columns.push(cast_with_options(
+                    array,
+                    field.data_type(),
+                    &CastOptions::default(),
+                )?);
+            } else {
+                // Create a new array with the default value for the field type
+                let default_value = match field.data_type() {
+                    DataType::Int64 => ScalarValue::Int64(Some(0)),
+                    DataType::Utf8 => ScalarValue::Utf8(Some("a".to_string())),
+                    _ => unimplemented!("Unsupported data type: {:?}", field.data_type()),
+                };
+                output_columns
+                    .push(default_value.to_array_of_size(batch.num_rows()).unwrap());
+            }
+        }
+        let batch = RecordBatch::try_new_with_options(
+            Arc::clone(&self.logical_file_schema),
+            output_columns,
+            &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
+        )
+        .unwrap();
+        Ok(batch)
+    }
+
+    fn map_column_statistics(
+        &self,
+        _file_col_statistics: &[ColumnStatistics],
+    ) -> Result<Vec<ColumnStatistics>> {
+        Ok(vec![
+            ColumnStatistics::new_unknown();
+            self.logical_file_schema.fields().len()
+        ])
+    }
+}
+
+// Implement a custom PhysicalExprAdapterFactory that fills in missing columns with the default value for the field type
+#[derive(Debug)]
+struct CustomPhysicalExprAdapterFactory;
+
+impl PhysicalExprAdapterFactory for CustomPhysicalExprAdapterFactory {
+    fn create(
+        &self,
+        logical_file_schema: SchemaRef,
+        physical_file_schema: SchemaRef,
+    ) -> Arc<dyn PhysicalExprAdapter> {
+        Arc::new(CustomPhysicalExprAdapter {
+            logical_file_schema: Arc::clone(&logical_file_schema),
+            physical_file_schema: Arc::clone(&physical_file_schema),
+            inner: Arc::new(DefaultPhysicalExprAdapter::new(
+                logical_file_schema,
+                physical_file_schema,
+            )),
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+struct CustomPhysicalExprAdapter {
+    logical_file_schema: SchemaRef,
+    physical_file_schema: SchemaRef,
+    inner: Arc<dyn PhysicalExprAdapter>,
+}
+
+impl PhysicalExprAdapter for CustomPhysicalExprAdapter {
+    fn rewrite(&self, mut expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
+        expr = expr
+            .transform(|expr| {
+                if let Some(column) = expr.as_any().downcast_ref::<Column>() {
+                    let field_name = column.name();
+                    if self
+                        .physical_file_schema
+                        .field_with_name(field_name)
+                        .ok()
+                        .is_none()
+                    {
+                        let field = self
+                            .logical_file_schema
+                            .field_with_name(field_name)
+                            .map_err(|_| {
+                                datafusion_common::DataFusionError::Plan(format!(
+                                    "Field '{}' not found in logical file schema",
+                                    field_name
+                                ))
+                            })?;
+                        // If the field does not exist, create a default value expression
+                        // Note that we use slightly different logic here to create a default value so that we can see different behavior in tests
+                        let default_value = match field.data_type() {
+                            DataType::Int64 => ScalarValue::Int64(Some(1)),
+                            DataType::Utf8 => ScalarValue::Utf8(Some("b".to_string())),
+                            _ => unimplemented!(
+                                "Unsupported data type: {:?}",
+                                field.data_type()
+                            ),
+                        };
+                        return Ok(Transformed::yes(Arc::new(expressions::Literal::new(
+                            default_value,
+                        ))));
+                    }
+                }
+
+                Ok(Transformed::no(expr))
+            })
+            .data()?;
+        self.inner.rewrite(expr)
+    }
+
+    fn with_partition_values(
+        &self,
+        partition_values: Vec<(FieldRef, ScalarValue)>,
+    ) -> Arc<dyn PhysicalExprAdapter> {
+        assert!(
+            partition_values.is_empty(),
+            "Partition values are not supported in this test"
+        );
+        Arc::new(self.clone())
+    }
+}
+
 #[tokio::test]
 async fn single_file() {
     let batch =
@@ -56,8 +241,22 @@ async fn single_file() {
         Field::new("c2", DataType::Utf8, true),
     ]));
 
-    let ctx = SessionContext::new();
+    let mut cfg = SessionConfig::new()
+        // Disable statistics collection for this test otherwise early pruning makes it hard to demonstrate data adaptation
+        .with_collect_statistics(false)
+        .with_parquet_pruning(false)
+        .with_parquet_page_index_pruning(false);
+    cfg.options_mut().execution.parquet.pushdown_filters = true;
+    let ctx = SessionContext::new_with_config(cfg);
     ctx.register_object_store(store_url.as_ref(), Arc::clone(&store));
+    assert!(
+        !ctx.state()
+            .config_mut()
+            .options_mut()
+            .execution
+            .collect_statistics
+    );
+    assert!(!ctx.state().config().collect_statistics());
 
     let listing_table_config =
         ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap())
@@ -89,4 +288,92 @@ async fn single_file() {
         "+----+----+",
     ];
     assert_batches_eq!(expected, &batches);
+
+    // Test using a custom schema adapter and no explicit physical expr adapter
+    // This should use the custom schema adapter both for projections and predicate pushdown
+    let listing_table_config =
+        ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap())
+            .infer_options(&ctx.state())
+            .await
+            .unwrap()
+            .with_schema(table_schema.clone())
+            .with_schema_adapter_factory(Arc::new(CustomSchemaAdapterFactory));
+    let table = ListingTable::try_new(listing_table_config).unwrap();
+    ctx.deregister_table("t").unwrap();
+    ctx.register_table("t", Arc::new(table)).unwrap();
+    let batches = ctx
+        .sql("SELECT c2, c1 FROM t WHERE c1 = 2 AND c2 = 'a'")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    let expected = [
+        "+----+----+",
+        "| c2 | c1 |",
+        "+----+----+",
+        "| a  | 2  |",
+        "+----+----+",
+    ];
+    assert_batches_eq!(expected, &batches);
+
+    // Do the same test but with a custom physical expr adapter
+    // Now the default schema adapter will be used for projections, but the custom physical expr adapter will be used for predicate pushdown
+    let listing_table_config =
+        ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap())
+            .infer_options(&ctx.state())
+            .await
+            .unwrap()
+            .with_schema(table_schema.clone())
+            .with_physical_expr_adapter_factory(Arc::new(
+                CustomPhysicalExprAdapterFactory,
+            ));
+    let table = ListingTable::try_new(listing_table_config).unwrap();
+    ctx.deregister_table("t").unwrap();
+    ctx.register_table("t", Arc::new(table)).unwrap();
+    let batches = ctx
+        .sql("SELECT c2, c1 FROM t WHERE c1 = 2 AND c2 = 'b'")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    let expected = [
+        "+----+----+",
+        "| c2 | c1 |",
+        "+----+----+",
+        "|    | 2  |",
+        "+----+----+",
+    ];
+    assert_batches_eq!(expected, &batches);
+
+    // If we use both then the custom physical expr adapter will be used for predicate pushdown and the custom schema adapter will be used for projections
+    let listing_table_config =
+        ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap())
+            .infer_options(&ctx.state())
+            .await
+            .unwrap()
+            .with_schema(table_schema.clone())
+            .with_schema_adapter_factory(Arc::new(CustomSchemaAdapterFactory))
+            .with_physical_expr_adapter_factory(Arc::new(
+                CustomPhysicalExprAdapterFactory,
+            ));
+    let table = ListingTable::try_new(listing_table_config).unwrap();
+    ctx.deregister_table("t").unwrap();
+    ctx.register_table("t", Arc::new(table)).unwrap();
+    let batches = ctx
+        .sql("SELECT c2, c1 FROM t WHERE c1 = 2 AND c2 = 'b'")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    let expected = [
+        "+----+----+",
+        "| c2 | c1 |",
+        "+----+----+",
+        "| a  | 2  |",
+        "+----+----+",
+    ];
+    assert_batches_eq!(expected, &batches);
 }