VortexScanExec reports statistics to datafusion

vortex-array/src/array/chunked/variants.rs

@@ -1,5 +1,6 @@
+use vortex_dtype::field::Field;
 use vortex_dtype::DType;
-use vortex_error::vortex_panic;
+use vortex_error::{vortex_err, vortex_panic, VortexResult};
 
 use crate::array::chunked::ChunkedArray;
 use crate::variants::{
@@ -65,8 +66,7 @@ impl StructArrayTrait for ChunkedArray {
     fn field(&self, idx: usize) -> Option<Array> {
         let mut chunks = Vec::with_capacity(self.nchunks());
         for chunk in self.chunks() {
-            let array = chunk.with_dyn(|a| a.as_struct_array().and_then(|s| s.field(idx)))?;
-            chunks.push(array);
+            chunks.push(chunk.with_dyn(|a| a.as_struct_array().and_then(|s| s.field(idx)))?);
         }
 
         let projected_dtype = self.dtype().as_struct().and_then(|s| s.dtypes().get(idx))?;
@@ -81,6 +81,28 @@ impl StructArrayTrait for ChunkedArray {
             .into_array();
         Some(chunked)
     }
+
+    fn project(&self, projection: &[Field]) -> VortexResult<Array> {
+        let mut chunks = Vec::with_capacity(self.nchunks());
+        for chunk in self.chunks() {
+            chunks.push(chunk.with_dyn(|a| {
+                a.as_struct_array()
+                    .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
+                    .project(projection)
+            })?);
+        }
+
+        let projected_dtype = self
+            .dtype()
+            .as_struct()
+            .ok_or_else(|| vortex_err!("Not a struct dtype"))?
+            .project(projection)?;
+        ChunkedArray::try_new(
+            chunks,
+            DType::Struct(projected_dtype, self.dtype().nullability()),
+        )
+        .map(|a| a.into_array())
+    }
 }
 
 impl ListArrayTrait for ChunkedArray {}

vortex-array/src/array/constant/variants.rs

@@ -1,8 +1,9 @@
 use std::iter;
 use std::sync::Arc;
 
+use vortex_dtype::field::Field;
 use vortex_dtype::{DType, PType};
-use vortex_error::{vortex_panic, VortexError, VortexExpect as _};
+use vortex_error::{vortex_panic, VortexError, VortexExpect as _, VortexResult};
 use vortex_scalar::{ExtScalar, Scalar, ScalarValue, StructScalar};
 
 use crate::array::constant::ConstantArray;
@@ -190,6 +191,14 @@ impl StructArrayTrait for ConstantArray {
             .field_by_idx(idx)
             .map(|scalar| ConstantArray::new(scalar, self.len()).into_array())
     }
+
+    fn project(&self, projection: &[Field]) -> VortexResult<Array> {
+        Ok(ConstantArray::new(
+            StructScalar::try_from(self.scalar())?.project(projection)?,
+            self.len(),
+        )
+        .into_array())
+    }
 }
 
 impl ListArrayTrait for ConstantArray {}

vortex-array/src/array/sparse/variants.rs

@@ -1,5 +1,6 @@
+use vortex_dtype::field::Field;
 use vortex_dtype::DType;
-use vortex_error::VortexExpect;
+use vortex_error::{vortex_err, VortexExpect, VortexResult};
 use vortex_scalar::StructScalar;
 
 use crate::array::sparse::SparseArray;
@@ -83,6 +84,24 @@ impl StructArrayTrait for SparseArray {
             .into_array(),
         )
     }
+
+    fn project(&self, projection: &[Field]) -> VortexResult<Array> {
+        let values = self.values().with_dyn(|s| {
+            s.as_struct_array()
+                .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
+                .project(projection)
+        })?;
+        let scalar = StructScalar::try_from(self.fill_value())?.project(projection)?;
+
+        SparseArray::try_new_with_offset(
+            self.indices().clone(),
+            values,
+            self.len(),
+            self.indices_offset(),
+            scalar,
+        )
+        .map(|a| a.into_array())
+    }
 }
 
 impl ListArrayTrait for SparseArray {}

