Docs for DataFusion crate (#2064)

gatesn · web-flow · commit 1d36132df3e9 · 2025-01-23T20:05:58.000Z
And remove unused option to disable pushdown Part of #1905
diff --git a/bench-vortex/benches/datafusion.rs b/bench-vortex/benches/datafusion.rs
@@ -23,7 +23,7 @@ use vortex::sampling_compressor::compressors::r#for::FoRCompressor;
 use vortex::sampling_compressor::compressors::CompressorRef;
 use vortex::sampling_compressor::SamplingCompressor;
 use vortex::{ArrayData, Context};
-use vortex_datafusion::memory::{VortexMemTable, VortexMemTableOptions};
+use vortex_datafusion::memory::VortexMemTable;
 
 pub static CTX: LazyLock<Context> = LazyLock::new(|| {
     Context::default().with_encodings([
@@ -157,51 +157,29 @@ fn bench_arrow<M: Measurement>(mut group: BenchmarkGroup<M>, session: &SessionCo
 fn bench_vortex<M: Measurement>(
     mut group: BenchmarkGroup<M>,
     session: &SessionContext,
-    enable_pushdown: bool,
     compress: bool,
 ) {
     let vortex_dataset = toy_dataset_vortex(compress);
-    let vortex_table = Arc::new(VortexMemTable::new(
-        vortex_dataset,
-        VortexMemTableOptions::default().with_pushdown(enable_pushdown),
-    ));
+    let vortex_table = Arc::new(VortexMemTable::new(vortex_dataset));
 
     measure_provider(&mut group, session, vortex_table);
 }
 
 fn bench_datafusion(c: &mut Criterion) {
     bench_arrow(c.benchmark_group("arrow"), &SessionContext::new());
 
-    // compress=true, pushdown enabled
+    // compress=true
     bench_vortex(
-        c.benchmark_group("vortex-pushdown-compressed"),
+        c.benchmark_group("vortex-compressed"),
         &SessionContext::new(),
         true,
-        true,
-    );
-
-    // compress=false, pushdown enabled
-    bench_vortex(
-        c.benchmark_group("vortex-pushdown-uncompressed"),
-        &SessionContext::new(),
-        true,
-        false,
     );
 
-    // compress=true, pushdown disabled
+    // compress=false
     bench_vortex(
-        c.benchmark_group("vortex-nopushdown-compressed"),
+        c.benchmark_group("vortex-uncompressed"),
         &SessionContext::new(),
         false,
-        true,
-    );
-
-    // compress=false, pushdown disabled
-    bench_vortex(
-        c.benchmark_group("vortex-nopushdown-uncompressed"),
-        &SessionContext::new(),
-        false,
-        false,
     );
 }
 
diff --git a/bench-vortex/benches/tpch.rs b/bench-vortex/benches/tpch.rs
@@ -15,13 +15,7 @@ fn benchmark(c: &mut Criterion) {
     let data_dir = DBGen::new(DBGenOptions::default()).generate().unwrap();
 
     let vortex_ctx = runtime
-        .block_on(load_datasets(
-            &data_dir,
-            Format::InMemoryVortex {
-                enable_pushdown: true,
-            },
-            false,
-        ))
+        .block_on(load_datasets(&data_dir, Format::InMemoryVortex, false))
         .unwrap();
     let arrow_ctx = runtime
         .block_on(load_datasets(&data_dir, Format::Arrow, false))
@@ -50,9 +44,7 @@ fn benchmark(c: &mut Criterion) {
                     &vortex_ctx,
                     &sql_queries,
                     q,
-                    Format::InMemoryVortex {
-                        enable_pushdown: true,
-                    },
+                    Format::InMemoryVortex,
                 )
                 .await;
                 assert_eq!(expected_row_count, row_count, "Mismatched row count {row_count} instead of {expected_row_count} in query {q} for in memory pushdown format");
diff --git a/bench-vortex/src/lib.rs b/bench-vortex/src/lib.rs
@@ -61,7 +61,7 @@ pub enum Format {
     Csv,
     Arrow,
     Parquet,
-    InMemoryVortex { enable_pushdown: bool },
+    InMemoryVortex,
     OnDiskVortex { enable_compression: bool },
 }
 
@@ -71,8 +71,8 @@ impl std::fmt::Display for Format {
             Format::Csv => write!(f, "csv"),
             Format::Arrow => write!(f, "arrow"),
             Format::Parquet => write!(f, "parquet"),
-            Format::InMemoryVortex { enable_pushdown } => {
-                write!(f, "in_memory_vortex(pushdown={enable_pushdown})")
+            Format::InMemoryVortex => {
+                write!(f, "in_memory_vortex")
             }
             Format::OnDiskVortex { enable_compression } => {
                 write!(f, "on_disk_vortex(compressed={enable_compression})")
@@ -87,12 +87,7 @@ impl Format {
             Format::Csv => "csv".to_string(),
             Format::Arrow => "arrow".to_string(),
             Format::Parquet => "parquet".to_string(),
-            Format::InMemoryVortex { enable_pushdown } => if *enable_pushdown {
-                "vortex-in-memory-pushdown"
-            } else {
-                "vortex-in-memory"
-            }
-            .to_string(),
+            Format::InMemoryVortex => "vortex-in-memory".to_string(),
             Format::OnDiskVortex { enable_compression } => if *enable_compression {
                 "vortex-file-compressed"
             } else {
diff --git a/bench-vortex/src/tpch/mod.rs b/bench-vortex/src/tpch/mod.rs
@@ -20,7 +20,6 @@ use vortex::file::{VortexWriteOptions, VORTEX_FILE_EXTENSION};
 use vortex::sampling_compressor::SamplingCompressor;
 use vortex::variants::StructArrayTrait;
 use vortex::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant};
-use vortex_datafusion::memory::VortexMemTableOptions;
 use vortex_datafusion::persistent::VortexFormat;
 use vortex_datafusion::SessionContextExt;
 
@@ -64,17 +63,8 @@ pub async fn load_datasets<P: AsRef<Path>>(
                 Format::Parquet => {
                     register_parquet(&context, stringify!($name), &$name, $schema).await
                 }
-                Format::InMemoryVortex {
-                    enable_pushdown, ..
-                } => {
-                    register_vortex(
-                        &context,
-                        stringify!($name),
-                        &$name,
-                        $schema,
-                        enable_pushdown,
-                    )
-                    .await
+                Format::InMemoryVortex => {
+                    register_vortex(&context, stringify!($name), &$name, $schema).await
                 }
                 Format::OnDiskVortex { enable_compression } => {
                     register_vortex_file(
@@ -305,7 +295,6 @@ async fn register_vortex(
     name: &str,
     file: &Path,
     schema: &Schema,
-    enable_pushdown: bool,
 ) -> anyhow::Result<()> {
     let record_batches = session
         .read_csv(
@@ -330,11 +319,7 @@ async fn register_vortex(
     let dtype = chunks[0].dtype().clone();
     let chunked_array = ChunkedArray::try_new(chunks, dtype)?.into_array();
 
-    session.register_mem_vortex_opts(
-        name,
-        chunked_array,
-        VortexMemTableOptions::default().with_pushdown(enable_pushdown),
-    )?;
+    session.register_mem_vortex(name, chunked_array)?;
 
     Ok(())
 }
diff --git a/vortex-datafusion/src/lib.rs b/vortex-datafusion/src/lib.rs
@@ -1,5 +1,5 @@
 //! Connectors to enable DataFusion to read Vortex data.
-
+#![deny(missing_docs)]
 #![allow(clippy::nonminimal_bool)]
 #![allow(clippy::cast_possible_truncation)]
 
@@ -12,7 +12,7 @@ use datafusion_expr::{Expr, Operator};
 use vortex_array::{ArrayDType, ArrayData};
 use vortex_error::vortex_err;
 
-use crate::memory::{VortexMemTable, VortexMemTableOptions};
+use crate::memory::VortexMemTable;
 
 pub mod memory;
 pub mod persistent;
@@ -50,36 +50,17 @@ fn supported_data_types(dt: DataType) -> bool {
     is_supported
 }
 
+/// Extension function to the DataFusion [`SessionContext`] for registering Vortex tables.
 pub trait SessionContextExt {
-    fn register_mem_vortex<S: AsRef<str>>(&self, name: S, array: ArrayData) -> DFResult<()> {
-        self.register_mem_vortex_opts(name, array, VortexMemTableOptions::default())
-    }
+    /// Register an in-memory Vortex [`ArrayData`] as a DataFusion table.
+    fn register_mem_vortex<S: AsRef<str>>(&self, name: S, array: ArrayData) -> DFResult<()>;
 
-    fn register_mem_vortex_opts<S: AsRef<str>>(
-        &self,
-        name: S,
-        array: ArrayData,
-        options: VortexMemTableOptions,
-    ) -> DFResult<()>;
-
-    fn read_mem_vortex(&self, array: ArrayData) -> DFResult<DataFrame> {
-        self.read_mem_vortex_opts(array, VortexMemTableOptions::default())
-    }
-
-    fn read_mem_vortex_opts(
-        &self,
-        array: ArrayData,
-        options: VortexMemTableOptions,
-    ) -> DFResult<DataFrame>;
+    /// Read an in-memory Vortex [`ArrayData`] into a DataFusion [`DataFrame`].
+    fn read_mem_vortex(&self, array: ArrayData) -> DFResult<DataFrame>;
 }
 
 impl SessionContextExt for SessionContext {
-    fn register_mem_vortex_opts<S: AsRef<str>>(
-        &self,
-        name: S,
-        array: ArrayData,
-        options: VortexMemTableOptions,
-    ) -> DFResult<()> {
+    fn register_mem_vortex<S: AsRef<str>>(&self, name: S, array: ArrayData) -> DFResult<()> {
         if !array.dtype().is_struct() {
             return Err(vortex_err!(
                 "Vortex arrays must have struct type, found {}",
@@ -88,16 +69,12 @@ impl SessionContextExt for SessionContext {
             .into());
         }
 
-        let vortex_table = VortexMemTable::new(array, options);
+        let vortex_table = VortexMemTable::new(array);
         self.register_table(name.as_ref(), Arc::new(vortex_table))
             .map(|_| ())
     }
 
-    fn read_mem_vortex_opts(
-        &self,
-        array: ArrayData,
-        options: VortexMemTableOptions,
-    ) -> DFResult<DataFrame> {
+    fn read_mem_vortex(&self, array: ArrayData) -> DFResult<DataFrame> {
         if !array.dtype().is_struct() {
             return Err(vortex_err!(
                 "Vortex arrays must have struct type, found {}",
@@ -106,7 +83,7 @@ impl SessionContextExt for SessionContext {
             .into());
         }
 
-        let vortex_table = VortexMemTable::new(array, options);
+        let vortex_table = VortexMemTable::new(array);
 
         self.read_table(Arc::new(vortex_table))
     }
diff --git a/vortex-datafusion/src/memory/mod.rs b/vortex-datafusion/src/memory/mod.rs
@@ -1,7 +1,8 @@
+//! In-memory implementation of a Vortex table provider.
 mod exec;
 mod plans;
 mod provider;
 mod statistics;
 mod stream;
 
-pub use provider::{VortexMemTable, VortexMemTableOptions};
+pub use provider::VortexMemTable;
diff --git a/vortex-datafusion/src/memory/provider.rs b/vortex-datafusion/src/memory/provider.rs
@@ -32,7 +32,6 @@ use crate::memory::plans::{RowSelectorExec, TakeRowsExec};
 pub struct VortexMemTable {
     array: ChunkedArray,
     schema_ref: SchemaRef,
-    options: VortexMemTableOptions,
 }
 
 impl VortexMemTable {
@@ -41,7 +40,7 @@ impl VortexMemTable {
     /// # Panics
     ///
     /// Creation will panic if the provided array is not of `DType::Struct` type.
-    pub fn new(array: ArrayData, options: VortexMemTableOptions) -> Self {
+    pub fn new(array: ArrayData) -> Self {
         let arrow_schema = infer_schema(array.dtype()).vortex_expect("schema is inferable");
         let schema_ref = SchemaRef::new(arrow_schema);
 
@@ -54,11 +53,7 @@ impl VortexMemTable {
             }
         };
 
-        Self {
-            array,
-            schema_ref,
-            options,
-        }
+        Self { array, schema_ref }
     }
 }
 
@@ -141,15 +136,6 @@ impl TableProvider for VortexMemTable {
         &self,
         filters: &[&Expr],
     ) -> DFResult<Vec<TableProviderFilterPushDown>> {
-        // In the case the caller has configured this provider with filter pushdown disabled,
-        // do not attempt to apply any filters at scan time.
-        if !self.options.enable_pushdown {
-            return Ok(filters
-                .iter()
-                .map(|_| TableProviderFilterPushDown::Unsupported)
-                .collect());
-        }
-
         filters
             .iter()
             .map(|expr| {
@@ -163,27 +149,6 @@ impl TableProvider for VortexMemTable {
     }
 }
 
-/// Optional configurations to pass when loading a [VortexMemTable].
-#[derive(Debug, Clone)]
-pub struct VortexMemTableOptions {
-    pub enable_pushdown: bool,
-}
-
-impl Default for VortexMemTableOptions {
-    fn default() -> Self {
-        Self {
-            enable_pushdown: true,
-        }
-    }
-}
-
-impl VortexMemTableOptions {
-    pub fn with_pushdown(mut self, enable_pushdown: bool) -> Self {
-        self.enable_pushdown = enable_pushdown;
-        self
-    }
-}
-
 /// Construct an operator plan that executes in two stages.
 ///
 /// The first plan stage only materializes the columns related to the provided set of filter
@@ -220,7 +185,6 @@ mod test {
     use vortex_array::array::{PrimitiveArray, StructArray, VarBinViewArray};
     use vortex_array::{ArrayData, IntoArrayData};
 
-    use crate::memory::VortexMemTableOptions;
     use crate::{can_be_pushed_down, SessionContextExt as _};
 
     fn presidents_array() -> ArrayData {
@@ -271,44 +235,6 @@ mod test {
         );
     }
 
-    #[tokio::test]
-    #[cfg_attr(miri, ignore)]
-    async fn test_datafusion_no_pushdown() {
-        let ctx = SessionContext::new();
-
-        let df = ctx
-            .read_mem_vortex_opts(
-                presidents_array(),
-                // Disable pushdown. We run this test to make sure that the naive codepath also
-                // produces correct results and does not panic anywhere.
-                VortexMemTableOptions::default().with_pushdown(false),
-            )
-            .unwrap();
-
-        let distinct_names = df
-            .filter(col("term_start").gt_eq(lit(1795)))
-            .unwrap()
-            .filter(col("term_start").lt(lit(2000)))
-            .unwrap()
-            .aggregate(vec![], vec![count_distinct(col("president"))])
-            .unwrap()
-            .collect()
-            .await
-            .unwrap();
-
-        assert_eq!(distinct_names.len(), 1);
-
-        assert_eq!(
-            *distinct_names[0]
-                .column(0)
-                .as_primitive::<Int64Type>()
-                .values()
-                .first()
-                .unwrap(),
-            4i64
-        );
-    }
-
     #[test]
     fn test_can_be_pushed_down0() {
         let e = BinaryExpr {
diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs
diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs
