diff --git a/vortex-datafusion/src/persistent/access_plan.rs b/vortex-datafusion/src/persistent/access_plan.rs index 0e14b6730b0..3fa6ebec280 100644 --- a/vortex-datafusion/src/persistent/access_plan.rs +++ b/vortex-datafusion/src/persistent/access_plan.rs @@ -9,8 +9,8 @@ use vortex::scan::selection::Selection; /// /// `VortexAccessPlan` is the hook to use when an external index or planner /// already knows that only part of a file needs to be scanned. The plan is -/// attached as `extensions` on `PartitionedFile`, and the internal -/// `VortexOpener` applies it before building the Vortex scan. +/// attached as `extensions` on `PartitionedFile`, and the internal Vortex +/// morsel planner applies it before building the Vortex scan. /// /// The current access plan surface is intentionally small: it lets callers /// provide a [`Selection`] that narrows the rows considered by the scan. @@ -56,7 +56,7 @@ impl VortexAccessPlan { /// Applies this access plan to a [`ScanBuilder`]. /// - /// This is used internally by the file opener after it has translated a + /// This is used internally by the morsel planner after it has translated a /// `PartitionedFile` into a Vortex scan. pub fn apply_to_builder(&self, mut scan_builder: ScanBuilder) -> ScanBuilder { let Self { selection } = self; diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs index 53fefacafa4..3d98ffd43d8 100644 --- a/vortex-datafusion/src/persistent/mod.rs +++ b/vortex-datafusion/src/persistent/mod.rs @@ -26,7 +26,7 @@ mod access_plan; mod cache; mod format; pub mod metrics; -mod opener; +pub mod morsel; pub mod reader; mod sink; mod source; diff --git a/vortex-datafusion/src/persistent/morsel/mod.rs b/vortex-datafusion/src/persistent/morsel/mod.rs new file mode 100644 index 00000000000..d281fa4bdc6 --- /dev/null +++ b/vortex-datafusion/src/persistent/morsel/mod.rs @@ -0,0 +1,1030 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![allow(missing_docs)] + +//! Morsel-driven I/O support for Vortex files. + +use std::ops::Range; +use std::sync::Arc; +use std::sync::Weak; + +use arrow_array::RecordBatch; +use arrow_array::RecordBatchOptions; +use arrow_schema::Field; +use arrow_schema::Schema; +use arrow_schema::SchemaRef; +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::exec_datafusion_err; +use datafusion_common::format::MetricCategory; +use datafusion_datasource::PartitionedFile; +use datafusion_datasource::TableSchema; +use datafusion_datasource::morsel::Morsel; +use datafusion_datasource::morsel::MorselPlan; +use datafusion_datasource::morsel::MorselPlanner; +use datafusion_datasource::morsel::Morselizer; +use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; +use datafusion_execution::cache::cache_manager::FileMetadataCache; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::projection::Projector; +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_plan::metrics::ExecutionPlanMetricsSet; +use datafusion_physical_plan::metrics::MetricBuilder; +use datafusion_pruning::FilePruner; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream::BoxStream; +use object_store::path::Path; +use tokio::sync::OnceCell; +use tracing::Instrument; +use vortex::array::VortexSessionExecute; +use vortex::dtype::FieldMask; +use vortex::error::VortexError; +use vortex::error::VortexExpect; +use vortex::file::Footer; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; +use vortex::io::InstrumentedReadAt; +use vortex::io::VortexReadAt; +use vortex::layout::LayoutReader; +use vortex::layout::scan::repeated_scan::RepeatedScan; +use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::split_by::SplitBy; +use vortex::metrics::Label; +use vortex::metrics::MetricsRegistry; +use vortex::session::VortexSession; +use vortex_arrow::ArrowSessionExt; +use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::dash_map::Entry; + +use crate::VortexAccessPlan; +use crate::convert::ExpressionConvertor; +use crate::convert::exprs::ProcessedProjection; +use crate::convert::exprs::make_vortex_predicate; +use crate::convert::schema::calculate_physical_schema; +use crate::metrics::PARTITION_LABEL; +use crate::metrics::PATH_LABEL; +use crate::persistent::cache::CachedVortexMetadata; +use crate::persistent::stream::PrunableStream; +use crate::reader::VortexReaderFactory; + +pub(crate) type NaturalSplitCache = DashMap>>>; + +/// A file's natural split boundaries, with each split's byte assignment precomputed. +/// +/// `row_boundaries` is the sorted, deduplicated split point list covering the whole file. +/// `assignment_bytes` holds the byte that owns each split (see [`split_assignment_byte`]) and is +/// sorted, so translating a DataFusion byte range into rows is two binary searches rather than a +/// scan over every split. +#[derive(Debug)] +pub(crate) struct NaturalSplits { + row_boundaries: Arc<[u64]>, + assignment_bytes: Box<[u64]>, +} + +impl NaturalSplits { + fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self { + let row_count = row_boundaries.last().copied().unwrap_or_default(); + let assignment_bytes = if row_count == 0 { + Box::default() + } else { + row_boundaries + .windows(2) + .enumerate() + .map(|(idx, boundaries)| { + split_assignment_byte( + idx, + &(boundaries[0]..boundaries[1]), + row_count, + total_size, + ) + }) + .collect() + }; + + debug_assert!(assignment_bytes.is_sorted()); + debug_assert_eq!( + assignment_bytes.len() + usize::from(!row_boundaries.is_empty()), + row_boundaries.len() + ); + + Self { + row_boundaries, + assignment_bytes, + } + } +} + +/// Creates morsel planners for Vortex files. +pub struct VortexMorselizer { + pub partition: usize, + pub session: VortexSession, + pub vortex_reader_factory: Arc, + /// Optional table schema projection. The indices are w.r.t. the `table_schema`, which is + /// all fields in the final scan result not including the partition columns. + pub projection: ProjectionExprs, + /// Filter expression optimized for pushdown into Vortex scan operations. + /// This may be a subset of file_pruning_predicate containing only expressions + /// that Vortex can efficiently evaluate. + pub filter: Option, + /// Filter expression used by DataFusion's FilePruner to eliminate files based on + /// statistics and partition values without opening them. + pub file_pruning_predicate: Option, + pub expr_adapter_factory: Arc, + /// 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, + /// Desired row count for record batches returned from the scan. + /// 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 + /// a file reader the first time we read a file. + pub layout_readers: Arc>>, + /// Shared full-file natural split ranges keyed by file path. + pub(crate) natural_splits: Arc, + /// Whether the query has output ordering specified + pub has_output_ordering: bool, + + pub expression_convertor: Arc, + pub file_metadata_cache: Option>, + /// Whether to enable expression pushdown into the underlying Vortex scan. + pub projection_pushdown: bool, + pub scan_concurrency: Option, +} + +impl std::fmt::Debug for VortexMorselizer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("VortexMorselizer") + .field("partition", &self.partition) + .field("session", &self.session) + .field("vortex_reader_factory", &self.vortex_reader_factory) + .field("projection", &self.projection) + .field("filter", &self.filter) + .field("file_pruning_predicate", &self.file_pruning_predicate) + .field("expr_adapter_factory", &self.expr_adapter_factory) + .field("table_schema", &self.table_schema) + .field("limit", &self.limit) + .field("df_metrics", &self.df_metrics) + .field("layout_readers", &self.layout_readers) + .field("natural_splits", &self.natural_splits) + .field("has_output_ordering", &self.has_output_ordering) + .field("file_metadata_cache", &self.file_metadata_cache) + .field("projection_pushdown", &self.projection_pushdown) + .field("scan_concurrency", &self.scan_concurrency) + .finish_non_exhaustive() + } +} + +impl Morselizer for VortexMorselizer { + fn plan_file(&self, file: PartitionedFile) -> DFResult> { + Ok(Box::new(VortexMorselPlanner::try_new(self, file)?)) + } +} + +/// Plans morsels for a Vortex file. +#[derive(Debug)] +pub struct VortexMorselPlanner { + state: State, +} + +impl VortexMorselPlanner { + /// Create new [`VortexMorselPlanner`] + pub fn try_new(morselizer: &VortexMorselizer, 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( + morselizer + .projection + .project_schema(morselizer.table_schema.table_schema())?, + ); + let session = morselizer.session.clone(); + let metrics_registry = Arc::clone(&morselizer.metrics_registry); + let labels = vec![ + Label::new(PATH_LABEL, file.path().to_string()), + Label::new(PARTITION_LABEL, morselizer.partition.to_string()), + ]; + + let mut projection = morselizer.projection.clone(); + let mut filter = morselizer.filter.clone(); + + let reader = morselizer + .vortex_reader_factory + .create_reader(&file, &session)?; + + let reader = + InstrumentedReadAt::new_with_labels(reader, metrics_registry.as_ref(), labels.clone()); + + let mut file_pruning_predicate = morselizer.file_pruning_predicate.clone(); + let expr_adapter_factory = Arc::clone(&morselizer.expr_adapter_factory); + let file_metadata_cache = morselizer.file_metadata_cache.clone(); + + let unified_file_schema = Arc::clone(morselizer.table_schema.file_schema()); + let limit = morselizer.limit; + let layout_readers = Arc::clone(&morselizer.layout_readers); + let natural_splits = Arc::clone(&morselizer.natural_splits); + let has_output_ordering = morselizer.has_output_ordering; + let scan_concurrency = morselizer.scan_concurrency; + + let expr_convertor = Arc::clone(&morselizer.expression_convertor); + let projection_pushdown = morselizer.projection_pushdown; + + let predicate_creation_errors = MetricBuilder::new(&morselizer.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 = morselizer + .table_schema + .table_partition_cols() + .iter() + .map(|f| f.name()) + .cloned() + .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) + })?; + 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()?; + } + + // 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 file_pruner = file_pruning_predicate + .filter(|_| file.has_statistics() || predicate_uses_partition_columns) + .and_then(|predicate| { + FilePruner::try_new( + Arc::clone(&predicate), + &unified_file_schema, + pruning_file, + predicate_creation_errors, + ) + }); + + Ok(Self { + state: State::Start { + state: FileOpenState { + file, + output_schema, + session, + metrics_registry, + labels, + projection, + filter, + reader, + file_pruner, + expr_adapter_factory, + file_metadata_cache, + unified_file_schema, + limit, + layout_readers, + natural_splits, + has_output_ordering, + expr_convertor, + projection_pushdown, + scan_concurrency, + }, + }, + }) + } +} + +struct FileOpenState { + file: PartitionedFile, + output_schema: SchemaRef, + session: VortexSession, + metrics_registry: Arc, + labels: Vec