Skip to main content

vortex_datafusion/persistent/
format.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Formatter;
6use std::sync::Arc;
7
8use arrow_schema::Schema;
9use arrow_schema::SchemaRef;
10use async_trait::async_trait;
11use datafusion_catalog::Session;
12use datafusion_common::ColumnStatistics;
13use datafusion_common::DataFusionError;
14use datafusion_common::GetExt;
15use datafusion_common::Result as DFResult;
16use datafusion_common::ScalarValue as DFScalarValue;
17use datafusion_common::Statistics;
18use datafusion_common::config::ConfigExtension;
19use datafusion_common::config::ConfigField;
20use datafusion_common::extensions_options;
21use datafusion_common::internal_datafusion_err;
22use datafusion_common::not_impl_err;
23use datafusion_common::parsers::CompressionTypeVariant;
24use datafusion_common::stats::Precision as DFPrecision;
25use datafusion_common_runtime::SpawnedTask;
26use datafusion_datasource::TableSchema;
27use datafusion_datasource::file::FileSource;
28use datafusion_datasource::file_compression_type::FileCompressionType;
29use datafusion_datasource::file_format::FileFormat;
30use datafusion_datasource::file_format::FileFormatFactory;
31use datafusion_datasource::file_scan_config::FileScanConfig;
32use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
33use datafusion_datasource::file_sink_config::FileSinkConfig;
34use datafusion_datasource::sink::DataSinkExec;
35use datafusion_datasource::source::DataSourceExec;
36use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry;
37use datafusion_expr::dml::InsertOp;
38use datafusion_physical_expr::LexRequirement;
39use datafusion_physical_plan::ExecutionPlan;
40use futures::FutureExt;
41use futures::StreamExt as _;
42use futures::TryStreamExt as _;
43use futures::stream;
44use object_store::ObjectMeta;
45use object_store::ObjectStore;
46use vortex::VortexSessionDefault;
47use vortex::array::arrow::ArrowSessionExt;
48use vortex::array::memory::MemorySessionExt;
49use vortex::dtype::DType;
50use vortex::dtype::Nullability;
51use vortex::dtype::PType;
52use vortex::error::VortexExpect;
53use vortex::error::VortexResult;
54use vortex::error::vortex_err;
55use vortex::expr::stats::Precision;
56use vortex::expr::stats::Stat;
57use vortex::file::EOF_SIZE;
58use vortex::file::MAX_POSTSCRIPT_SIZE;
59use vortex::file::OpenOptionsSessionExt;
60use vortex::file::VORTEX_FILE_EXTENSION;
61use vortex::io::object_store::ObjectStoreReadAt;
62use vortex::io::session::RuntimeSessionExt;
63use vortex::scalar::Scalar;
64use vortex::scalar::ScalarValue as VortexScalarValue;
65use vortex::session::VortexSession;
66
67use super::cache::CachedVortexMetadata;
68use super::sink::VortexSink;
69use super::source::VortexSource;
70use crate::PrecisionExt as _;
71use crate::convert::TryToDataFusion;
72use crate::convert::stats::is_constant_to_distinct_count;
73
74const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE;
75
76/// DataFusion [`FileFormat`] implementation for `.vortex` files.
77///
78/// Most applications do not construct `VortexFormat` directly. Instead, they
79/// register [`VortexFormatFactory`] with a [`SessionContext`] and let
80/// DataFusion instantiate `VortexFormat` as tables are planned.
81///
82/// Construct `VortexFormat` directly when you are wiring a [`ListingTable`] by
83/// hand and need to pass a file format into [`ListingOptions`].
84///
85/// # Example
86///
87/// ```no_run
88/// use std::sync::Arc;
89///
90/// use datafusion::datasource::listing::ListingOptions;
91/// use datafusion::datasource::listing::ListingTable;
92/// use datafusion::datasource::listing::ListingTableConfig;
93/// use datafusion::datasource::listing::ListingTableUrl;
94/// use datafusion::prelude::SessionContext;
95/// use tempfile::tempdir;
96/// use vortex::VortexSessionDefault;
97/// use vortex::session::VortexSession;
98/// use vortex_datafusion::VortexFormat;
99///
100/// # #[tokio::main]
101/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
102/// let ctx = SessionContext::new();
103/// let dir = tempdir()?;
104///
105/// let format = Arc::new(VortexFormat::new(VortexSession::default()));
106/// let table_url = ListingTableUrl::parse(dir.path().to_str().unwrap())?;
107/// let config = ListingTableConfig::new(table_url)
108///     .with_listing_options(
109///         ListingOptions::new(format).with_session_config_options(ctx.state().config()),
110///     )
111///     .infer_schema(&ctx.state())
112///     .await?;
113///
114/// let table = ListingTable::try_new(config)?;
115/// # let _ = table;
116/// # Ok(())
117/// # }
118/// ```
119///
120/// [`SessionContext`]: https://docs.rs/datafusion/latest/datafusion/prelude/struct.SessionContext.html
121/// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html
122/// [`ListingOptions`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingOptions.html
123pub struct VortexFormat {
124    session: VortexSession,
125    opts: VortexTableOptions,
126}
127
128impl Debug for VortexFormat {
129    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
130        f.debug_struct("VortexFormat")
131            .field("opts", &self.opts)
132            .finish()
133    }
134}
135
136extensions_options! {
137    /// Options to configure [`VortexFormat`] and [`VortexSource`].
138    ///
139    /// The API follows DataFusion's built-in Parquet and JSON format factories:
140    /// a format factory may carry customized defaults, the session may carry
141    /// format defaults, and `CREATE EXTERNAL TABLE ... OPTIONS(...)` can
142    /// override individual fields for one table.
143    ///
144    /// [`FileFormatFactory::create`] builds the `VortexTableOptions` copied into
145    /// each [`VortexFormat`] as follows:
146    ///
147    /// 1. If the factory has explicit options from
148    ///    [`VortexFormatFactory::with_options`] or
149    ///    [`VortexFormatFactory::new_with_options`], start from that complete
150    ///    `VortexTableOptions` value. This matches
151    ///    [`ParquetFormatFactory::new_with_options`] and
152    ///    [`JsonFormatFactory::new_with_options`]: factory options replace
153    ///    session defaults; they are not merged with them field-by-field.
154    /// 2. If the factory does not have explicit options, read the session's
155    ///    `vortex` extension at the time `create` is called. This is the value
156    ///    changed by `SET vortex.<option> = ...`.
157    /// 3. If the session has no `vortex` extension, start from
158    ///    `VortexTableOptions::default()`.
159    /// 4. Apply table `OPTIONS(...)` last. Each option overwrites only its
160    ///    matching field, so per-table settings can override either the factory
161    ///    options or the session/default value.
162    ///
163    /// In SQL, session settings use the `vortex.` prefix. Table options use the
164    /// field names directly, the same style as Parquet or JSON table options:
165    ///
166    /// ```text
167    /// SET vortex.predicate_pushdown = false;
168    ///
169    /// CREATE EXTERNAL TABLE t (x BIGINT)
170    /// STORED AS vortex
171    /// LOCATION 's3://bucket/path/'
172    /// OPTIONS(predicate_pushdown 'true');
173    /// ```
174    ///
175    /// # Example
176    ///
177    /// ```rust
178    /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions};
179    ///
180    /// let mut options = VortexTableOptions::default();
181    /// options.predicate_pushdown = true;
182    /// options.projection_pushdown = true;
183    /// options.scan_concurrency = Some(8);
184    ///
185    /// let factory = VortexFormatFactory::new().with_options(options);
186    /// # let _ = factory;
187    /// ```
188    ///
189    /// [`SessionConfig`]: https://docs.rs/datafusion/latest/datafusion/prelude/struct.SessionConfig.html
190    /// [`ParquetFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/parquet/struct.ParquetFormatFactory.html#method.new_with_options
191    /// [`JsonFormatFactory::new_with_options`]: https://docs.rs/datafusion/latest/datafusion/datasource/file_format/json/struct.JsonFormatFactory.html#method.new_with_options
192    pub struct VortexTableOptions {
193        /// The number of bytes to read when parsing a file footer.
194        ///
195        /// Values smaller than `MAX_POSTSCRIPT_SIZE + EOF_SIZE` will be clamped to that minimum
196        /// during footer parsing.
197        pub footer_initial_read_size_bytes: usize, default = DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES
198        /// Whether to enable projection pushdown into the underlying Vortex scan.
199        ///
200        /// When enabled, projection expressions may be partially evaluated during
201        /// the scan. When disabled, Vortex reads only the referenced columns and
202        /// all expressions are evaluated after the scan.
203        ///
204        /// Disabled by default.
205        pub projection_pushdown: bool, default = false
206        /// Whether to enable predicate pushdown into the underlying Vortex scan.
207        ///
208        /// When enabled, supported filters are evaluated during the scan. When
209        /// disabled, DataFusion evaluates filters after the scan, while
210        /// `VortexSource` can still use the full predicate for file pruning.
211        ///
212        /// Enabled by default.
213        pub predicate_pushdown: bool, default = true
214        /// The intra-partition scan concurrency, controlling the number of row splits to process
215        /// concurrently per-thread within each file.
216        ///
217        /// This does not affect the overall parallelism
218        /// across partitions, which is controlled by DataFusion's execution configuration.
219        ///
220        /// Leave as `None` to use Vortex's scan default. Override per session
221        /// with `SET vortex.scan_concurrency = <n>`, or per table with
222        /// `OPTIONS(scan_concurrency '<n>')`.
223        pub scan_concurrency: Option<usize>, default = None
224    }
225}
226
227impl ConfigExtension for VortexTableOptions {
228    const PREFIX: &'static str = "vortex";
229}
230
231/// Registration entry point for the file-backed Vortex integration.
232///
233/// `VortexFormatFactory` is the type most applications use. Register it with a
234/// DataFusion session, and DataFusion will create [`VortexFormat`] values for
235/// `CREATE EXTERNAL TABLE`, [`ListingTable`], and URL-table scans.
236///
237/// The factory stores a [`VortexSession`] and optional factory-level
238/// [`VortexTableOptions`]. When options are set on the factory they act like
239/// customized format defaults, matching DataFusion's Parquet and JSON factory
240/// APIs. Otherwise, `VortexFormatFactory::create` uses the session's `vortex`
241/// options. In both cases, table `OPTIONS(...)` are applied last for the table
242/// being created.
243///
244/// # Example
245///
246/// ```no_run
247/// use std::sync::Arc;
248///
249/// use datafusion::datasource::provider::DefaultTableFactory;
250/// use datafusion::execution::SessionStateBuilder;
251/// use datafusion_common::GetExt;
252/// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions};
253///
254/// let mut options = VortexTableOptions::default();
255/// options.predicate_pushdown = true;
256/// options.projection_pushdown = true;
257///
258/// let factory = Arc::new(VortexFormatFactory::new().with_options(options));
259///
260/// let mut state_builder = SessionStateBuilder::new()
261///     .with_default_features()
262///     .with_table_factory(
263///         factory.get_ext().to_uppercase(),
264///         Arc::new(DefaultTableFactory::new()),
265///     );
266///
267/// if let Some(file_formats) = state_builder.file_formats() {
268///     file_formats.push(factory.clone() as _);
269/// }
270/// ```
271///
272/// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html
273#[derive(Debug)]
274pub struct VortexFormatFactory {
275    session: VortexSession,
276    options: Option<VortexTableOptions>,
277}
278
279impl GetExt for VortexFormatFactory {
280    fn get_ext(&self) -> String {
281        VORTEX_FILE_EXTENSION.to_string()
282    }
283}
284
285impl VortexFormatFactory {
286    /// Creates a factory with a default [`VortexSession`] and no factory-level
287    /// options.
288    ///
289    /// Formats created by this factory start from the session's `vortex`
290    /// options, or from [`VortexTableOptions::default`] if the session does not
291    /// contain them. Table-level `OPTIONS(...)` are still applied last.
292    #[expect(
293        clippy::new_without_default,
294        reason = "FormatFactory defines `default` method, so having `Default` implementation is confusing"
295    )]
296    pub fn new() -> Self {
297        Self {
298            session: VortexSession::default(),
299            options: None,
300        }
301    }
302
303    /// Creates a factory with an explicit session and factory-level options.
304    ///
305    /// The supplied options become the complete starting value for every
306    /// [`VortexFormat`] created by this factory. Session `SET vortex.*` values
307    /// are ignored for these formats, matching DataFusion's built-in
308    /// `new_with_options` factories. Table-level `OPTIONS(...)` are still
309    /// applied last.
310    pub fn new_with_options(session: VortexSession, options: VortexTableOptions) -> Self {
311        Self {
312            session,
313            options: Some(options),
314        }
315    }
316
317    /// Sets factory-level options.
318    ///
319    /// This is the usual way to customize Vortex defaults for every table
320    /// created through the factory. These options replace, rather than merge
321    /// with, session `SET vortex.*` values. Table-level `OPTIONS(...)` are still
322    /// applied last.
323    ///
324    /// # Example
325    ///
326    /// ```rust
327    /// use vortex_datafusion::{VortexFormatFactory, VortexTableOptions};
328    ///
329    /// let mut options = VortexTableOptions::default();
330    /// options.predicate_pushdown = true;
331    /// options.projection_pushdown = true;
332    ///
333    /// let factory = VortexFormatFactory::new().with_options(options);
334    /// # let _ = factory;
335    /// ```
336    pub fn with_options(mut self, options: VortexTableOptions) -> Self {
337        self.options = Some(options);
338        self
339    }
340}
341
342impl FileFormatFactory for VortexFormatFactory {
343    #[expect(clippy::disallowed_types, reason = "required by trait signature")]
344    fn create(
345        &self,
346        state: &dyn Session,
347        format_options: &std::collections::HashMap<String, String>,
348    ) -> DFResult<Arc<dyn FileFormat>> {
349        // This mirrors DataFusion's Parquet/JSON file-format factories:
350        //
351        // 1. Factory options are a complete customized default when present.
352        // 2. Without factory options, use the session's `vortex` extension
353        //    (`SET vortex.* = ...`), falling back to built-in defaults.
354        // 3. Table-level `CREATE EXTERNAL TABLE ... OPTIONS(...)` values apply
355        //    last. DataFusion prefixes file-format options with `format.`
356        //    before passing them to this factory; SQL users write the field
357        //    name directly, e.g. `OPTIONS(predicate_pushdown 'false')`.
358        let mut opts = self
359            .options
360            .clone()
361            .or_else(|| {
362                state
363                    .config_options()
364                    .extensions
365                    .get::<VortexTableOptions>()
366                    .cloned()
367            })
368            .unwrap_or_default();
369        for (key, value) in format_options {
370            if let Some(key) = key.strip_prefix("format.") {
371                ConfigField::set(&mut opts, key, value)?;
372            } else {
373                tracing::trace!("Ignoring option '{key}'");
374            }
375        }
376
377        Ok(Arc::new(VortexFormat::new_with_options(
378            self.session.clone(),
379            opts,
380        )))
381    }
382
383    fn default(&self) -> Arc<dyn FileFormat> {
384        Arc::new(VortexFormat::new(self.session.clone()))
385    }
386}
387
388impl VortexFormat {
389    /// Creates a format with default [`VortexTableOptions`].
390    ///
391    /// Prefer [`VortexFormatFactory`] when registering with a session. Construct
392    /// `VortexFormat` directly when building [`ListingOptions`] manually.
393    ///
394    /// [`ListingOptions`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingOptions.html
395    pub fn new(session: VortexSession) -> Self {
396        Self::new_with_options(session, VortexTableOptions::default())
397    }
398
399    /// Creates a format with explicit [`VortexTableOptions`].
400    pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self {
401        Self { session, opts }
402    }
403
404    /// Returns the format-specific configuration that will be copied into the
405    /// [`VortexSource`] created for a scan.
406    pub fn options(&self) -> &VortexTableOptions {
407        &self.opts
408    }
409}
410
411#[async_trait]
412impl FileFormat for VortexFormat {
413    fn compression_type(&self) -> Option<FileCompressionType> {
414        None
415    }
416
417    fn get_ext(&self) -> String {
418        VORTEX_FILE_EXTENSION.to_string()
419    }
420
421    fn get_ext_with_compression(
422        &self,
423        file_compression_type: &FileCompressionType,
424    ) -> DFResult<String> {
425        match file_compression_type.get_variant() {
426            CompressionTypeVariant::UNCOMPRESSED => Ok(self.get_ext()),
427            _ => Err(DataFusionError::Internal(
428                "Vortex does not support file level compression.".into(),
429            )),
430        }
431    }
432
433    async fn infer_schema(
434        &self,
435        state: &dyn Session,
436        store: &Arc<dyn ObjectStore>,
437        objects: &[ObjectMeta],
438    ) -> DFResult<SchemaRef> {
439        let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
440
441        let mut file_schemas = stream::iter(objects.iter().cloned())
442            .map(|object| {
443                let store = Arc::clone(store);
444                let session = self.session.clone();
445                let opts = self.opts.clone();
446                let cache = Arc::clone(&file_metadata_cache);
447
448                SpawnedTask::spawn(async move {
449                    // Check if we have entry metadata for this file
450                    if let Some(entry) = cache.get(&object.location)
451                        && entry.is_valid_for(&object)
452                        && let Some(cached_vortex) = entry
453                            .file_metadata
454                            .as_any()
455                            .downcast_ref::<CachedVortexMetadata>()
456                    {
457                        let inferred_schema = session
458                            .arrow()
459                            .to_arrow_schema(cached_vortex.footer().dtype())?;
460                        return VortexResult::Ok((object.location, inferred_schema));
461                    }
462
463                    // Not entry or invalid - open the file
464                    let reader = Arc::new(ObjectStoreReadAt::new_with_allocator(
465                        store,
466                        object.location.clone(),
467                        session.handle(),
468                        session.allocator(),
469                    ));
470
471                    let vxf = session
472                        .open_options()
473                        .with_initial_read_size(opts.footer_initial_read_size_bytes)
474                        .with_file_size(object.size)
475                        .open_read(reader)
476                        .await?;
477
478                    // Cache the metadata
479                    let cached_metadata = Arc::new(CachedVortexMetadata::new(&vxf));
480                    let entry = CachedFileMetadataEntry::new(object.clone(), cached_metadata);
481                    cache.put(&object.location, entry);
482
483                    let inferred_schema = session.arrow().to_arrow_schema(vxf.dtype())?;
484                    VortexResult::Ok((object.location, inferred_schema))
485                })
486                .map(|f| f.vortex_expect("Failed to spawn infer_schema"))
487            })
488            .buffer_unordered(state.config_options().execution.meta_fetch_concurrency)
489            .try_collect::<Vec<_>>()
490            .await
491            .map_err(|e| DataFusionError::Execution(format!("Failed to infer schema: {e}")))?;
492
493        // Get consistent order of schemas for `Schema::try_merge`, as some filesystems don't have deterministic listing orders
494        file_schemas.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
495        let file_schemas = file_schemas.into_iter().map(|(_, schema)| schema);
496
497        Ok(Arc::new(Schema::try_merge(file_schemas)?))
498    }
499
500    #[tracing::instrument(skip_all, fields(location = object.location.as_ref()))]
501    async fn infer_stats(
502        &self,
503        state: &dyn Session,
504        store: &Arc<dyn ObjectStore>,
505        table_schema: SchemaRef,
506        object: &ObjectMeta,
507    ) -> DFResult<Statistics> {
508        let object = object.clone();
509        let store = Arc::clone(store);
510        let session = self.session.clone();
511        let opts = self.opts.clone();
512        let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache();
513
514        SpawnedTask::spawn(async move {
515            // Try to get entry metadata first
516            let cached_metadata = file_metadata_cache
517                .get(&object.location)
518                .filter(|entry| entry.is_valid_for(&object))
519                .and_then(|entry| {
520                    entry
521                        .file_metadata
522                        .as_any()
523                        .downcast_ref::<CachedVortexMetadata>()
524                        .map(|m| {
525                            (
526                                m.footer().dtype().clone(),
527                                m.footer().statistics().cloned(),
528                                m.footer().row_count(),
529                            )
530                        })
531                });
532
533            let (dtype, file_stats, row_count) = match cached_metadata {
534                Some(metadata) => metadata,
535                None => {
536                    // Not entry - open the file
537                    let reader = Arc::new(ObjectStoreReadAt::new_with_allocator(
538                        store,
539                        object.location.clone(),
540                        session.handle(),
541                        session.allocator(),
542                    ));
543
544                    let vxf = session
545                        .open_options()
546                        .with_initial_read_size(opts.footer_initial_read_size_bytes)
547                        .with_file_size(object.size)
548                        .open_read(reader)
549                        .await
550                        .map_err(|e| {
551                            DataFusionError::Execution(format!(
552                                "Failed to open Vortex file {}: {e}",
553                                object.location
554                            ))
555                        })?;
556
557                    // Cache the metadata
558                    let file_metadata = Arc::new(CachedVortexMetadata::new(&vxf));
559                    let entry = CachedFileMetadataEntry::new(object.clone(), file_metadata);
560                    file_metadata_cache.put(&object.location, entry);
561
562                    (
563                        vxf.dtype().clone(),
564                        vxf.file_stats().cloned(),
565                        vxf.row_count(),
566                    )
567                }
568            };
569
570            let struct_dtype = dtype
571                .as_struct_fields_opt()
572                .vortex_expect("dtype is not a struct");
573
574            // Evaluate the statistics for each column that we are able to return to DataFusion.
575            let Some(file_stats) = file_stats else {
576                // If the file has no column stats, the best we can do is return a row count.
577                return Ok(Statistics {
578                    num_rows: DFPrecision::Exact(
579                        usize::try_from(row_count)
580                            .map_err(|_| vortex_err!("Row count overflow"))
581                            .vortex_expect("Row count overflow"),
582                    ),
583                    total_byte_size: DFPrecision::Absent,
584                    column_statistics: vec![
585                        ColumnStatistics::default();
586                        table_schema.fields().len()
587                    ],
588                });
589            };
590
591            let mut column_statistics = Vec::with_capacity(table_schema.fields().len());
592
593            for field in table_schema.fields().iter() {
594                // If the column does not exist, continue. This can happen if the schema has evolved
595                // but we have not yet updated the Vortex file.
596                let Some(col_idx) = struct_dtype.find(field.name()) else {
597                    // The default sets all statistics to `Precision<Absent>`.
598                    column_statistics.push(ColumnStatistics::default());
599                    continue;
600                };
601                let (stats_set, stats_dtype) = file_stats.get(col_idx);
602
603                // Update the total size in bytes.
604                let column_size =
605                    stats_set.get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into());
606
607                let target_dtype =
608                    session
609                        .arrow()
610                        .from_arrow_field(field.as_ref())
611                        .map_err(|e| {
612                            DataFusionError::Execution(format!(
613                                "Failed to derive Vortex DType for field {}: {e}",
614                                field.name()
615                            ))
616                        })?;
617                let min = scalar_stat_to_df(
618                    Stat::Min,
619                    stats_set.get(Stat::Min),
620                    stats_dtype,
621                    &target_dtype,
622                );
623
624                let max = scalar_stat_to_df(
625                    Stat::Max,
626                    stats_set.get(Stat::Max),
627                    stats_dtype,
628                    &target_dtype,
629                );
630
631                let null_count = stats_set.get_as::<usize>(Stat::NullCount, &PType::U64.into());
632
633                column_statistics.push(ColumnStatistics {
634                    null_count: null_count.to_df(),
635                    min_value: min.to_df(),
636                    max_value: max.to_df(),
637                    sum_value: DFPrecision::Absent,
638                    distinct_count: is_constant_to_distinct_count(
639                        stats_set.get_as::<bool>(
640                            Stat::IsConstant,
641                            &DType::Bool(Nullability::NonNullable),
642                        ),
643                    ),
644                    byte_size: column_size.to_df(),
645                })
646            }
647
648            let total_byte_size = column_statistics
649                .iter()
650                .fold(DFPrecision::Exact(0), |acc, cs| acc.add(&cs.byte_size));
651
652            Ok(Statistics {
653                num_rows: DFPrecision::Exact(
654                    usize::try_from(row_count)
655                        .map_err(|_| vortex_err!("Row count overflow"))
656                        .vortex_expect("Row count overflow"),
657                ),
658                total_byte_size,
659                column_statistics,
660            })
661        })
662        .await
663        .vortex_expect("Failed to spawn infer_stats")
664    }
665
666    async fn create_physical_plan(
667        &self,
668        state: &dyn Session,
669        file_scan_config: FileScanConfig,
670    ) -> DFResult<Arc<dyn ExecutionPlan>> {
671        let mut source = file_scan_config
672            .file_source()
673            .downcast_ref::<VortexSource>()
674            .cloned()
675            .ok_or_else(|| internal_datafusion_err!("Expected VortexSource"))?;
676
677        source = source
678            .with_file_metadata_cache(state.runtime_env().cache_manager.get_file_metadata_cache());
679
680        let conf = FileScanConfigBuilder::from(file_scan_config)
681            .with_source(Arc::new(source))
682            .build();
683
684        Ok(DataSourceExec::from_data_source(conf))
685    }
686
687    async fn create_writer_physical_plan(
688        &self,
689        input: Arc<dyn ExecutionPlan>,
690        _state: &dyn Session,
691        conf: FileSinkConfig,
692        order_requirements: Option<LexRequirement>,
693    ) -> DFResult<Arc<dyn ExecutionPlan>> {
694        if conf.insert_op != InsertOp::Append {
695            return not_impl_err!("Overwrites are not implemented yet for Vortex");
696        }
697
698        let schema = Arc::clone(conf.output_schema());
699        let sink = Arc::new(VortexSink::new(conf, schema, self.session.clone()));
700
701        Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
702    }
703
704    fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
705        Arc::new(
706            VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()),
707        ) as _
708    }
709}
710
711fn scalar_stat_to_df(
712    stat: Stat,
713    value: Precision<VortexScalarValue>,
714    stats_dtype: &DType,
715    target_dtype: &DType,
716) -> Precision<DFScalarValue> {
717    let Some(stat_dtype) = stat.dtype(stats_dtype) else {
718        return Precision::Absent;
719    };
720
721    value
722        .map(|stat_value| {
723            Scalar::try_new(stat_dtype, Some(stat_value))?
724                .cast(target_dtype)?
725                .try_to_df()
726        })
727        .transpose()
728        .unwrap_or(Precision::Absent)
729}
730
731#[cfg(test)]
732mod tests {
733
734    use super::*;
735    use crate::common_tests::TestSessionContext;
736
737    #[tokio::test]
738    async fn create_table() -> anyhow::Result<()> {
739        let ctx = TestSessionContext::default();
740
741        ctx.session
742            .sql(
743                "CREATE EXTERNAL TABLE my_tbl \
744                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
745                STORED AS vortex  \
746                LOCATION 'table/'",
747            )
748            .await?;
749
750        assert!(ctx.session.table_exist("my_tbl")?);
751
752        Ok(())
753    }
754
755    #[tokio::test]
756    async fn configure_format_source() -> anyhow::Result<()> {
757        let ctx = TestSessionContext::default();
758
759        ctx.session
760            .sql(
761                "CREATE EXTERNAL TABLE my_tbl \
762                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
763                STORED AS vortex \
764                LOCATION 'table/' \
765                OPTIONS( footer_initial_read_size_bytes '12345', predicate_pushdown 'false', scan_concurrency '3' );",
766            )
767            .await?
768            .collect()
769            .await?;
770
771        Ok(())
772    }
773
774    #[test]
775    fn format_plumbs_footer_initial_read_size() {
776        let mut opts = VortexTableOptions::default();
777        ConfigField::set(&mut opts, "footer_initial_read_size_bytes", "12345").unwrap();
778
779        let format = VortexFormat::new_with_options(VortexSession::default(), opts);
780        assert_eq!(format.options().footer_initial_read_size_bytes, 12345);
781    }
782
783    #[test]
784    fn format_plumbs_source_options() -> anyhow::Result<()> {
785        let opts = VortexTableOptions {
786            projection_pushdown: true,
787            predicate_pushdown: false,
788            scan_concurrency: Some(3),
789            ..Default::default()
790        };
791        let format = VortexFormat::new_with_options(VortexSession::default(), opts.clone());
792        let table_schema = TableSchema::from_file_schema(Arc::new(Schema::empty()));
793
794        let source = format.file_source(table_schema);
795        let source = source
796            .downcast_ref::<VortexSource>()
797            .ok_or_else(|| anyhow::anyhow!("expected VortexSource"))?;
798
799        assert_eq!(
800            source.options().projection_pushdown,
801            opts.projection_pushdown
802        );
803        assert_eq!(source.options().predicate_pushdown, opts.predicate_pushdown);
804        assert_eq!(source.options().scan_concurrency, opts.scan_concurrency);
805        Ok(())
806    }
807}