vortex-array/src/array/struct_/mod.rs

@@ -8,7 +8,9 @@ use crate::stats::{ArrayStatisticsCompute, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
 use crate::variants::{ArrayVariants, StructArrayTrait};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use crate::{impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoCanonical};
+use crate::{
+    impl_encoding, Array, ArrayDType, ArrayDef, ArrayTrait, Canonical, IntoArray, IntoCanonical,
+};
 
 mod compute;
 
@@ -99,10 +101,7 @@ impl StructArray {
     /// which specifies the new ordering of columns in the struct. The projection can be used to
     /// perform column re-ordering, deletion, or duplication at a logical level, without any data
     /// copying.
-    ///
-    /// # Panics
-    /// This function will panic an error if the projection references columns not within the
-    /// schema boundaries.
+    #[allow(clippy::same_name_method)]
     pub fn project(&self, projection: &[Field]) -> VortexResult<Self> {
         let mut children = Vec::with_capacity(projection.len());
         let mut names = Vec::with_capacity(projection.len());
@@ -149,6 +148,10 @@ impl StructArrayTrait for StructArray {
                 .unwrap_or_else(|e| vortex_panic!(e, "StructArray: field {} not found", idx))
         })
     }
+
+    fn project(&self, projection: &[Field]) -> VortexResult<Array> {
+        self.project(projection).map(|a| a.into_array())
+    }
 }
 
 impl IntoCanonical for StructArray {

vortex-array/src/variants.rs

@@ -3,8 +3,9 @@
 //! When callers only want to make assumptions about the DType, and not about any specific
 //! encoding, they can use these traits to write encoding-agnostic code.
 
+use vortex_dtype::field::Field;
 use vortex_dtype::{DType, ExtDType, FieldNames};
-use vortex_error::{vortex_panic, VortexExpect as _};
+use vortex_error::{vortex_panic, VortexExpect as _, VortexResult};
 
 use crate::iter::{AccessorRef, VectorizedArrayIter};
 use crate::{Array, ArrayTrait};
@@ -227,6 +228,8 @@ pub trait StructArrayTrait: ArrayTrait {
 
         field_idx.and_then(|field_idx| self.field(field_idx))
     }
+
+    fn project(&self, projection: &[Field]) -> VortexResult<Array>;
 }
 
 pub trait ListArrayTrait: ArrayTrait {}

vortex-datafusion/src/lib.rs

@@ -12,7 +12,11 @@ use arrow_array::RecordBatch;
 use arrow_schema::{DataType, Schema, SchemaRef};
 use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext};
 use datafusion::prelude::{DataFrame, SessionContext};
-use datafusion_common::{exec_datafusion_err, DataFusionError, Result as DFResult};
+use datafusion_common::stats::Precision;
+use datafusion_common::{
+    exec_datafusion_err, ColumnStatistics, DataFusionError, Result as DFResult, ScalarValue,
+    Statistics,
+};
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_expr::{Expr, Operator};
 use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
@@ -21,9 +25,10 @@ use memory::{VortexMemTable, VortexMemTableOptions};
 use persistent::config::VortexTableOptions;
 use persistent::provider::VortexFileTableProvider;
 use vortex::array::ChunkedArray;
+use vortex::stats::{ArrayStatistics, Stat};
 use vortex::{Array, ArrayDType, IntoArrayVariant};
 use vortex_dtype::field::Field;
-use vortex_error::vortex_err;
+use vortex_error::{vortex_err, VortexExpect};
 
 pub mod memory;
 pub mod persistent;
@@ -209,23 +214,21 @@ pub(crate) struct VortexRecordBatchStream {
 impl Stream for VortexRecordBatchStream {
     type Item = DFResult<RecordBatch>;
 
-    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.get_mut();
-        if this.idx >= this.num_chunks {
+    fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        if self.idx >= self.num_chunks {
             return Poll::Ready(None);
         }
 
         // Grab next chunk, project and convert to Arrow.
-        let chunk = this.chunks.chunk(this.idx)?;
-        this.idx += 1;
+        let chunk = self.chunks.chunk(self.idx)?;
+        self.idx += 1;
 
         let struct_array = chunk
-            .clone()
             .into_struct()
             .map_err(|vortex_error| DataFusionError::Execution(format!("{}", vortex_error)))?;
 
         let projected_struct = struct_array
-            .project(&this.projection)
+            .project(&self.projection)
             .map_err(|vortex_err| {
                 exec_datafusion_err!("projection pushdown to Vortex failed: {vortex_err}")
             })?;
@@ -288,4 +291,64 @@ impl ExecutionPlan for VortexScanExec {
                 .collect(),
         }))
     }
