add example for RecordBatch adaptation

Author: adriangb

datafusion/physical-expr/src/schema_rewriter.rs

@@ -198,12 +198,16 @@ impl<'a> PhysicalExprSchemaRewriter<'a> {
 
 #[cfg(test)]
 mod tests {
-    use crate::expressions::lit;
+    use crate::expressions::{col, lit};
 
     use super::*;
-    use arrow::datatypes::{DataType, Field, Schema};
-    use datafusion_common::ScalarValue;
+    use arrow::{
+        array::{RecordBatch, RecordBatchOptions},
+        datatypes::{DataType, Field, Schema, SchemaRef},
+    };
+    use datafusion_common::{record_batch, ScalarValue};
     use datafusion_expr::Operator;
+    use itertools::Itertools;
     use std::sync::Arc;
 
     fn create_test_schema() -> (Schema, Schema) {
@@ -369,4 +373,95 @@ mod tests {
             .to_string()
             .contains("Non-nullable column 'b' is missing"));
     }
+
+    /// Stolen from ProjectionExec
+    fn batch_project(
+        expr: Vec<Arc<dyn PhysicalExpr>>,
+        batch: &RecordBatch,
+        schema: SchemaRef,
+    ) -> Result<RecordBatch> {
+        // Records time on drop
+        let arrays = expr
+            .iter()
+            .map(|expr| {
+                expr.evaluate(batch)
+                    .and_then(|v| v.into_array(batch.num_rows()))
+            })
+            .collect::<Result<Vec<_>>>()?;
+
+        if arrays.is_empty() {
+            let options =
+                RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+            RecordBatch::try_new_with_options(Arc::clone(&schema), arrays, &options)
+                .map_err(Into::into)
+        } else {
+            RecordBatch::try_new(Arc::clone(&schema), arrays).map_err(Into::into)
+        }
+    }
+
+    /// Example showing how we can use the `PhysicalExprSchemaRewriter` to adapt RecordBatches during a scan
+    /// to apply projections, type conversions and handling of missing columns all at once.
+    #[test]
+    fn test_adapt_batches() {
+        let physical_batch = record_batch!(
+            ("a", Int32, vec![Some(1), None, Some(3)]),
+            ("extra", Utf8, vec![Some("x"), Some("y"), None])
+        )
+        .unwrap();
+
+        let physical_schema = physical_batch.schema();
+
+        let logical_schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int64, true), // Different type
+            Field::new("b", DataType::Utf8, true),   // Missing from physical
+        ]));
+
+        let projection = vec![
+            col("b", &logical_schema).unwrap(),
+            col("a", &logical_schema).unwrap(),
+        ];
+
+        let rewriter = PhysicalExprSchemaRewriter::new(&physical_schema, &logical_schema);
+
+        let adapted_projection = projection
+            .into_iter()
+            .map(|expr| rewriter.rewrite(expr).unwrap())
+            .collect_vec();
+
+        let adapted_schema = Arc::new(Schema::new(
+            adapted_projection
+                .iter()
+                .map(|expr| expr.return_field(&physical_schema).unwrap())
+                .collect_vec()
+        ));
+
+        let res = batch_project(
+            adapted_projection,
+            &physical_batch,
+            Arc::clone(&adapted_schema),
+        )
+        .unwrap();
+
+        assert_eq!(res.num_columns(), 2);
+        assert_eq!(res.column(0).data_type(), &DataType::Utf8);
+        assert_eq!(res.column(1).data_type(), &DataType::Int64);
+        assert_eq!(
+            res.column(0)
+                .as_any()
+                .downcast_ref::<arrow::array::StringArray>()
+                .unwrap()
+                .iter()
+                .collect_vec(),
+            vec![None, None, None]
+        );
+        assert_eq!(
+            res.column(1)
+                .as_any()
+                .downcast_ref::<arrow::array::Int64Array>()
+                .unwrap()
+                .iter()
+                .collect_vec(),
+            vec![Some(1), None, Some(3)]
+        );
+    }
 }