Merge branch 'main' into union-propagation-distinct-count

Author: buraksenn

datafusion-cli/src/object_storage/instrumented.rs

@@ -727,8 +727,8 @@ impl ExecutionPlan for SampleExec {
         Some(self.metrics.clone_inner())
     }
 
-    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
-        let mut stats = self.input.partition_statistics(partition)?;
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+        let mut stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?);
         let ratio = self.upper_bound - self.lower_bound;
 
         // Scale statistics by sampling ratio (inexact due to randomness)
@@ -741,7 +741,7 @@ impl ExecutionPlan for SampleExec {
             .map(|n| (n as f64 * ratio) as usize)
             .to_inexact();
 
-        Ok(stats)
+        Ok(Arc::new(stats))
     }
 
     fn apply_expressions(

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -2218,7 +2218,9 @@ mod tests {
     use crate::test;
     use crate::test_util::{plan_and_collect, populate_csv_partitions};
     use arrow::datatypes::{DataType, TimeUnit};
+    use arrow_schema::FieldRef;
     use datafusion_common::DataFusionError;
+    use datafusion_common::datatype::DataTypeExt;
     use std::error::Error;
     use std::path::PathBuf;
 
@@ -2735,7 +2737,7 @@ mod tests {
     struct MyTypePlanner {}
 
     impl TypePlanner for MyTypePlanner {
-        fn plan_type(&self, sql_type: &ast::DataType) -> Result<Option<DataType>> {
+        fn plan_type_field(&self, sql_type: &ast::DataType) -> Result<Option<FieldRef>> {
             match sql_type {
                 ast::DataType::Datetime(precision) => {
                     let precision = match precision {
@@ -2745,7 +2747,9 @@ mod tests {
                         None | Some(9) => TimeUnit::Nanosecond,
                         _ => unreachable!(),
                     };
-                    Ok(Some(DataType::Timestamp(precision, None)))
+                    Ok(Some(
+                        DataType::Timestamp(precision, None).into_nullable_field_ref(),
+                    ))
                 }
                 _ => Ok(None),
             }

datafusion/core/src/execution/context/mod.rs

@@ -184,12 +184,12 @@ impl ExecutionPlan for CustomExecutionPlan {
         Ok(Box::pin(TestCustomRecordBatchStream { nb_batch: 1 }))
     }
 
-    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
         if partition.is_some() {
-            return Ok(Statistics::new_unknown(&self.schema()));
+            return Ok(Arc::new(Statistics::new_unknown(&self.schema())));
         }
         let batch = TEST_CUSTOM_RECORD_BATCH!().unwrap();
-        Ok(Statistics {
+        Ok(Arc::new(Statistics {
             num_rows: Precision::Exact(batch.num_rows()),
             total_byte_size: Precision::Absent,
             column_statistics: self
@@ -208,7 +208,7 @@ impl ExecutionPlan for CustomExecutionPlan {
                     ..Default::default()
                 })
                 .collect(),
-        })
+        }))
     }
 
     fn apply_expressions(

datafusion/core/tests/custom_sources_cases/mod.rs

@@ -182,11 +182,11 @@ impl ExecutionPlan for StatisticsValidation {
         unimplemented!("This plan only serves for testing statistics")
     }
 
-    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
         if partition.is_some() {
-            Ok(Statistics::new_unknown(&self.schema))
+            Ok(Arc::new(Statistics::new_unknown(&self.schema)))
         } else {
-            Ok(self.stats.clone())
+            Ok(Arc::new(self.stats.clone()))
         }
     }
 
@@ -255,7 +255,7 @@ async fn sql_basic() -> Result<()> {
     let physical_plan = df.create_physical_plan().await.unwrap();
 
     // the statistics should be those of the source
-    assert_eq!(stats, physical_plan.partition_statistics(None)?);
+    assert_eq!(stats, *physical_plan.partition_statistics(None)?);
 
     Ok(())
 }
@@ -295,7 +295,7 @@ async fn sql_limit() -> Result<()> {
                 .collect(),
             total_byte_size: Precision::Absent
         },
-        physical_plan.partition_statistics(None)?
+        *physical_plan.partition_statistics(None)?
     );
 
     let df = ctx
@@ -304,7 +304,7 @@ async fn sql_limit() -> Result<()> {
         .unwrap();
     let physical_plan = df.create_physical_plan().await.unwrap();
     // when the limit is larger than the original number of lines, statistics remain unchanged
-    assert_eq!(stats, physical_plan.partition_statistics(None)?);
+    assert_eq!(stats, *physical_plan.partition_statistics(None)?);
 
     Ok(())
 }
@@ -324,7 +324,7 @@ async fn sql_window() -> Result<()> {
     let result = physical_plan.partition_statistics(None)?;
 
     assert_eq!(stats.num_rows, result.num_rows);
-    let col_stats = result.column_statistics;
+    let col_stats = &result.column_statistics;
     assert_eq!(2, col_stats.len());
     assert_eq!(stats.column_statistics[1], col_stats[0]);
 

datafusion/core/tests/custom_sources_cases/statistics.rs

@@ -17,8 +17,11 @@
 
 use std::sync::Arc;
 
-use arrow::array::{RecordBatch, record_batch};
-use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use arrow::array::{
+    Array, ArrayRef, BooleanArray, Int32Array, Int64Array, RecordBatch, StringArray,
+    StructArray, record_batch,
+};
+use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef};
 use bytes::{BufMut, BytesMut};
 use datafusion::assert_batches_eq;
 use datafusion::common::Result;
@@ -320,6 +323,145 @@ async fn test_physical_expr_adapter_with_non_null_defaults() {
     assert_batches_eq!(expected, &batches);
 }
 
+#[tokio::test]
+async fn test_struct_schema_evolution_projection_and_filter() -> Result<()> {
+    use std::collections::HashMap;
+
+    // Physical struct: {id: Int32, name: Utf8}
+    let physical_struct_fields: Fields = vec![
+        Arc::new(Field::new("id", DataType::Int32, false)),
+        Arc::new(Field::new("name", DataType::Utf8, true)),
+    ]
+    .into();
+
+    let struct_array = StructArray::new(
+        physical_struct_fields.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
+            Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
+        ],
+        None,
+    );
+
+    let physical_schema = Arc::new(Schema::new(vec![Field::new(
+        "s",
+        DataType::Struct(physical_struct_fields),
+        true,
+    )]));
+
+    let batch =
+        RecordBatch::try_new(Arc::clone(&physical_schema), vec![Arc::new(struct_array)])?;
+
+    let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+    let store_url = ObjectStoreUrl::parse("memory://").unwrap();
+    write_parquet(batch, store.clone(), "struct_evolution.parquet").await;
+
+    // Logical struct: {id: Int64?, name: Utf8?, extra: Boolean?} + metadata
+    let logical_struct_fields: Fields = vec![
+        Arc::new(Field::new("id", DataType::Int64, true)),
+        Arc::new(Field::new("name", DataType::Utf8, true)),
+        Arc::new(Field::new("extra", DataType::Boolean, true).with_metadata(
+            HashMap::from([("nested_meta".to_string(), "1".to_string())]),
+        )),
+    ]
+    .into();
+
+    let table_schema = Arc::new(Schema::new(vec![
+        Field::new("s", DataType::Struct(logical_struct_fields), false)
+            .with_metadata(HashMap::from([("top_meta".to_string(), "1".to_string())])),
+    ]));
+
+    let mut cfg = SessionConfig::new()
+        .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));
+
+    let listing_table_config =
+        ListingTableConfig::new(ListingTableUrl::parse("memory:///").unwrap())
+            .infer_options(&ctx.state())
+            .await
+            .unwrap()
+            .with_schema(table_schema.clone())
+            .with_expr_adapter_factory(Arc::new(DefaultPhysicalExprAdapterFactory));
+
+    let table = ListingTable::try_new(listing_table_config).unwrap();
+    ctx.register_table("t", Arc::new(table)).unwrap();
+
+    let batches = ctx
+        .sql("SELECT s FROM t")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(batches.len(), 1);
+
+    // Verify top-level metadata propagation
+    let output_schema = batches[0].schema();
+    let s_field = output_schema.field_with_name("s").unwrap();
+    assert_eq!(
+        s_field.metadata().get("top_meta").map(String::as_str),
+        Some("1")
+    );
+
+    // Verify nested struct type/field propagation + values
+    let s_array = batches[0]
+        .column(0)
+        .as_any()
+        .downcast_ref::<StructArray>()
+        .expect("expected struct array");
+
+    let id_array = s_array
+        .column_by_name("id")
+        .expect("id column")
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .expect("id should be cast to Int64");
+    assert_eq!(id_array.values(), &[1, 2, 3]);
+
+    let extra_array = s_array.column_by_name("extra").expect("extra column");
+    assert_eq!(extra_array.null_count(), 3);
+
+    // Verify nested field metadata propagation
+    let extra_field = match s_field.data_type() {
+        DataType::Struct(fields) => fields
+            .iter()
+            .find(|f| f.name() == "extra")
+            .expect("extra field"),
+        other => panic!("expected struct type for s, got {other:?}"),
+    };
+    assert_eq!(
+        extra_field
+            .metadata()
+            .get("nested_meta")
+            .map(String::as_str),
+        Some("1")
+    );
+
+    // Smoke test: filtering on a missing nested field evaluates correctly
+    let filtered = ctx
+        .sql("SELECT get_field(s, 'extra') AS extra FROM t WHERE get_field(s, 'extra') IS NULL")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(filtered.len(), 1);
+    assert_eq!(filtered[0].num_rows(), 3);
+    let extra = filtered[0]
+        .column(0)
+        .as_any()
+        .downcast_ref::<BooleanArray>()
+        .expect("extra should be a boolean array");
+    assert_eq!(extra.null_count(), 3);
+
+    Ok(())
+}
+
 /// Test demonstrating that a single PhysicalExprAdapterFactory instance can be
 /// reused across multiple ListingTable instances.
 ///

datafusion/core/tests/parquet/expr_adapter.rs

@@ -1191,12 +1191,12 @@ impl ExecutionPlan for StatisticsExec {
         unimplemented!("This plan only serves for testing statistics")
     }
 
-    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
-        Ok(if partition.is_some() {
+    fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
+        Ok(Arc::new(if partition.is_some() {
             Statistics::new_unknown(&self.schema)
         } else {
             self.stats.clone()
-        })
+        }))
     }
 
     fn apply_expressions(