+
+    fn statistics(&self) -> DFResult<Statistics> {
+        let projection = self
+            .scan_projection
+            .iter()
+            .copied()
+            .map(Field::from)
+            .collect::<Vec<_>>();
+        let projected = self
+            .array
+            .as_ref()
+            .with_dyn(|a| {
+                a.as_struct_array()
+                    .ok_or_else(|| vortex_err!("Not a struct array"))
+                    .and_then(|s| s.project(&projection))
+            })
+            .map_err(|vortex_err| {
+                exec_datafusion_err!("projection pushdown to Vortex failed: {vortex_err}")
+            })?;
+        let column_statistics = projected.with_dyn(|a| {
+            let struct_arr = a.as_struct_array_unchecked();
+            (0..struct_arr.nfields())
+                .map(|i| {
+                    let arr = struct_arr.field(i).vortex_expect("iterating over field");
+                    ColumnStatistics {
+                        null_count: arr
+                            .statistics()
+                            .get_as::<u64>(Stat::NullCount)
+                            .map(|n| n as usize)
+                            .map(Precision::Exact)
+                            .unwrap_or(Precision::Absent),
+                        max_value: arr
+                            .statistics()
+                            .get(Stat::Max)
+                            .map(|n| {
+                                ScalarValue::try_from(n)
+                                    .vortex_expect("cannot convert scalar to df scalar")
+                            })
+                            .map(Precision::Exact)
+                            .unwrap_or(Precision::Absent),
+                        min_value: arr
+                            .statistics()
+                            .get(Stat::Min)
+                            .map(|n| {
+                                ScalarValue::try_from(n)
+                                    .vortex_expect("cannot convert scalar to df scalar")
+                            })
+                            .map(Precision::Exact)
+                            .unwrap_or(Precision::Absent),
+                        distinct_count: Precision::Absent,
+                    }
+                })
+                .collect()
+        });
+        Ok(Statistics {
+            num_rows: Precision::Exact(self.array.len()),
+            total_byte_size: Precision::Exact(projected.nbytes()),
+            column_statistics,
+        })
+    }
 }

vortex-scalar/src/struct_.rs

@@ -1,7 +1,8 @@
 use std::sync::Arc;
 
+use vortex_dtype::field::Field;
 use vortex_dtype::DType;
-use vortex_error::{vortex_bail, VortexError, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, VortexError, VortexExpect, VortexResult};
 
 use crate::value::ScalarValue;
 use crate::Scalar;
@@ -83,6 +84,34 @@ impl<'a> StructScalar<'a> {
             Ok(Scalar::null(dtype.clone()))
         }
     }
+
+    pub fn project(&self, projection: &[Field]) -> VortexResult<Scalar> {
+        let struct_dtype = self
+            .dtype
+            .as_struct()
+            .ok_or_else(|| vortex_err!("Not a struct dtype"))?;
+        let projected_dtype = struct_dtype.project(projection)?;
+        let new_fields = if let Some(fs) = self.fields() {
+            ScalarValue::List(
+                projection
+                    .iter()
+                    .map(|p| match p {
+                        Field::Name(n) => struct_dtype
+                            .find_name(n)
+                            .vortex_expect("DType has been successfully projected already"),
+                        Field::Index(i) => *i,
+                    })
+                    .map(|i| fs[i].clone())
+                    .collect(),
+            )
+        } else {
+            ScalarValue::Null
+        };
+        Ok(Scalar::new(
+            DType::Struct(projected_dtype, self.dtype().nullability()),
+            new_fields,
+        ))
+    }
 }
 
 impl Scalar {