From 2af32c09951086c9a3f49e65f80c99ae9ed68ff6 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Fri, 10 Jul 2026 16:14:24 +0100 Subject: [PATCH 1/2] Some DF fixes Signed-off-by: Adam Gutglick --- vortex-datafusion/src/persistent/opener.rs | 212 ++++++++++++++++----- vortex-datafusion/src/persistent/sink.rs | 60 +++++- vortex-datafusion/src/persistent/source.rs | 26 +-- 3 files changed, 228 insertions(+), 70 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index b1c999656f6..bfd0369b67e 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -10,6 +10,7 @@ use arrow_schema::Schema; use datafusion_common::DataFusionError; use datafusion_common::Result as DFResult; use datafusion_common::ScalarValue; +use datafusion_common::Statistics; use datafusion_common::arrow::array::AsArray; use datafusion_common::arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; @@ -23,11 +24,13 @@ use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; use datafusion_physical_expr::split_conjunction; +use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_adapter::replace_columns_with_literals; -use datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr; -use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_pruning::FilePruner; use futures::FutureExt; use futures::StreamExt; @@ -83,12 +86,12 @@ pub(crate) struct VortexOpener { /// This is the table's schema without partition columns. It may contain fields which do /// not exist in the file, and are supplied by the `schema_adapter_factory`. pub table_schema: TableSchema, - /// A hint for the desired row count of record batches returned from the scan. - pub batch_size: usize, /// If provided, the scan will not return more than this many rows. pub limit: Option, /// A metrics object for tracking performance of the scan. pub metrics_registry: Arc, + /// DataFusion-native metrics exposed through `DataSourceExec`. + pub df_metrics: ExecutionPlanMetricsSet, /// A shared cache of file readers. /// /// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache @@ -123,12 +126,11 @@ impl FileOpener for VortexOpener { let reader = InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); - let file_pruning_predicate = self.file_pruning_predicate.clone(); + let mut file_pruning_predicate = self.file_pruning_predicate.clone(); let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let file_metadata_cache = self.file_metadata_cache.clone(); let unified_file_schema = Arc::clone(self.table_schema.file_schema()); - let batch_size = self.batch_size; let limit = self.limit; let layout_readers = Arc::clone(&self.layout_readers); let natural_split_ranges = Arc::clone(&self.natural_split_ranges); @@ -138,6 +140,10 @@ impl FileOpener for VortexOpener { let expr_convertor = Arc::clone(&self.expression_convertor); let projection_pushdown = self.projection_pushdown; + let predicate_creation_errors = MetricBuilder::new(&self.df_metrics) + .with_category(MetricCategory::Rows) + .global_counter("num_predicate_creation_errors"); + // Replace column access for partition columns with literals #[expect(clippy::disallowed_types)] let literal_value_cols = self @@ -149,6 +155,13 @@ impl FileOpener for VortexOpener { .zip(file.partition_values.clone()) .collect::>(); + let predicate_uses_partition_columns = + file_pruning_predicate.as_ref().is_some_and(|predicate| { + collect_columns(predicate) + .iter() + .any(|column| literal_value_cols.contains_key(column.name())) + }); + if !literal_value_cols.is_empty() { projection = projection.try_map_exprs(|expr| { replace_columns_with_literals(Arc::clone(&expr), &literal_value_cols) @@ -156,24 +169,30 @@ impl FileOpener for VortexOpener { filter = filter .map(|p| replace_columns_with_literals(p, &literal_value_cols)) .transpose()?; + file_pruning_predicate = file_pruning_predicate + .map(|p| replace_columns_with_literals(p, &literal_value_cols)) + .transpose()?; } Ok(async move { - // Create FilePruner when we have a predicate and either dynamic expressions - // or file statistics available. The pruner can eliminate files without - // opening them based on File-level statistics (min/max values per column) + // FilePruner requires a statistics object even when the rewritten predicate + // only contains partition literals. Supply unknown file-column statistics in + // that case so static and dynamic partition predicates can still prune. + let synthetic_statistics = (!file.has_statistics() && predicate_uses_partition_columns) + .then(|| { + file.clone() + .with_statistics(Arc::new(Statistics::new_unknown(&unified_file_schema))) + }); + let pruning_file = synthetic_statistics.as_ref().unwrap_or(&file); + let mut file_pruner = file_pruning_predicate - .filter(|p| { - // Only create pruner if we have dynamic expressions or file statistics - // to work with. Static predicates without stats won't benefit from pruning. - is_dynamic_physical_expr(p) || file.has_statistics() - }) + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) .and_then(|predicate| { FilePruner::try_new( Arc::clone(&predicate), &unified_file_schema, - &file, - Count::default(), + pruning_file, + predicate_creation_errors, ) }); @@ -424,34 +443,12 @@ impl FileOpener for VortexOpener { }) .into_stream() .map_err(|e| exec_datafusion_err!("Failed to create Vortex stream: {e}"))? - .map_ok(move |rb| { - // We try and slice the stream into respecting datafusion's configured batch size. - stream::iter( - (0..rb.num_rows().div_ceil(batch_size * 2)) - .flat_map(move |block_idx| { - let offset = block_idx * batch_size * 2; - - // If we have less than two batches worth of rows left, we keep them together as a single batch. - if rb.num_rows() - offset < 2 * batch_size { - let length = rb.num_rows() - offset; - [Some(rb.slice(offset, length)), None].into_iter() - } else { - let first = rb.slice(offset, batch_size); - let second = rb.slice(offset + batch_size, batch_size); - [Some(first), Some(second)].into_iter() - } - }) - .flatten() - .map(Ok), - ) - }) .map_err(move |e: VortexError| { DataFusionError::External(Box::new(e.with_context(format!( "Failed to read Vortex file: {}", file.object_meta.location )))) }) - .try_flatten() .map(move |batch| { if projector.projection().as_ref().is_empty() { batch @@ -564,6 +561,7 @@ fn split_midpoint_to_byte(split_range: &Range, row_count: u64, total_size: #[cfg(test)] mod tests { + use std::fmt; use std::sync::Arc; use std::sync::LazyLock; @@ -585,9 +583,12 @@ mod tests { use datafusion::physical_expr::planner::logical2physical; use datafusion::physical_expr_adapter::DefaultPhysicalExprAdapterFactory; use datafusion::scalar::ScalarValue; + use datafusion_common::stats::Precision; use datafusion_execution::cache::DefaultFilesMetadataCache; use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions as df_expr; + use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; use datafusion_physical_expr::projection::ProjectionExpr; use insta::assert_snapshot; use itertools::Itertools; @@ -612,6 +613,53 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::default); + #[derive(Debug, Eq, Hash, PartialEq)] + struct SnapshotErrorExpr; + + impl fmt::Display for SnapshotErrorExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "snapshot_error") + } + } + + impl PhysicalExpr for SnapshotErrorExpr { + fn data_type(&self, _input_schema: &Schema) -> DFResult { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> DFResult { + Ok(false) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(self, f) + } + + fn evaluate(&self, _batch: &RecordBatch) -> DFResult { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + Vec::new() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> DFResult { + assert!(children.is_empty()); + Ok(self) + } + + fn snapshot(&self) -> DFResult> { + Err(DataFusionError::Internal( + "intentional snapshot error".to_owned(), + )) + } + } + #[rstest] #[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))] #[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))] @@ -700,9 +748,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -766,6 +814,82 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_file_pruning_replaces_partition_columns_without_file_statistics() + -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let table_schema = TableSchema::new( + Arc::clone(&file_schema), + vec![Arc::new(Field::new("part", DataType::Int32, false))], + ); + + let partition_column = Arc::new(df_expr::Column::new("part", 1)) as PhysicalExprRef; + let predicate = Arc::new(df_expr::BinaryExpr::new( + Arc::clone(&partition_column), + Operator::Gt, + df_expr::lit(ScalarValue::Int32(Some(1))), + )) as PhysicalExprRef; + let dynamic_predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![partition_column], + predicate, + )) as PhysicalExprRef; + + let mut opener = make_opener(object_store, table_schema, None); + opener.file_pruning_predicate = Some(dynamic_predicate); + let df_metrics = opener.df_metrics.clone(); + + // The file does not exist and has no statistics. Replacing `part` with 1 + // makes the predicate false, so pruning must happen before any file I/O. + let mut file = PartitionedFile::new("missing.vortex", 1); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(batches.is_empty()); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(0) + ); + + Ok(()) + } + + #[tokio::test] + async fn test_file_pruning_creation_errors_are_reported() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "metrics/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)])).unwrap(); + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + let mut statistics = Statistics::new_unknown(batch.schema().as_ref()); + statistics.column_statistics[0].null_count = Precision::Exact(0); + let file = PartitionedFile::new(file_path, data_size).with_statistics(Arc::new(statistics)); + + let mut opener = make_opener( + object_store, + TableSchema::from_file_schema(batch.schema()), + None, + ); + opener.file_pruning_predicate = Some(Arc::new(SnapshotErrorExpr)); + let df_metrics = opener.df_metrics.clone(); + + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_eq!( + df_metrics + .clone_inner() + .sum_by_name("num_predicate_creation_errors") + .map(|metric| metric.as_usize()), + Some(1) + ); + + Ok(()) + } + #[tokio::test] async fn test_open_empty_file() -> anyhow::Result<()> { use futures::TryStreamExt; @@ -872,9 +996,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -959,9 +1083,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: TableSchema::from_file_schema(Arc::clone(&table_schema)), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1116,9 +1240,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: table_schema.clone(), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1176,9 +1300,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema: TableSchema::from_file_schema(schema), - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, @@ -1385,9 +1509,9 @@ mod tests { file_pruning_predicate: None, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), table_schema, - batch_size: 100, limit: None, metrics_registry: Arc::new(DefaultMetricsRegistry::default()), + df_metrics: ExecutionPlanMetricsSet::new(), layout_readers: Default::default(), natural_split_ranges: Default::default(), has_output_ordering: false, diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs index 32be0c6873b..ed655b7e5d0 100644 --- a/vortex-datafusion/src/persistent/sink.rs +++ b/vortex-datafusion/src/persistent/sink.rs @@ -20,6 +20,10 @@ use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::TaskContext; use datafusion_physical_plan::DisplayAs; use datafusion_physical_plan::DisplayFormatType; +use datafusion_physical_plan::metrics::Count; +use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_physical_plan::metrics::MetricCategory; use datafusion_physical_plan::metrics::MetricsSet; use futures::StreamExt; use object_store::ObjectStore; @@ -38,15 +42,29 @@ pub struct VortexSink { config: FileSinkConfig, schema: SchemaRef, session: VortexSession, + metrics: ExecutionPlanMetricsSet, + rows_written: Count, + bytes_written: Count, } impl VortexSink { /// Creates a new [`VortexSink`] instance. pub fn new(config: FileSinkConfig, schema: SchemaRef, session: VortexSession) -> Self { + let metrics = ExecutionPlanMetricsSet::new(); + let rows_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + let bytes_written = MetricBuilder::new(&metrics) + .with_category(MetricCategory::Bytes) + .global_counter("bytes_written"); + Self { config, schema, session, + metrics, + rows_written, + bytes_written, } } } @@ -75,7 +93,7 @@ impl DisplayAs for VortexSink { #[async_trait] impl DataSink for VortexSink { fn metrics(&self) -> Option { - None + Some(self.metrics.clone_inner()) } /// Returns the sink schema @@ -161,7 +179,12 @@ impl FileSink for VortexSink { Ok(r) => { let (path, summary) = r?; - row_count += summary.row_count(); + let rows = summary.row_count(); + row_count += rows; + self.rows_written + .add(usize::try_from(rows).unwrap_or(usize::MAX)); + self.bytes_written + .add(usize::try_from(summary.size()).unwrap_or(usize::MAX)); tracing::info!(path = %path, "Successfully written file"); } @@ -234,12 +257,29 @@ mod tests { ) .await?; - ctx.session + let insert = ctx + .session .sql("INSERT INTO my_tbl VALUES ('hello', 1), ('world', 2);") - .await? - .collect() + .await?; + let physical_plan = insert.create_physical_plan().await?; + datafusion_physical_plan::collect(Arc::clone(&physical_plan), ctx.session.task_ctx()) .await?; + let metrics = physical_plan + .metrics() + .ok_or_else(|| anyhow::anyhow!("Vortex insert did not expose sink metrics"))?; + assert_eq!( + metrics + .sum_by_name("rows_written") + .map(|metric| metric.as_usize()), + Some(2) + ); + assert!( + metrics + .sum_by_name("bytes_written") + .is_some_and(|metric| metric.as_usize() > 0) + ); + let batches = ctx .session .sql("SELECT * from my_tbl") @@ -482,10 +522,12 @@ mod tests { file_output_mode: FileOutputMode::SingleFile, }; - let get_sink = || VortexSink { - config: config.clone(), - schema: Arc::clone(table_schema.file_schema()), - session: session.clone(), + let get_sink = || { + VortexSink::new( + config.clone(), + Arc::clone(table_schema.file_schema()), + session.clone(), + ) }; insta::assert_snapshot!(DefaultDisplay(get_sink()).to_string(), @"VortexSink(file_groups=[])"); diff --git a/vortex-datafusion/src/persistent/source.rs b/vortex-datafusion/src/persistent/source.rs index 469759527f6..34746c0ea50 100644 --- a/vortex-datafusion/src/persistent/source.rs +++ b/vortex-datafusion/src/persistent/source.rs @@ -29,7 +29,6 @@ use datafusion_physical_plan::filter_pushdown::PushedDownPredicate; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use object_store::ObjectStore; use object_store::path::Path; -use vortex::error::VortexExpect; use vortex::file::VORTEX_FILE_EXTENSION; use vortex::layout::LayoutReader; use vortex::metrics::DefaultMetricsRegistry; @@ -190,8 +189,8 @@ pub struct VortexSource { /// Subset of predicates that can be pushed down into Vortex scan operations. /// These are expressions that Vortex can efficiently evaluate during scanning. pub(crate) vortex_predicate: Option, - pub(crate) batch_size: Option, - _unused_df_metrics: ExecutionPlanMetricsSet, + /// DataFusion-native metrics exposed through `DataSourceExec`. + df_metrics: ExecutionPlanMetricsSet, /// Shared layout readers, the source only lives as long as one scan. /// /// Sharing the readers allows us to only read every layout once from the file, even across partitions. @@ -228,8 +227,7 @@ impl VortexSource { projection, full_predicate: None, vortex_predicate: None, - batch_size: None, - _unused_df_metrics: Default::default(), + df_metrics: Default::default(), layout_readers: Arc::new(DashMap::default()), natural_split_ranges: Arc::new(DashMap::default()), expression_convertor, @@ -335,10 +333,6 @@ impl VortexSource { base_config: &FileScanConfig, partition: usize, ) -> DFResult { - let batch_size = self - .batch_size - .vortex_expect("batch_size must be supplied to VortexSource"); - let expr_adapter_factory = base_config .expr_adapter_factory .clone() @@ -358,9 +352,9 @@ impl VortexSource { file_pruning_predicate: self.full_predicate.clone(), expr_adapter_factory, table_schema: self.table_schema.clone(), - batch_size, limit: base_config.limit.map(|l| l as u64), metrics_registry: Arc::clone(&self.vx_metrics_registry), + df_metrics: self.df_metrics.clone(), layout_readers: Arc::clone(&self.layout_readers), natural_split_ranges: Arc::clone(&self.natural_split_ranges), has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered, @@ -388,10 +382,9 @@ impl FileSource for VortexSource { )?)) } - fn with_batch_size(&self, batch_size: usize) -> Arc { - let mut source = self.clone(); - source.batch_size = Some(batch_size); - Arc::new(source) + fn with_batch_size(&self, _batch_size: usize) -> Arc { + // DataSourceExec applies BatchSplitStream after the FileSource stream. + Arc::new(self.clone()) } fn filter(&self) -> Option> { @@ -399,7 +392,7 @@ impl FileSource for VortexSource { } fn metrics(&self) -> &ExecutionPlanMetricsSet { - &self._unused_df_metrics + &self.df_metrics } fn file_type(&self) -> &str { @@ -667,12 +660,11 @@ mod tests { inner: DefaultExpressionConvertor::default(), }) as Arc; - let mut source = VortexSource::new( + let source = VortexSource::new( TableSchema::from_file_schema(file_schema), VortexSession::default(), ) .with_expression_convertor(Arc::clone(&expression_convertor)); - source.batch_size = Some(100); let config = FileScanConfigBuilder::new( ObjectStoreUrl::local_filesystem(), From e6599f18cc38464b641a5743c6c476a54fe14872 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 15 Jul 2026 11:55:28 +0100 Subject: [PATCH 2/2] Maintain schema metadata on export Signed-off-by: Adam Gutglick --- vortex-datafusion/src/convert/schema.rs | 31 ++++++++- vortex-datafusion/src/persistent/opener.rs | 77 ++++++++++++++++++++-- 2 files changed, 100 insertions(+), 8 deletions(-) diff --git a/vortex-datafusion/src/convert/schema.rs b/vortex-datafusion/src/convert/schema.rs index 2ec17758593..907fd5cbbe1 100644 --- a/vortex-datafusion/src/convert/schema.rs +++ b/vortex-datafusion/src/convert/schema.rs @@ -56,7 +56,10 @@ pub fn calculate_physical_schema( }) .collect::>>()?; - Ok(Schema::new(fields)) + Ok(Schema::new_with_metadata( + fields, + reference_logical_schema.metadata().clone(), + )) } /// Calculate the physical Arrow type for a field, preferring the logical type when the @@ -246,6 +249,32 @@ mod tests { ); } + #[test] + fn test_schema_metadata_preserved() -> DFResult<()> { + let logical_schema = Schema::new_with_metadata( + vec![Field::new("col", DataType::Int32, false)], + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ); + let dtype = DType::Struct( + StructFields::from_iter([( + "col", + DType::Primitive(PType::I32, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ); + + let physical_schema = + calculate_physical_schema(&dtype, &logical_schema, &ArrowSession::default())?; + + assert_eq!( + physical_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + Ok(()) + } + #[test] fn test_utf8_variants_preserved() { // Non-view string types become view types after roundtrip through DType, diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index bfd0369b67e..61ad58a3e0b 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -111,6 +111,12 @@ pub(crate) struct VortexOpener { impl FileOpener for VortexOpener { fn open(&self, file: PartitionedFile) -> DFResult { + // Calculate the output schema before replacing partition columns with literals so it + // retains the table and partition-field metadata declared by the plan. + let output_schema = Arc::new( + self.projection + .project_schema(self.table_schema.table_schema())?, + ); let session = self.session.clone(); let metrics_registry = Arc::clone(&self.metrics_registry); let labels = vec![ @@ -259,8 +265,6 @@ impl FileOpener for VortexOpener { &session.arrow(), )?); - let projected_physical_schema = projection.project_schema(&unified_file_schema)?; - let expr_adapter = expr_adapter_factory.create( Arc::clone(&unified_file_schema), Arc::clone(&this_file_schema), @@ -287,7 +291,7 @@ impl FileOpener for VortexOpener { expr_convertor.split_projection( projection.clone(), &this_file_schema, - &projected_physical_schema, + output_schema.as_ref(), )? } else { // When projection pushdown is disabled, read only the required columns @@ -304,7 +308,7 @@ impl FileOpener for VortexOpener { // When projection pushdown is enabled, the scan outputs the projected columns. // When disabled, the scan outputs raw columns and the projection is applied after. let scan_reference_schema = if projection_pushdown { - projected_physical_schema + (*output_schema).clone() } else { // Build schema from the raw columns being read let column_indices = projection.column_indices(); @@ -312,7 +316,7 @@ impl FileOpener for VortexOpener { .into_iter() .map(|idx| this_file_schema.field(idx).clone()) .collect(); - Schema::new(fields) + Schema::new_with_metadata(fields, this_file_schema.metadata().clone()) }; let stream_schema = calculate_physical_schema(&scan_dtype, &scan_reference_schema, &session.arrow())?; @@ -450,11 +454,15 @@ impl FileOpener for VortexOpener { )))) }) .map(move |batch| { - if projector.projection().as_ref().is_empty() { + let batch = if projector.projection().as_ref().is_empty() { batch } else { batch.and_then(|b| projector.project_batch(&b)) - } + }?; + + batch + .with_schema(Arc::clone(&output_schema)) + .map_err(Into::into) }) .boxed(); @@ -814,6 +822,61 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_open_preserves_declared_schema_metadata() -> anyhow::Result<()> { + let object_store = Arc::new(InMemory::new()) as Arc; + let file_path = "part=1/file.vortex"; + let batch = record_batch!(("a", Int32, vec![Some(1), Some(2), Some(3)]))?; + let data_size = + write_arrow_to_vortex(Arc::clone(&object_store), file_path, batch.clone()).await?; + + let file_schema = Arc::new( + batch.schema().as_ref().clone().with_metadata( + [("table".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + ); + let table_schema = TableSchema::new( + file_schema, + vec![Arc::new( + Field::new("part", DataType::Int32, false).with_metadata( + [("partition".to_string(), "metadata".to_string())] + .into_iter() + .collect(), + ), + )], + ); + let projection = ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let expected_schema = Arc::new(projection.project_schema(table_schema.table_schema())?); + + assert_eq!( + expected_schema.metadata().get("table"), + Some(&"metadata".to_string()) + ); + assert_eq!( + expected_schema.field(1).metadata().get("partition"), + Some(&"metadata".to_string()) + ); + + for projection_pushdown in [false, true] { + let mut opener = make_opener(Arc::clone(&object_store), table_schema.clone(), None); + opener.projection = projection.clone(); + opener.projection_pushdown = projection_pushdown; + + let mut file = PartitionedFile::new(file_path.to_string(), data_size); + file.partition_values = vec![ScalarValue::Int32(Some(1))]; + let batches = opener.open(file)?.await?.try_collect::>().await?; + + assert!(!batches.is_empty()); + for batch in batches { + assert_eq!(batch.schema().as_ref(), expected_schema.as_ref()); + } + } + + Ok(()) + } + #[tokio::test] async fn test_file_pruning_replaces_partition_columns_without_file_statistics() -> anyhow::Result<()> {