diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index 469c9be6177..c7f92229cf5 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -45,10 +45,10 @@ pub fn get_session_context() -> SessionContext { .build_arc() .expect("could not build runtime environment"); - let factory = VortexFormatFactory::new().with_options(VortexTableOptions { - projection_pushdown: true, - ..Default::default() - }); + let mut options = VortexTableOptions::default(); + options.projection_pushdown = true; + + let factory = VortexFormatFactory::new().with_options(options); let mut session_state_builder = SessionStateBuilder::new() .with_config(SessionConfig::from_env().expect("shouldn't fail")) diff --git a/vortex-datafusion/public-api.lock b/vortex-datafusion/public-api.lock index 9d558c1de48..f2811a9054b 100644 --- a/vortex-datafusion/public-api.lock +++ b/vortex-datafusion/public-api.lock @@ -176,7 +176,7 @@ impl datafusion_datasource::file_format::FileFormatFactory for vortex_datafusion pub fn vortex_datafusion::VortexFormatFactory::as_any(&self) -> &dyn core::any::Any -pub fn vortex_datafusion::VortexFormatFactory::create(&self, _state: &dyn datafusion_session::session::Session, format_options: &std::collections::hash::map::HashMap) -> datafusion_common::error::Result> +pub fn vortex_datafusion::VortexFormatFactory::create(&self, state: &dyn datafusion_session::session::Session, format_options: &std::collections::hash::map::HashMap) -> datafusion_common::error::Result> pub fn vortex_datafusion::VortexFormatFactory::default(&self) -> alloc::sync::Arc @@ -232,7 +232,7 @@ pub fn vortex_datafusion::VortexSource::try_pushdown_projection(&self, projectio pub fn vortex_datafusion::VortexSource::with_batch_size(&self, batch_size: usize) -> alloc::sync::Arc -pub struct vortex_datafusion::VortexTableOptions +#[non_exhaustive] pub struct vortex_datafusion::VortexTableOptions pub vortex_datafusion::VortexTableOptions::footer_initial_read_size_bytes: usize @@ -244,12 +244,6 @@ impl core::clone::Clone for vortex_datafusion::VortexTableOptions pub fn vortex_datafusion::VortexTableOptions::clone(&self) -> vortex_datafusion::VortexTableOptions -impl core::cmp::Eq for vortex_datafusion::VortexTableOptions - -impl core::cmp::PartialEq for vortex_datafusion::VortexTableOptions - -pub fn vortex_datafusion::VortexTableOptions::eq(&self, other: &vortex_datafusion::VortexTableOptions) -> bool - impl core::default::Default for vortex_datafusion::VortexTableOptions pub fn vortex_datafusion::VortexTableOptions::default() -> Self @@ -258,15 +252,27 @@ impl core::fmt::Debug for vortex_datafusion::VortexTableOptions pub fn vortex_datafusion::VortexTableOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result -impl core::marker::StructuralPartialEq for vortex_datafusion::VortexTableOptions +impl datafusion_common::config::ConfigExtension for vortex_datafusion::VortexTableOptions -impl datafusion_common::config::ConfigField for vortex_datafusion::VortexTableOptions +pub const vortex_datafusion::VortexTableOptions::PREFIX: &'static str -pub fn vortex_datafusion::VortexTableOptions::reset(&mut self, key: &str) -> datafusion_common::error::Result<()> +impl datafusion_common::config::ConfigField for vortex_datafusion::VortexTableOptions pub fn vortex_datafusion::VortexTableOptions::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> -pub fn vortex_datafusion::VortexTableOptions::visit(&self, v: &mut V, key_prefix: &str, _description: &'static str) +pub fn vortex_datafusion::VortexTableOptions::visit(&self, v: &mut V, _key_prefix: &str, _description: &'static str) + +impl datafusion_common::config::ExtensionOptions for vortex_datafusion::VortexTableOptions + +pub fn vortex_datafusion::VortexTableOptions::as_any(&self) -> &dyn core::any::Any + +pub fn vortex_datafusion::VortexTableOptions::as_any_mut(&mut self) -> &mut dyn core::any::Any + +pub fn vortex_datafusion::VortexTableOptions::cloned(&self) -> alloc::boxed::Box + +pub fn vortex_datafusion::VortexTableOptions::entries(&self) -> alloc::vec::Vec + +pub fn vortex_datafusion::VortexTableOptions::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> pub trait vortex_datafusion::ExpressionConvertor: core::marker::Send + core::marker::Sync diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index d5c7fe913f1..8f3e49f95be 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -15,8 +15,9 @@ use datafusion_common::DataFusionError; use datafusion_common::GetExt; use datafusion_common::Result as DFResult; use datafusion_common::Statistics; -use datafusion_common::config::ConfigField; -use datafusion_common::config_namespace; +use datafusion_common::config::ConfigExtension; +use datafusion_common::config::ExtensionOptions; +use datafusion_common::extensions_options; use datafusion_common::internal_datafusion_err; use datafusion_common::not_impl_err; use datafusion_common::parsers::CompressionTypeVariant; @@ -130,7 +131,7 @@ impl Debug for VortexFormat { } } -config_namespace! { +extensions_options! { /// Options to configure [`VortexFormat`] and [`VortexSource`]. /// /// These options are usually set on a [`VortexFormatFactory`] and inherited @@ -142,13 +143,26 @@ config_namespace! { /// ```rust /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions}; /// - /// let factory = VortexFormatFactory::new().with_options(VortexTableOptions { - /// projection_pushdown: true, - /// scan_concurrency: Some(8), - /// ..Default::default() - /// }); - /// # let _ = factory; + /// let mut opts = VortexTableOptions::default(); + /// opts.projection_pushdown = true; + /// opts.scan_concurrency = Some(8); + /// + /// let factory = VortexFormatFactory::new().with_options(opts); /// ``` + /// If used through the SQL interface, additional setup is required, by adding it as an extensions to `ConfigOptions`: + /// + /// ```rust + /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions}; + /// use datafusion_common::config::{ConfigOptions, Extensions}; + /// + /// + /// let mut extensions = Extensions::new(); + /// extensions.insert(VortexTableOptions::default()); + /// let mut config = ConfigOptions::new().with_extensions(extensions); + /// config.set("vortex.scan_concurrency", "1"); + /// ``` + /// + /// Or directly on [`SessionConfig`] using `SessionConfig::with_option_extension`. /// /// [`SessionConfig`]: https://docs.rs/datafusion/latest/datafusion/prelude/struct.SessionConfig.html pub struct VortexTableOptions { @@ -172,8 +186,6 @@ config_namespace! { } } -impl Eq for VortexTableOptions {} - /// Registration entry point for the file-backed Vortex integration. /// /// `VortexFormatFactory` is the type most applications use. Register it with a @@ -194,10 +206,10 @@ impl Eq for VortexTableOptions {} /// use datafusion_common::GetExt; /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions}; /// -/// let factory = Arc::new(VortexFormatFactory::new().with_options(VortexTableOptions { -/// projection_pushdown: true, -/// ..Default::default() -/// })); +/// let mut opts = VortexTableOptions::default(); +/// opts.projection_pushdown = true; +/// +/// let factory = Arc::new(VortexFormatFactory::new().with_options(opts)); /// /// let mut state_builder = SessionStateBuilder::new() /// .with_default_features() @@ -259,11 +271,10 @@ impl VortexFormatFactory { /// ```rust /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions}; /// - /// let factory = VortexFormatFactory::new().with_options(VortexTableOptions { - /// projection_pushdown: true, - /// ..Default::default() - /// }); - /// # let _ = factory; + /// let mut opts = VortexTableOptions::default(); + /// opts.projection_pushdown = true; + /// + /// let factory = VortexFormatFactory::new().with_options(opts); /// ``` pub fn with_options(mut self, options: VortexTableOptions) -> Self { self.options = Some(options); @@ -271,14 +282,24 @@ impl VortexFormatFactory { } } +impl ConfigExtension for VortexTableOptions { + const PREFIX: &'static str = "vortex"; +} + impl FileFormatFactory for VortexFormatFactory { #[expect(clippy::disallowed_types, reason = "required by trait signature")] fn create( &self, - _state: &dyn Session, + state: &dyn Session, format_options: &std::collections::HashMap, ) -> DFResult> { - let mut opts = self.options.clone().unwrap_or_default(); + let mut opts = state + .config_options() + .extensions + .get::() + .cloned() + .unwrap_or_default(); + for (key, value) in format_options { if let Some(key) = key.strip_prefix("format.") { opts.set(key, value)?; diff --git a/vortex-datafusion/src/persistent/tests.rs b/vortex-datafusion/src/persistent/tests.rs index 194184e904e..478531c5b43 100644 --- a/vortex-datafusion/src/persistent/tests.rs +++ b/vortex-datafusion/src/persistent/tests.rs @@ -148,7 +148,6 @@ async fn test_query_file(#[values(Some(1), None)] limit: Option) -> anyho #[tokio::test] async fn test_addition_pushdown() -> anyhow::Result<()> { let ctx = TestSessionContext::default(); - dbg!(&ctx.store); ctx.session .sql( diff --git a/vortex-sqllogictest/bin/sqllogictests-runner.rs b/vortex-sqllogictest/bin/sqllogictests-runner.rs index 1dde50ada35..ee98d6cb311 100644 --- a/vortex-sqllogictest/bin/sqllogictests-runner.rs +++ b/vortex-sqllogictest/bin/sqllogictests-runner.rs @@ -7,6 +7,7 @@ use clap::Parser; use datafusion::common::GetExt; use datafusion::datasource::provider::DefaultTableFactory; use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; use datafusion_sqllogictest::DataFusion; use datafusion_sqllogictest::df_value_validator; @@ -24,6 +25,7 @@ use sqllogictest::parse_file; use sqllogictest::strict_column_validator; use vortex::error::VortexExpect; use vortex_datafusion::VortexFormatFactory; +use vortex_datafusion::VortexTableOptions; use vortex_sqllogictest::args::Args; use vortex_sqllogictest::duckdb::DuckDB; use vortex_sqllogictest::duckdb::DuckDBTestError; @@ -77,7 +79,11 @@ async fn main() -> anyhow::Result<()> { let mut errors = vec![]; let factory = Arc::new(VortexFormatFactory::new()); + + let config = SessionConfig::new().with_option_extension(VortexTableOptions::default()); + let session_state_builder = SessionStateBuilder::new() + .with_config(config) .with_default_features() .with_table_factory( factory.get_ext().to_uppercase(), diff --git a/vortex-sqllogictest/slt/tpch/datafusion/tcph.slt b/vortex-sqllogictest/slt/tpch/datafusion/tcph.slt index 7fee809dc9f..a8d2cfc6bba 100644 --- a/vortex-sqllogictest/slt/tpch/datafusion/tcph.slt +++ b/vortex-sqllogictest/slt/tpch/datafusion/tcph.slt @@ -7,4 +7,14 @@ include ../../setup.slt include ./create.slt.no include ../results/*.slt.no -# include ../drop.slt.no \ No newline at end of file + +include ../drop.slt.no + +statement ok +set vortex.projection_pushdown = true + +include ./create.slt.no + +include ../results/*.slt.no + +include ../drop.slt.no \ No newline at end of file