16972846388

Committed 14 Aug 2025 05:58PM UTC coverage: 87.72% (+0.03%) from 87.694%

Build # 16972846388

Build Type

push

github

Committed by

web-flow

Commit Message

chore(deps): update rust crate async-trait to v0.1.89 (#4245)

This PR contains the following updates:

| Package | Type | Update | Change |
|---|---|---|---|
| [async-trait](https://redirect.github.com/dtolnay/async-trait) |
workspace.dependencies | patch | `0.1.88` -> `0.1.89` |

---

### Release Notes

<details>
<summary>dtolnay/async-trait (async-trait)</summary>

###
[`v0.1.89`](https://redirect.github.com/dtolnay/async-trait/releases/tag/0.1.89)

[Compare
Source](https://redirect.github.com/dtolnay/async-trait/compare/0.1.88...0.1.89)

- Improve IDE functionality
([#&#8203;293](https://redirect.github.com/dtolnay/async-trait/issues/293),
thanks [@&#8203;Veykril](https://redirect.github.com/Veykril))

</details>

---

### Configuration

📅 **Schedule**: Branch creation - At any time (no schedule defined),
Automerge - At any time (no schedule defined).

🚦 **Automerge**: Enabled.

♻ **Rebasing**: Whenever PR becomes conflicted, or you tick the
rebase/retry checkbox.

🔕 **Ignore**: Close this PR and you won't be reminded about this update
again.

---

- [ ] <!-- rebase-check -->If you want to rebase/retry this PR, check
this box

---

This PR was generated by [Mend Renovate](https://mend.io/renovate/).
View the [repository job
log](https://developer.mend.io/github/vortex-data/vortex).

<!--renovate-debug:eyJjcmVhdGVkSW5WZXIiOiI0MS43MS4xIiwidXBkYXRlZEluVmVyIjoiNDEuNzEuMSIsInRhcmdldEJyYW5jaCI6ImRldmVsb3AiLCJsYWJlbHMiOlsiY2hvcmUiXX0=-->

Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com>

Run Details

56267 of 64144 relevant lines covered (87.72%)

630214.06 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

79.73

/vortex-datafusion/src/persistent/format.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::any::Any;
use std::fmt::{Debug, Formatter};
use std::sync::Arc;

use async_trait::async_trait;
use datafusion::arrow::datatypes::{Schema, SchemaRef};
use datafusion::catalog::Session;
use datafusion::common::parsers::CompressionTypeVariant;
use datafusion::common::runtime::SpawnedTask;
use datafusion::common::stats::Precision;
use datafusion::common::{
    ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics,
    config_datafusion_err, not_impl_err,
};
use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
use datafusion::datasource::file_format::{FileFormat, FileFormatFactory};
use datafusion::datasource::physical_plan::{
    FileScanConfig, FileScanConfigBuilder, FileSinkConfig, FileSource,
};
use datafusion::datasource::sink::DataSinkExec;
use datafusion::datasource::source::DataSourceExec;
use datafusion::logical_expr::dml::InsertOp;
use datafusion::physical_expr::LexRequirement;
use datafusion::physical_plan::ExecutionPlan;
use futures::{FutureExt, StreamExt as _, TryStreamExt as _, stream};
use itertools::Itertools;
use object_store::{ObjectMeta, ObjectStore};
use vortex::dtype::arrow::FromArrowType;
use vortex::dtype::{DType, Nullability, PType};
use vortex::error::{VortexExpect, VortexResult, vortex_err};
use vortex::file::VORTEX_FILE_EXTENSION;
use vortex::metrics::VortexMetrics;
use vortex::scalar::Scalar;
use vortex::session::VortexSession;
use vortex::stats;
use vortex::stats::{Stat, StatsSet};

use super::cache::VortexFileCache;
use super::sink::VortexSink;
use super::source::VortexSource;
use crate::PrecisionExt as _;
use crate::convert::TryToDataFusion;

/// Vortex implementation of a DataFusion [`FileFormat`].
pub struct VortexFormat {
    session: Arc<VortexSession>,
    file_cache: VortexFileCache,
    opts: VortexFormatOptions,
}

impl Debug for VortexFormat {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VortexFormat")
            .field("opts", &self.opts)
            .finish()
    }
}

/// Options to configure the [`VortexFormat`].
#[derive(Debug)]
pub struct VortexFormatOptions {
    /// The size of the in-memory [`vortex::file::Footer`] cache.
    pub footer_cache_size_mb: usize,
    /// The size of the in-memory segment cache.
    pub segment_cache_size_mb: usize,
}

impl Default for VortexFormatOptions {
    fn default() -> Self {
        Self {
            footer_cache_size_mb: 64,
            segment_cache_size_mb: 0,
        }
    }
}

/// Minimal factory to create [`VortexFormat`] instances.
#[derive(Default, Debug)]
pub struct VortexFormatFactory {
    session: Arc<VortexSession>,
}

impl GetExt for VortexFormatFactory {
    fn get_ext(&self) -> String {
        VORTEX_FILE_EXTENSION.to_string()
    }
}

impl FileFormatFactory for VortexFormatFactory {
    #[allow(clippy::disallowed_types)]
    fn create(
        &self,
        _state: &dyn Session,
        format_options: &std::collections::HashMap<String, String>,
    ) -> DFResult<Arc<dyn FileFormat>> {
        if !format_options.is_empty() {
            return Err(config_datafusion_err!(
                "Vortex tables don't support any options"
            ));
        }

        Ok(Arc::new(VortexFormat::new(self.session.clone())))
    }

    fn default(&self) -> Arc<dyn FileFormat> {
        Arc::new(VortexFormat::default())
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

impl Default for VortexFormat {
    fn default() -> Self {
        Self::new(Arc::new(VortexSession::default()))
    }
}

impl VortexFormat {
    /// Create a new instance of the [`VortexFormat`].
    pub fn new(session: Arc<VortexSession>) -> Self {
        let opts = VortexFormatOptions::default();
        Self {
            session: session.clone(),
            file_cache: VortexFileCache::new(
                opts.footer_cache_size_mb,
                opts.segment_cache_size_mb,
                session,
            ),
            opts,
        }
    }

    /// Return the format specific configuration
    pub fn options(&self) -> &VortexFormatOptions {
        &self.opts
    }
}

#[async_trait]
impl FileFormat for VortexFormat {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn compression_type(&self) -> Option<FileCompressionType> {
        None
    }

    fn get_ext(&self) -> String {
        VORTEX_FILE_EXTENSION.to_string()
    }

    fn get_ext_with_compression(
        &self,
        file_compression_type: &FileCompressionType,
    ) -> DFResult<String> {
        match file_compression_type.get_variant() {
            CompressionTypeVariant::UNCOMPRESSED => Ok(self.get_ext()),
            _ => Err(DataFusionError::Internal(
                "Vortex does not support file level compression.".into(),
            )),
        }
    }

    async fn infer_schema(
        &self,
        state: &dyn Session,
        store: &Arc<dyn ObjectStore>,
        objects: &[ObjectMeta],
    ) -> DFResult<SchemaRef> {
        let mut file_schemas = stream::iter(objects.iter().cloned())
            .map(|o| {
                let store = store.clone();
                let cache = self.file_cache.clone();
                SpawnedTask::spawn(async move {
                    let vxf = cache.try_get(&o, store).await?;
                    let inferred_schema = vxf.dtype().to_arrow_schema()?;
                    VortexResult::Ok((o.location, inferred_schema))
                })
                .map(|f| f.vortex_expect("Failed to spawn infer_schema"))
            })
            .buffer_unordered(state.config_options().execution.meta_fetch_concurrency)
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| DataFusionError::Execution(format!("Failed to infer schema: {e}")))?;

        // Get consistent order of schemas for `Schema::try_merge`, as some filesystems don't have deterministic listing orders
        file_schemas.sort_by(|(l1, _), (l2, _)| l1.cmp(l2));
        let file_schemas = file_schemas.into_iter().map(|(_, schema)| schema);

        Ok(Arc::new(Schema::try_merge(file_schemas)?))
    }

    #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, fields(location = object.location.as_ref()
    )))]
    async fn infer_stats(
        &self,
        _state: &dyn Session,
        store: &Arc<dyn ObjectStore>,
        table_schema: SchemaRef,
        object: &ObjectMeta,
    ) -> DFResult<Statistics> {
        let object = object.clone();
        let store = store.clone();
        let cache = self.file_cache.clone();
        SpawnedTask::spawn(async move {
            let vxf = cache.try_get(&object, store.clone()).await.map_err(|e| {
                DataFusionError::Execution(format!(
                    "Failed to open Vortex file {}: {e}",
                    object.location
                ))
            })?;

            let struct_dtype = vxf
                .dtype()
                .as_struct()
                .vortex_expect("dtype is not a struct");

            // Evaluate the statistics for each column that we are able to return to DataFusion.
            let Some(file_stats) = vxf.file_stats() else {
                // If the file has no column stats, the best we can do is return a row count.
                return Ok(Statistics {
                    num_rows: Precision::Exact(
                        usize::try_from(vxf.row_count())
                            .map_err(|_| vortex_err!("Row count overflow"))
                            .vortex_expect("Row count overflow"),
                    ),
                    total_byte_size: Precision::Absent,
                    column_statistics: vec![ColumnStatistics::default(); struct_dtype.nfields()],
                });
            };

            let stats = table_schema
                .fields()
                .iter()
                .map(|field| struct_dtype.find(field.name()))
                .map(|idx| match idx {
                    None => StatsSet::default(),
                    Some(id) => file_stats[id].clone(),
                })
                .collect_vec();

            let total_byte_size = stats
                .iter()
                .map(|stats_set| {
                    stats_set
                        .get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into())
                        .unwrap_or_else(|| stats::Precision::inexact(0_usize))
                })
                .fold(stats::Precision::exact(0_usize), |acc, stats_set| {
                    acc.zip(stats_set).map(|(acc, stats_set)| acc + stats_set)
                });

            // Sum up the total byte size across all the columns.
            let total_byte_size = total_byte_size.to_df();

            let column_statistics = stats
                .into_iter()
                .zip(table_schema.fields().iter())
                .map(|(stats_set, field)| {
                    let null_count = stats_set.get_as::<usize>(Stat::NullCount, &PType::U64.into());
                    let min = stats_set.get(Stat::Min).and_then(|n| {
                        n.map(|n| {
                            Scalar::new(
                                Stat::Min
                                    .dtype(&DType::from_arrow(field.as_ref()))
                                    .vortex_expect("must have a valid dtype"),
                                n,
                            )
                            .try_to_df()
                            .ok()
                        })
                        .transpose()
                    });

                    let max = stats_set.get(Stat::Max).and_then(|n| {
                        n.map(|n| {
                            Scalar::new(
                                Stat::Max
                                    .dtype(&DType::from_arrow(field.as_ref()))
                                    .vortex_expect("must have a valid dtype"),
                                n,
                            )
                            .try_to_df()
                            .ok()
                        })
                        .transpose()
                    });

                    ColumnStatistics {
                        null_count: null_count.to_df(),
                        max_value: max.to_df(),
                        min_value: min.to_df(),
                        sum_value: Precision::Absent,
                        distinct_count: stats_set
                            .get_as::<bool>(
                                Stat::IsConstant,
                                &DType::Bool(Nullability::NonNullable),
                            )
                            .and_then(|is_constant| {
                                is_constant.as_exact().map(|_| Precision::Exact(1))
                            })
                            .unwrap_or(Precision::Absent),
                    }
                })
                .collect::<Vec<_>>();

            Ok(Statistics {
                num_rows: Precision::Exact(
                    usize::try_from(vxf.row_count())
                        .map_err(|_| vortex_err!("Row count overflow"))
                        .vortex_expect("Row count overflow"),
                ),
                total_byte_size,
                column_statistics,
            })
        })
        .await
        .vortex_expect("Failed to spawn infer_stats")
    }

    async fn create_physical_plan(
        &self,
        _state: &dyn Session,
        file_scan_config: FileScanConfig,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        if file_scan_config
            .file_groups
            .iter()
            .flat_map(|fg| fg.files())
            .any(|f| f.range.is_some())
        {
            return not_impl_err!("File level partitioning isn't implemented yet for Vortex");
        }

        if !file_scan_config.table_partition_cols.is_empty() {
            return not_impl_err!("Hive style partitioning isn't implemented yet for Vortex");
        }

        if !file_scan_config.output_ordering.is_empty() {
            return not_impl_err!("Vortex doesn't support output ordering");
        }

        let source = VortexSource::new(self.file_cache.clone(), self.session.metrics().clone());
        Ok(DataSourceExec::from_data_source(
            FileScanConfigBuilder::from(file_scan_config)
                .with_source(Arc::new(source))
                .build(),
        ))
    }

    async fn create_writer_physical_plan(
        &self,
        input: Arc<dyn ExecutionPlan>,
        _state: &dyn Session,
        conf: FileSinkConfig,
        order_requirements: Option<LexRequirement>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        if conf.insert_op != InsertOp::Append {
            return not_impl_err!("Overwrites are not implemented yet for Vortex");
        }

        if !conf.table_partition_cols.is_empty() {
            return not_impl_err!("Hive style partitioning isn't implemented yet for Vortex");
        }

        let schema = conf.output_schema().clone();
        let sink = Arc::new(VortexSink::new(conf, schema));

        Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
    }

    fn file_source(&self) -> Arc<dyn FileSource> {
        Arc::new(VortexSource::new(
            self.file_cache.clone(),
            VortexMetrics::default(),
        ))
    }
}

#[cfg(test)]
mod tests {
    use datafusion::execution::SessionStateBuilder;
    use datafusion::prelude::SessionContext;
    use tempfile::TempDir;

    use super::*;
    use crate::persistent::register_vortex_format_factory;

    #[tokio::test]
    async fn create_table() {
        let dir = TempDir::new().unwrap();

        let factory: VortexFormatFactory = Default::default();
        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
        register_vortex_format_factory(factory, &mut session_state_builder);
        let session = SessionContext::new_with_state(session_state_builder.build());

        let df = session
            .sql(&format!(
                "CREATE EXTERNAL TABLE my_tbl \
                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                STORED AS vortex LOCATION '{}'",
                dir.path().to_str().unwrap()
            ))
            .await
            .unwrap();

        assert_eq!(df.count().await.unwrap(), 0);
    }

    #[tokio::test]
    #[should_panic]
    async fn fail_table_config() {
        let dir = TempDir::new().unwrap();

        let factory: VortexFormatFactory = Default::default();
        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
        register_vortex_format_factory(factory, &mut session_state_builder);
        let session = SessionContext::new_with_state(session_state_builder.build());

        session
            .sql(&format!(
                "CREATE EXTERNAL TABLE my_tbl \
                (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                STORED AS vortex LOCATION '{}' \
                OPTIONS( some_key 'value' );",
                dir.path().to_str().unwrap()
            ))
            .await
            .unwrap()
            .collect()
            .await
            .unwrap();
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::any::Any;
5	use std::fmt::{Debug, Formatter};
6	use std::sync::Arc;
7
8	use async_trait::async_trait;
9	use datafusion::arrow::datatypes::{Schema, SchemaRef};
10	use datafusion::catalog::Session;
11	use datafusion::common::parsers::CompressionTypeVariant;
12	use datafusion::common::runtime::SpawnedTask;
13	use datafusion::common::stats::Precision;
14	use datafusion::common::{
15	ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics,
16	config_datafusion_err, not_impl_err,
17	};
18	use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
19	use datafusion::datasource::file_format::{FileFormat, FileFormatFactory};
20	use datafusion::datasource::physical_plan::{
21	FileScanConfig, FileScanConfigBuilder, FileSinkConfig, FileSource,
22	};
23	use datafusion::datasource::sink::DataSinkExec;
24	use datafusion::datasource::source::DataSourceExec;
25	use datafusion::logical_expr::dml::InsertOp;
26	use datafusion::physical_expr::LexRequirement;
27	use datafusion::physical_plan::ExecutionPlan;
28	use futures::{FutureExt, StreamExt as _, TryStreamExt as _, stream};
29	use itertools::Itertools;
30	use object_store::{ObjectMeta, ObjectStore};
31	use vortex::dtype::arrow::FromArrowType;
32	use vortex::dtype::{DType, Nullability, PType};
33	use vortex::error::{VortexExpect, VortexResult, vortex_err};
34	use vortex::file::VORTEX_FILE_EXTENSION;
35	use vortex::metrics::VortexMetrics;
36	use vortex::scalar::Scalar;
37	use vortex::session::VortexSession;
38	use vortex::stats;
39	use vortex::stats::{Stat, StatsSet};
40
41	use super::cache::VortexFileCache;
42	use super::sink::VortexSink;
43	use super::source::VortexSource;
44	use crate::PrecisionExt as _;
45	use crate::convert::TryToDataFusion;
46
47	/// Vortex implementation of a DataFusion [`FileFormat`].
48	pub struct VortexFormat {
49	session: Arc<VortexSession>,
50	file_cache: VortexFileCache,
51	opts: VortexFormatOptions,
52	}
53
54	impl Debug for VortexFormat {
55	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {	×
56	f.debug_struct("VortexFormat")	×
57	.field("opts", &self.opts)	×
58	.finish()	×
59	}	×
60	}
61
62	/// Options to configure the [`VortexFormat`].
63	#[derive(Debug)]
64	pub struct VortexFormatOptions {
65	/// The size of the in-memory [`vortex::file::Footer`] cache.
66	pub footer_cache_size_mb: usize,
67	/// The size of the in-memory segment cache.
68	pub segment_cache_size_mb: usize,
69	}
70
71	impl Default for VortexFormatOptions {
72	fn default() -> Self {	20✔
73	Self {	20✔
74	footer_cache_size_mb: 64,	20✔
75	segment_cache_size_mb: 0,	20✔
76	}	20✔
77	}	20✔
78	}
79
80	/// Minimal factory to create [`VortexFormat`] instances.
81	#[derive(Default, Debug)]
82	pub struct VortexFormatFactory {
83	session: Arc<VortexSession>,
84	}
85
86	impl GetExt for VortexFormatFactory {
87	fn get_ext(&self) -> String {	18✔
88	VORTEX_FILE_EXTENSION.to_string()	18✔
89	}	18✔
90	}
91
92	impl FileFormatFactory for VortexFormatFactory {
93	#[allow(clippy::disallowed_types)]
94	fn create(	3✔
95	&self,	3✔
96	_state: &dyn Session,	3✔
97	format_options: &std::collections::HashMap<String, String>,	3✔
98	) -> DFResult<Arc<dyn FileFormat>> {	3✔
99	if !format_options.is_empty() {	3✔
100	return Err(config_datafusion_err!(	1✔
101	"Vortex tables don't support any options"	1✔
102	));	1✔
103	}	2✔
104
105	Ok(Arc::new(VortexFormat::new(self.session.clone())))	2✔
106	}	3✔
107
108	fn default(&self) -> Arc<dyn FileFormat> {	×
109	Arc::new(VortexFormat::default())	×
110	}	×
111
112	fn as_any(&self) -> &dyn Any {	×
113	self	×
114	}	×
115	}
116
117	impl Default for VortexFormat {
118	fn default() -> Self {	18✔
119	Self::new(Arc::new(VortexSession::default()))	18✔
120	}	18✔
121	}
122
123	impl VortexFormat {
124	/// Create a new instance of the [`VortexFormat`].
125	pub fn new(session: Arc<VortexSession>) -> Self {	20✔
126	let opts = VortexFormatOptions::default();	20✔
127	Self {	20✔
128	session: session.clone(),	20✔
129	file_cache: VortexFileCache::new(	20✔
130	opts.footer_cache_size_mb,	20✔
131	opts.segment_cache_size_mb,	20✔
132	session,	20✔
133	),	20✔
134	opts,	20✔
135	}	20✔
136	}	20✔
137
138	/// Return the format specific configuration
139	pub fn options(&self) -> &VortexFormatOptions {	×
140	&self.opts	×
141	}	×
142	}
143
144	#[async_trait]
145	impl FileFormat for VortexFormat {
146	fn as_any(&self) -> &dyn Any {	×
147	self	×
148	}	×
149
150	fn compression_type(&self) -> Option<FileCompressionType> {	×
151	None	×
152	}	×
153
154	fn get_ext(&self) -> String {	22✔
155	VORTEX_FILE_EXTENSION.to_string()	22✔
156	}	22✔
157
158	fn get_ext_with_compression(	×
159	&self,	×
160	file_compression_type: &FileCompressionType,	×
161	) -> DFResult<String> {	×
162	match file_compression_type.get_variant() {	×
163	CompressionTypeVariant::UNCOMPRESSED => Ok(self.get_ext()),	×
164	_ => Err(DataFusionError::Internal(	×
165	"Vortex does not support file level compression.".into(),	×
166	)),	×
167	}
168	}	×
169
170	async fn infer_schema(
171	&self,
172	state: &dyn Session,
173	store: &Arc<dyn ObjectStore>,
174	objects: &[ObjectMeta],
175	) -> DFResult<SchemaRef> {	2✔
176	let mut file_schemas = stream::iter(objects.iter().cloned())
177	.map(\|o\| {	2✔
178	let store = store.clone();	2✔
179	let cache = self.file_cache.clone();	2✔
180	SpawnedTask::spawn(async move {	2✔
181	let vxf = cache.try_get(&o, store).await?;	2✔
182	let inferred_schema = vxf.dtype().to_arrow_schema()?;	2✔
183	VortexResult::Ok((o.location, inferred_schema))	2✔
184	})	2✔
185	.map(\|f\| f.vortex_expect("Failed to spawn infer_schema"))	2✔
186	})	2✔
187	.buffer_unordered(state.config_options().execution.meta_fetch_concurrency)
188	.try_collect::<Vec<_>>()
189	.await
190	.map_err(\|e\| DataFusionError::Execution(format!("Failed to infer schema: {e}")))?;	×
191
192	// Get consistent order of schemas for `Schema::try_merge`, as some filesystems don't have deterministic listing orders
193	file_schemas.sort_by(\|(l1, _), (l2, _)\| l1.cmp(l2));	×
194	let file_schemas = file_schemas.into_iter().map(\|(_, schema)\| schema);
195
196	Ok(Arc::new(Schema::try_merge(file_schemas)?))
197	}	2✔
198
199	#[cfg_attr(feature = "tracing", tracing::instrument(skip_all, fields(location = object.location.as_ref()
200	)))]
201	async fn infer_stats(
202	&self,
203	_state: &dyn Session,
204	store: &Arc<dyn ObjectStore>,
205	table_schema: SchemaRef,
206	object: &ObjectMeta,
207	) -> DFResult<Statistics> {	24✔
208	let object = object.clone();
209	let store = store.clone();
210	let cache = self.file_cache.clone();
211	SpawnedTask::spawn(async move {	28✔
212	let vxf = cache.try_get(&object, store.clone()).await.map_err(\|e\| {	28✔
213	DataFusionError::Execution(format!(	×
214	"Failed to open Vortex file {}: {e}",	×
215	object.location	×
216	))	×
217	})?;	×
218
219	let struct_dtype = vxf	28✔
220	.dtype()	28✔
221	.as_struct()	28✔
222	.vortex_expect("dtype is not a struct");	28✔
223
224	// Evaluate the statistics for each column that we are able to return to DataFusion.
225	let Some(file_stats) = vxf.file_stats() else {	28✔
226	// If the file has no column stats, the best we can do is return a row count.
227	return Ok(Statistics {
228	num_rows: Precision::Exact(
229	usize::try_from(vxf.row_count())	×
230	.map_err(\|_\| vortex_err!("Row count overflow"))	×
231	.vortex_expect("Row count overflow"),	×
232	),
233	total_byte_size: Precision::Absent,	×
234	column_statistics: vec![ColumnStatistics::default(); struct_dtype.nfields()],	×
235	});
236	};
237
238	let stats = table_schema	28✔
239	.fields()	28✔
240	.iter()	28✔
241	.map(\|field\| struct_dtype.find(field.name()))	168✔
242	.map(\|idx\| match idx {	168✔
243	None => StatsSet::default(),	×
244	Some(id) => file_stats[id].clone(),	168✔
245	})	168✔
246	.collect_vec();	28✔
247
248	let total_byte_size = stats	28✔
249	.iter()	28✔
250	.map(\|stats_set\| {	168✔
251	stats_set	168✔
252	.get_as::<usize>(Stat::UncompressedSizeInBytes, &PType::U64.into())	168✔
253	.unwrap_or_else(\|\| stats::Precision::inexact(0_usize))	168✔
254	})	168✔
255	.fold(stats::Precision::exact(0_usize), \|acc, stats_set\| {	168✔
256	acc.zip(stats_set).map(\|(acc, stats_set)\| acc + stats_set)	168✔
257	});	168✔
258
259	// Sum up the total byte size across all the columns.
260	let total_byte_size = total_byte_size.to_df();	28✔
261
262	let column_statistics = stats	28✔
263	.into_iter()	28✔
264	.zip(table_schema.fields().iter())	28✔
265	.map(\|(stats_set, field)\| {	168✔
266	let null_count = stats_set.get_as::<usize>(Stat::NullCount, &PType::U64.into());	168✔
267	let min = stats_set.get(Stat::Min).and_then(\|n\| {	168✔
268	n.map(\|n\| {	168✔
269	Scalar::new(	168✔
270	Stat::Min	168✔
271	.dtype(&DType::from_arrow(field.as_ref()))	168✔
272	.vortex_expect("must have a valid dtype"),	168✔
273	n,	168✔
274	)	168✔
275	.try_to_df()	168✔
276	.ok()	168✔
277	})	168✔
278	.transpose()	168✔
279	});	168✔
280
281	let max = stats_set.get(Stat::Max).and_then(\|n\| {	168✔
282	n.map(\|n\| {	168✔
283	Scalar::new(	168✔
284	Stat::Max	168✔
285	.dtype(&DType::from_arrow(field.as_ref()))	168✔
286	.vortex_expect("must have a valid dtype"),	168✔
287	n,	168✔
288	)	168✔
289	.try_to_df()	168✔
290	.ok()	168✔
291	})	168✔
292	.transpose()	168✔
293	});	168✔
294
295	ColumnStatistics {
296	null_count: null_count.to_df(),	168✔
297	max_value: max.to_df(),	168✔
298	min_value: min.to_df(),	168✔
299	sum_value: Precision::Absent,	168✔
300	distinct_count: stats_set	168✔
301	.get_as::<bool>(	168✔
302	Stat::IsConstant,	168✔
303	&DType::Bool(Nullability::NonNullable),	168✔
304	)
305	.and_then(\|is_constant\| {	168✔
306	is_constant.as_exact().map(\|_\| Precision::Exact(1))	×
307	})	×
308	.unwrap_or(Precision::Absent),	168✔
309	}
310	})	168✔
311	.collect::<Vec<_>>();	28✔
312
313	Ok(Statistics {
314	num_rows: Precision::Exact(
315	usize::try_from(vxf.row_count())	28✔
316	.map_err(\|_\| vortex_err!("Row count overflow"))	28✔
317	.vortex_expect("Row count overflow"),	28✔
318	),
319	total_byte_size,	28✔
320	column_statistics,	28✔
321	})
322	})	28✔
323	.await
324	.vortex_expect("Failed to spawn infer_stats")
325	}	24✔
326
327	async fn create_physical_plan(
328	&self,
329	_state: &dyn Session,
330	file_scan_config: FileScanConfig,
331	) -> DFResult<Arc<dyn ExecutionPlan>> {	180✔
332	if file_scan_config
333	.file_groups
334	.iter()
335	.flat_map(\|fg\| fg.files())	182✔
336	.any(\|f\| f.range.is_some())	182✔
337	{
338	return not_impl_err!("File level partitioning isn't implemented yet for Vortex");
339	}
340
341	if !file_scan_config.table_partition_cols.is_empty() {
342	return not_impl_err!("Hive style partitioning isn't implemented yet for Vortex");
343	}
344
345	if !file_scan_config.output_ordering.is_empty() {
346	return not_impl_err!("Vortex doesn't support output ordering");
347	}
348
349	let source = VortexSource::new(self.file_cache.clone(), self.session.metrics().clone());
350	Ok(DataSourceExec::from_data_source(
351	FileScanConfigBuilder::from(file_scan_config)
352	.with_source(Arc::new(source))
353	.build(),
354	))
355	}	180✔
356
357	async fn create_writer_physical_plan(
358	&self,
359	input: Arc<dyn ExecutionPlan>,
360	_state: &dyn Session,
361	conf: FileSinkConfig,
362	order_requirements: Option<LexRequirement>,
363	) -> DFResult<Arc<dyn ExecutionPlan>> {	2✔
364	if conf.insert_op != InsertOp::Append {
365	return not_impl_err!("Overwrites are not implemented yet for Vortex");
366	}
367
368	if !conf.table_partition_cols.is_empty() {
369	return not_impl_err!("Hive style partitioning isn't implemented yet for Vortex");
370	}
371
372	let schema = conf.output_schema().clone();
373	let sink = Arc::new(VortexSink::new(conf, schema));
374
375	Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
376	}	2✔
377
378	fn file_source(&self) -> Arc<dyn FileSource> {	180✔
379	Arc::new(VortexSource::new(	180✔
380	self.file_cache.clone(),	180✔
381	VortexMetrics::default(),	180✔
382	))	180✔
383	}	180✔
384	}
385
386	#[cfg(test)]
387	mod tests {
388	use datafusion::execution::SessionStateBuilder;
389	use datafusion::prelude::SessionContext;
390	use tempfile::TempDir;
391
392	use super::*;
393	use crate::persistent::register_vortex_format_factory;
394
395	#[tokio::test]
396	async fn create_table() {	1✔
397	let dir = TempDir::new().unwrap();	1✔
398
399	let factory: VortexFormatFactory = Default::default();	1✔
400	let mut session_state_builder = SessionStateBuilder::new().with_default_features();	1✔
401	register_vortex_format_factory(factory, &mut session_state_builder);	1✔
402	let session = SessionContext::new_with_state(session_state_builder.build());	1✔
403
404	let df = session	1✔
405	.sql(&format!(	1✔
406	"CREATE EXTERNAL TABLE my_tbl \	1✔
407	(c1 VARCHAR NOT NULL, c2 INT NOT NULL) \	1✔
408	STORED AS vortex LOCATION '{}'",	1✔
409	dir.path().to_str().unwrap()	1✔
410	))	1✔
411	.await	1✔
412	.unwrap();	1✔
413
414	assert_eq!(df.count().await.unwrap(), 0);	1✔
415	}	1✔
416
417	#[tokio::test]
418	#[should_panic]
419	async fn fail_table_config() {	1✔
420	let dir = TempDir::new().unwrap();	1✔
421
422	let factory: VortexFormatFactory = Default::default();	1✔
423	let mut session_state_builder = SessionStateBuilder::new().with_default_features();	1✔
424	register_vortex_format_factory(factory, &mut session_state_builder);	1✔
425	let session = SessionContext::new_with_state(session_state_builder.build());	1✔
426
427	session	1✔
428	.sql(&format!(	1✔
429	"CREATE EXTERNAL TABLE my_tbl \	1✔
430	(c1 VARCHAR NOT NULL, c2 INT NOT NULL) \	1✔
431	STORED AS vortex LOCATION '{}' \	1✔
432	OPTIONS( some_key 'value' );",	1✔
433	dir.path().to_str().unwrap()	1✔
434	))	1✔
435	.await	1✔
436	.unwrap()	1✔
437	.collect()	1✔
438	.await	1✔
439	.unwrap();	1✔
440	}	1✔
441	}

vortex-data / vortex / 16972846388

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous