diff --git a/Cargo.lock b/Cargo.lock index 8df28fc6..d3bb44a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4481,7 +4481,6 @@ dependencies = [ "ahash", "arrow", "arrow-schema", - "bytes", "datafusion", "datafusion-datasource", "divan", diff --git a/src/datafusion/Cargo.toml b/src/datafusion/Cargo.toml index 4fc53407..53e89258 100644 --- a/src/datafusion/Cargo.toml +++ b/src/datafusion/Cargo.toml @@ -16,7 +16,6 @@ datafusion-datasource = { workspace = true } futures = { workspace = true } tokio = { workspace = true } ahash = { workspace = true } -bytes = { workspace = true } log = { workspace = true } object_store = { workspace = true, features = ["http"] } liquid-cache-common = { workspace = true } diff --git a/src/datafusion/src/cache/stats.rs b/src/datafusion/src/cache/stats.rs index ac4b6e73..d1947fc6 100644 --- a/src/datafusion/src/cache/stats.rs +++ b/src/datafusion/src/cache/stats.rs @@ -160,7 +160,7 @@ impl LiquidCacheParquet { #[cfg(test)] mod tests { - use std::io::Read; + use std::fs::File; use crate::cache::{ParquetFileIdentity, id::BatchID}; @@ -169,7 +169,6 @@ mod tests { array::{Array, AsArray}, datatypes::UInt64Type, }; - use bytes::Bytes; use liquid_cache::{ cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, @@ -235,14 +234,11 @@ mod tests { } } - let mut tmp_file = NamedTempFile::new()?; + let tmp_file = NamedTempFile::new()?; cache.write_stats(tmp_file.path())?; // Read and verify stats - let mut bytes = Vec::new(); - tmp_file.read_to_end(&mut bytes)?; - let bytes = Bytes::from(bytes); - let reader = ParquetRecordBatchReader::try_new(bytes, 8192)?; + let reader = ParquetRecordBatchReader::try_new(File::open(tmp_file.path())?, 8192)?; let batch = reader.into_iter().next().unwrap()?; assert_eq!(batch.num_rows(), num_rows); diff --git a/src/datafusion/src/reader/plantime/mod.rs b/src/datafusion/src/reader/plantime/mod.rs index ddc34ead..e9d22f93 100644 --- a/src/datafusion/src/reader/plantime/mod.rs +++ b/src/datafusion/src/reader/plantime/mod.rs @@ -1,11 +1,8 @@ -#[cfg(test)] -pub(crate) use source::CachedMetaReaderFactory; pub use source::LiquidParquetSource; -pub(crate) use source::ParquetMetadataCacheReader; mod morselizer; mod row_filter; mod source; -pub(crate) use morselizer::{LiquidFileMetrics, LiquidMorselizer}; +pub(crate) use morselizer::{LiquidFileMetrics, LiquidFileReaderFactory, LiquidMorselizer}; pub use row_filter::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}; diff --git a/src/datafusion/src/reader/plantime/morselizer.rs b/src/datafusion/src/reader/plantime/morselizer.rs index 6334d39f..a8ff2b80 100644 --- a/src/datafusion/src/reader/plantime/morselizer.rs +++ b/src/datafusion/src/reader/plantime/morselizer.rs @@ -6,7 +6,7 @@ use datafusion::{ datasource::{ listing::{FileRange, PartitionedFile}, physical_plan::{ - ParquetFileMetrics, + ParquetFileMetrics, ParquetFileReaderFactory, parquet::{ BloomFilterStatistics, PagePruningAccessPlanFilter, ParquetAccessPlan, RowGroupAccessPlanFilter, @@ -15,6 +15,7 @@ use datafusion::{ table_schema::TableSchema, }, error::Result, + execution::object_store::ObjectStoreUrl, physical_expr::{ DynamicFilterTracking, PhysicalExpr, PhysicalExprSimplifier, projection::ProjectionExprs, utils::reassign_expr_columns, @@ -32,12 +33,13 @@ use parquet::{ arrow::{ ParquetRecordBatchStreamBuilder, ProjectionMask, arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions, RowSelection}, + async_reader::AsyncFileReader, parquet_column, }, + errors::ParquetError, file::metadata::PageIndexPolicy, }; -use super::source::{CachedMetaReaderFactory, ParquetMetadataCacheReader}; use crate::{ cache::{ BatchID, ColumnLineages, InsertArrowArrayError, LiquidCacheParquetRef, ParquetFileIdentity, @@ -62,7 +64,8 @@ pub(crate) struct LiquidMorselizer { pub(crate) predicate: Option>, pub(crate) table_schema: TableSchema, pub(crate) metrics: ExecutionPlanMetricsSet, - pub(crate) parquet_file_reader_factory: Arc, + pub(crate) parquet_file_reader_factory: Arc, + pub(crate) object_store_url: ObjectStoreUrl, pub(crate) reorder_filters: bool, pub(crate) liquid_cache: LiquidCacheParquetRef, pub(crate) expr_adapter_factory: Arc, @@ -71,6 +74,32 @@ pub(crate) struct LiquidMorselizer { pub(crate) prefetch: bool, } +type ParquetInput = Box; + +#[derive(Clone)] +pub(crate) struct LiquidFileReaderFactory { + pub(crate) factory: Arc, + pub(crate) partition_index: usize, + pub(crate) partitioned_file: PartitionedFile, + pub(crate) metadata_size_hint: Option, + pub(crate) metrics: ExecutionPlanMetricsSet, +} + +impl LiquidFileReaderFactory { + pub(crate) fn create(&self) -> parquet::errors::Result { + let reader = self + .factory + .create_reader( + self.partition_index, + self.partitioned_file.clone(), + self.metadata_size_hint, + &self.metrics, + ) + .map_err(|error| ParquetError::External(Box::new(error)))?; + Ok(reader) + } +} + impl fmt::Debug for LiquidMorselizer { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("LiquidMorselizer") @@ -88,15 +117,16 @@ impl Morselizer for LiquidMorselizer { let metrics = LiquidFileMetrics::new(self.partition_index, &file_name, &self.metrics); let metadata_size_hint = partitioned_file.metadata_size_hint; let file_identity = ParquetFileIdentity::new( - self.parquet_file_reader_factory.object_store_url().clone(), + self.object_store_url.clone(), partitioned_file.object_meta.location.to_string(), ); - let reader = self.parquet_file_reader_factory.create_liquid_reader( - self.partition_index, - partitioned_file.clone(), + let reader_factory = Arc::new(LiquidFileReaderFactory { + factory: Arc::clone(&self.parquet_file_reader_factory), + partition_index: self.partition_index, + partitioned_file: partitioned_file.clone(), metadata_size_hint, - &self.metrics, - ); + metrics: self.metrics.clone(), + }); let logical_file_schema = Arc::clone(self.table_schema.file_schema()); let output_schema = Arc::new( @@ -151,7 +181,7 @@ impl Morselizer for LiquidMorselizer { file_name, metrics, file_pruner, - reader, + reader_factory, batch_size: self.batch_size, logical_file_schema, output_schema, @@ -201,7 +231,7 @@ struct PreparedLiquidOpen { file_name: String, metrics: LiquidFileMetrics, file_pruner: Option, - reader: ParquetMetadataCacheReader, + reader_factory: Arc, batch_size: usize, logical_file_schema: SchemaRef, output_schema: SchemaRef, @@ -232,7 +262,7 @@ struct RowGroupPlanningContext { reader_metadata: ArrowReaderMetadata, physical_file_schema: SchemaRef, cache_full_schema: SchemaRef, - builder: ParquetRecordBatchStreamBuilder, + builder: ParquetRecordBatchStreamBuilder, projection_mask: ProjectionMask, row_filter: Option, pruning_predicate: Option>, @@ -300,9 +330,9 @@ impl LiquidOpenState { let metadata_load_time = prepared.metrics.file_metrics.metadata_load_time.clone(); let mut timer = metadata_load_time.timer(); + let mut reader = prepared.reader_factory.create()?; let reader_metadata = - ArrowReaderMetadata::load_async(&mut prepared.reader, options.clone()) - .await?; + ArrowReaderMetadata::load_async(&mut reader, options.clone()).await?; timer.stop(); Ok(MetadataLoadedLiquidOpen { prepared, @@ -353,11 +383,6 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • 1, - "meta data must be cached already" - ); - let rewriter = loaded.prepared.expr_adapter_factory.create( Arc::clone(&loaded.prepared.logical_file_schema), Arc::clone(&physical_file_schema), @@ -381,7 +406,7 @@ fn prepare_and_prune_by_stats(mut loaded: MetadataLoadedLiquidOpen) -> Result
  • Result, + builder: &mut ParquetRecordBatchStreamBuilder, predicate: &PruningPredicate, file_metrics: &ParquetFileMetrics, row_groups: &RowGroupAccessPlanFilter, @@ -921,9 +946,12 @@ mod tests { common::ScalarValue, datasource::{ listing::PartitionedFile, - physical_plan::{FileScanConfigBuilder, FileSource, ParquetSource}, + physical_plan::{ + FileScanConfigBuilder, FileSource, ParquetSource, + parquet::{CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory}, + }, }, - execution::object_store::ObjectStoreUrl, + execution::{object_store::ObjectStoreUrl, runtime_env::RuntimeEnv}, logical_expr::Operator, physical_expr::{ PhysicalExpr, @@ -938,8 +966,8 @@ mod tests { cache::{AlwaysHydrate, Evict}, cache_policies::LiquidPolicy, }; - use object_store::local::LocalFileSystem; - use parquet::arrow::{ArrowWriter, async_reader::AsyncFileReader}; + use object_store::{ObjectStore, local::LocalFileSystem, path::Path}; + use parquet::arrow::ArrowWriter; use crate::{ cache::{BatchID, CachedFileRef, CachedRowGroupRef, LiquidCacheParquet}, @@ -1094,6 +1122,7 @@ mod tests { ) .await; let metrics = ExecutionPlanMetricsSet::new(); + let object_store_url = ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(); let morselizer = LiquidMorselizer { partition_index: 0, projection: ProjectionExprs::from_indices(&options.projection_columns, schema.as_ref()), @@ -1101,10 +1130,10 @@ mod tests { predicate: options.predicate, table_schema: TableSchema::from(Arc::clone(&schema)), metrics: metrics.clone(), - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new( object_store, - ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(), )), + object_store_url: object_store_url.clone(), reorder_filters: false, liquid_cache: cache.clone(), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), @@ -1113,10 +1142,7 @@ mod tests { prefetch: true, }; let cached_file = cache.register_or_get_file( - ParquetFileIdentity::new( - ObjectStoreUrl::parse(format!("test-{file_id}:///")).unwrap(), - file_name, - ), + ParquetFileIdentity::new(object_store_url, file_name), schema, ); TestFilePlanner { @@ -1156,6 +1182,59 @@ mod tests { } } + async fn metadata_cache_hits(cache_limit: usize) -> (Option, Option) { + let schema = schema(); + let tmp_dir = tempfile::tempdir().unwrap(); + let parquet_path = tmp_dir.path().join("data.parquet"); + write_two_row_group_file(&parquet_path, Arc::clone(&schema)); + let file = PartitionedFile::new( + "data.parquet", + std::fs::metadata(&parquet_path).unwrap().len(), + ); + let object_store: Arc = + Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path()).unwrap()); + let metadata_cache = RuntimeEnv::default() + .cache_manager + .get_file_metadata_cache(); + metadata_cache.update_cache_limit(cache_limit); + let parquet_file_reader_factory = Arc::new(CachedParquetFileReaderFactory::new( + Arc::clone(&object_store), + Arc::clone(&metadata_cache), + )); + let parquet_source = ParquetSource::new(Arc::clone(&schema)) + .with_parquet_file_reader_factory(parquet_file_reader_factory); + let liquid_cache = create_test_cache(tmp_dir.path(), usize::MAX, usize::MAX).await; + let source = LiquidParquetSource::from_parquet_source(parquet_source, liquid_cache); + let base_config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(source.clone()), + ) + .with_file(file.clone()) + .build(); + let morselizer = source + .create_morselizer(object_store, &base_config, 0) + .unwrap(); + let path = Path::from("data.parquet"); + + advance_to_row_group_chain(morselizer.plan_file(file.clone()).unwrap()).await; + let first_hits = metadata_cache + .list_entries() + .get(&path) + .map(|entry| entry.hits); + advance_to_row_group_chain(morselizer.plan_file(file).unwrap()).await; + let second_hits = metadata_cache + .list_entries() + .get(&path) + .map(|entry| entry.hits); + (first_hits, second_hits) + } + + #[tokio::test] + async fn datafusion_metadata_cache_ablation() { + assert_eq!(metadata_cache_hits(usize::MAX).await, (Some(0), Some(1))); + assert_eq!(metadata_cache_hits(0).await, (None, None)); + } + fn gt_expr(column_name: &str, column_index: usize, literal: i32) -> Arc { Arc::new(BinaryExpr::new( Arc::new(Column::new(column_name, column_index)), @@ -1172,44 +1251,6 @@ mod tests { )) } - #[tokio::test] - async fn metadata_cache_is_scoped_to_object_store() { - let schema = schema(); - let dir_a = tempfile::tempdir().unwrap(); - let dir_b = tempfile::tempdir().unwrap(); - let path_a = dir_a.path().join("data.parquet"); - let path_b = dir_b.path().join("data.parquet"); - write_single_row_group_file(&path_a, schema.clone(), vec![1]); - write_single_row_group_file(&path_b, schema, vec![1, 2]); - let metrics = ExecutionPlanMetricsSet::new(); - let mut reader_a = CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_a.path()).unwrap()), - ObjectStoreUrl::parse("store-a:///").unwrap(), - ) - .create_liquid_reader( - 0, - PartitionedFile::new("data.parquet", std::fs::metadata(path_a).unwrap().len()), - None, - &metrics, - ); - let mut reader_b = CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_b.path()).unwrap()), - ObjectStoreUrl::parse("store-b:///").unwrap(), - ) - .create_liquid_reader( - 0, - PartitionedFile::new("data.parquet", std::fs::metadata(path_b).unwrap().len()), - None, - &metrics, - ); - - let metadata_a = reader_a.get_metadata(None).await.unwrap(); - let metadata_b = reader_b.get_metadata(None).await.unwrap(); - - assert_eq!(metadata_a.file_metadata().num_rows(), 1); - assert_eq!(metadata_b.file_metadata().num_rows(), 2); - } - #[tokio::test] async fn data_cache_is_scoped_to_object_store() { let schema = schema(); @@ -1231,10 +1272,10 @@ mod tests { predicate: None, table_schema: TableSchema::from(Arc::clone(&schema)), metrics: metrics.clone(), - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_a.path()).unwrap()), - ObjectStoreUrl::parse("data-cache-a:///").unwrap(), - )), + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::new( + LocalFileSystem::new_with_prefix(dir_a.path()).unwrap(), + ))), + object_store_url: ObjectStoreUrl::parse("data-cache-a:///").unwrap(), reorder_filters: false, liquid_cache: Arc::clone(&cache), expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), @@ -1249,10 +1290,10 @@ mod tests { predicate: None, table_schema: TableSchema::from(Arc::clone(&schema)), metrics, - parquet_file_reader_factory: Arc::new(CachedMetaReaderFactory::new( - Arc::new(LocalFileSystem::new_with_prefix(dir_b.path()).unwrap()), - ObjectStoreUrl::parse("data-cache-b:///").unwrap(), - )), + parquet_file_reader_factory: Arc::new(DefaultParquetFileReaderFactory::new(Arc::new( + LocalFileSystem::new_with_prefix(dir_b.path()).unwrap(), + ))), + object_store_url: ObjectStoreUrl::parse("data-cache-b:///").unwrap(), reorder_filters: false, liquid_cache: cache, expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), diff --git a/src/datafusion/src/reader/plantime/source.rs b/src/datafusion/src/reader/plantime/source.rs index 5f0c764e..69e0415b 100644 --- a/src/datafusion/src/reader/plantime/source.rs +++ b/src/datafusion/src/reader/plantime/source.rs @@ -1,20 +1,16 @@ use super::LiquidMorselizer; use crate::cache::{ColumnLineages, LiquidCacheParquetRef}; -use ahash::{HashMap, HashMapExt}; -use bytes::Bytes; use datafusion::{ common::{internal_err, tree_node::TreeNodeRecursion}, config::{ConfigOptions, TableParquetOptions}, datasource::{ - listing::PartitionedFile, physical_plan::{ - FileScanConfig, FileSource, ParquetFileMetrics, ParquetFileReaderFactory, - ParquetSource, parquet::can_expr_be_pushed_down_with_schemas, + FileScanConfig, FileSource, ParquetFileReaderFactory, ParquetSource, + parquet::{DefaultParquetFileReaderFactory, can_expr_be_pushed_down_with_schemas}, }, table_schema::TableSchema, }, error::Result, - execution::object_store::ObjectStoreUrl, physical_expr::projection::ProjectionExprs, physical_expr::utils::conjunction, physical_expr_adapter::DefaultPhysicalExprAdapterFactory, @@ -25,168 +21,11 @@ use datafusion::{ }, }; use datafusion_datasource::morsel::Morselizer; -use futures::{FutureExt, future::BoxFuture}; -use object_store::{ObjectStore, ObjectStoreExt, path::Path}; -use parquet::{ - arrow::{arrow_reader::ArrowReaderOptions, async_reader::AsyncFileReader}, - errors::ParquetError, - file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader}, -}; +use object_store::ObjectStore; use std::{ fmt::{self, Formatter}, - ops::Range, - sync::{Arc, LazyLock}, + sync::Arc, }; -use tokio::sync::RwLock; - -static META_CACHE: LazyLock = LazyLock::new(MetadataCache::new); - -#[derive(Debug)] -pub(crate) struct CachedMetaReaderFactory { - store: Arc, - store_url: ObjectStoreUrl, -} - -impl CachedMetaReaderFactory { - pub(crate) fn new(store: Arc, store_url: ObjectStoreUrl) -> Self { - Self { store, store_url } - } - - pub(crate) fn object_store_url(&self) -> &ObjectStoreUrl { - &self.store_url - } - - pub(crate) fn create_liquid_reader( - &self, - partition_index: usize, - partitioned_file: PartitionedFile, - metadata_size_hint: Option, - metrics: &ExecutionPlanMetricsSet, - ) -> ParquetMetadataCacheReader { - let path = partitioned_file.object_meta.location.clone(); - - ParquetMetadataCacheReader { - file_metrics: ParquetFileMetrics::new(partition_index, path.as_ref(), metrics), - store: Arc::clone(&self.store), - store_url: self.store_url.clone(), - file_size: partitioned_file.object_meta.size, - metadata_size_hint, - path, - } - } -} - -impl ParquetFileReaderFactory for CachedMetaReaderFactory { - fn create_reader( - &self, - partition_index: usize, - partitioned_file: PartitionedFile, - metadata_size_hint: Option, - metrics: &ExecutionPlanMetricsSet, - ) -> Result> { - let reader = self.create_liquid_reader( - partition_index, - partitioned_file, - metadata_size_hint, - metrics, - ); - Ok(Box::new(reader)) - } -} - -struct MetadataCache { - val: RwLock>>, -} - -impl MetadataCache { - fn new() -> Self { - Self { - val: RwLock::new(HashMap::new()), - } - } -} - -#[derive(Clone)] -pub struct ParquetMetadataCacheReader { - file_metrics: ParquetFileMetrics, - store: Arc, - store_url: ObjectStoreUrl, - file_size: u64, - metadata_size_hint: Option, - path: Path, -} - -fn to_parquet_err(error: object_store::Error) -> ParquetError { - ParquetError::External(Box::new(error)) -} - -impl AsyncFileReader for ParquetMetadataCacheReader { - fn get_byte_ranges( - &mut self, - ranges: Vec>, - ) -> BoxFuture<'_, parquet::errors::Result>> { - let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); - self.file_metrics.bytes_scanned.add(total as usize); - async move { - self.store - .get_ranges(&self.path, &ranges) - .await - .map_err(to_parquet_err) - } - .boxed() - } - - fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { - self.file_metrics - .bytes_scanned - .add((range.end - range.start) as usize); - async move { - self.store - .get_range(&self.path, range) - .await - .map_err(to_parquet_err) - } - .boxed() - } - - fn get_metadata( - &mut self, - options: Option<&ArrowReaderOptions>, - ) -> BoxFuture<'_, parquet::errors::Result>> { - let cache_key = (self.store_url.clone(), self.path.clone()); - let options = options.cloned(); - async move { - // First check with read lock - { - let cache = META_CACHE.val.read().await; - if let Some(meta) = cache.get(&cache_key) { - return Ok(meta.clone()); - } - } - - // Upgrade to write lock and double-check - let mut cache = META_CACHE.val.write().await; - match cache.entry(cache_key) { - std::collections::hash_map::Entry::Occupied(entry) => Ok(entry.get().clone()), - std::collections::hash_map::Entry::Vacant(entry) => { - let file_size = self.file_size; - let meta = ParquetMetaDataReader::new() - .with_arrow_reader_options(options.as_ref()) - .with_prefetch_hint(self.metadata_size_hint) - .load_and_finish(&mut *self, file_size) - .await?; - let mut reader = ParquetMetaDataReader::new_with_metadata(meta.clone()) - .with_page_index_policy(PageIndexPolicy::Optional); - reader.load_page_index(&mut *self).await?; - let meta = Arc::new(reader.finish()?); - entry.insert(meta.clone()); - Ok(meta) - } - } - } - .boxed() - } -} /// The data source for LiquidCache #[derive(Clone)] @@ -198,6 +37,7 @@ pub struct LiquidParquetSource { batch_size: Option, projection: ProjectionExprs, table_schema: TableSchema, + parquet_file_reader_factory: Option>, span: Option>, lineages: Arc, prefetch: bool, @@ -253,6 +93,7 @@ impl LiquidParquetSource { /// Create a new LiquidParquetSource from a ParquetSource pub fn from_parquet_source(source: ParquetSource, liquid_cache: LiquidCacheParquetRef) -> Self { let predicate = source.filter(); + let parquet_file_reader_factory = source.parquet_file_reader_factory().cloned(); let table_schema = source.table_schema().clone(); let projection = source.projection().cloned().unwrap_or_else(|| { @@ -269,6 +110,7 @@ impl LiquidParquetSource { liquid_cache, projection, metrics: source.metrics().clone(), + parquet_file_reader_factory, predicate: None, span: None, lineages: Arc::default(), @@ -311,10 +153,10 @@ impl FileSource for LiquidParquetSource { .clone() .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); - let reader_factory = Arc::new(CachedMetaReaderFactory::new( - object_store, - base_config.object_store_url.clone(), - )); + let parquet_file_reader_factory = self + .parquet_file_reader_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultParquetFileReaderFactory::new(object_store))); let execution_span = self .span @@ -330,7 +172,8 @@ impl FileSource for LiquidParquetSource { table_schema: self.table_schema.clone(), metrics: self.metrics.clone(), liquid_cache: self.liquid_cache.clone(), - parquet_file_reader_factory: reader_factory, + parquet_file_reader_factory, + object_store_url: base_config.object_store_url.clone(), reorder_filters: self.reorder_filters(), expr_adapter_factory, span: execution_span.map(Arc::new), diff --git a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs index 37babe4a..84a03b55 100644 --- a/src/datafusion/src/reader/runtime/liquid_cache_reader.rs +++ b/src/datafusion/src/reader/runtime/liquid_cache_reader.rs @@ -17,7 +17,7 @@ use parquet::errors::ParquetError; use parquet::file::metadata::ParquetMetaData; use crate::cache::{BatchID, CachedRowGroupRef, InsertArrowArrayError}; -use crate::reader::plantime::{LiquidRowFilter, ParquetMetadataCacheReader}; +use crate::reader::plantime::{LiquidFileReaderFactory, LiquidRowFilter}; use crate::reader::runtime::utils::take_next_batch; use crate::utils::{boolean_buffer_and_then, row_selector_to_boolean_buffer}; @@ -71,7 +71,7 @@ pub(crate) struct LiquidCacheReaderConfig { pub(crate) struct ParquetFallbackConfig { pub(crate) row_group_idx: usize, pub(crate) metadata: Arc, - pub(crate) input: ParquetMetadataCacheReader, + pub(crate) reader_factory: Arc, pub(crate) cache_projection: ProjectionMask, pub(crate) cache_column_ids: Vec, pub(crate) cache_batch_size: usize, @@ -81,7 +81,7 @@ pub(crate) struct ParquetFallbackConfig { pub(crate) struct ParquetFallback { row_group_idx: usize, metadata: Arc, - input: ParquetMetadataCacheReader, + reader_factory: Arc, cache_projection: ProjectionMask, cache_column_ids: Vec, cache_batch_size: usize, @@ -162,7 +162,7 @@ impl ParquetFallback { Self { row_group_idx: config.row_group_idx, metadata: config.metadata, - input: config.input, + reader_factory: config.reader_factory, cache_projection: config.cache_projection, cache_column_ids: config.cache_column_ids, cache_batch_size: config.cache_batch_size, @@ -199,14 +199,16 @@ impl ParquetFallback { let row_selection = build_row_selection_from(batch_id, self.cache_batch_size, self.row_count); - let stream = - ParquetRecordBatchStreamBuilder::new_with_metadata(self.input.clone(), reader_metadata) - .with_projection(self.cache_projection.clone()) - .with_row_groups(vec![self.row_group_idx]) - .with_batch_size(self.cache_batch_size) - .with_row_selection(row_selection) - .build()? - .boxed(); + let stream = ParquetRecordBatchStreamBuilder::new_with_metadata( + self.reader_factory.create()?, + reader_metadata, + ) + .with_projection(self.cache_projection.clone()) + .with_row_groups(vec![self.row_group_idx]) + .with_batch_size(self.cache_batch_size) + .with_row_selection(row_selection) + .build()? + .boxed(); self.stream = Some(stream); self.next_batch_id = batch_id; @@ -496,13 +498,15 @@ mod tests { use super::*; use crate::{ cache::LiquidCacheParquet, - reader::plantime::CachedMetaReaderFactory, + reader::plantime::LiquidFileReaderFactory, reader::{FilterCandidateBuilder, LiquidPredicate, LiquidRowFilter}, }; use arrow::array::{ArrayRef, Int32Array}; use arrow::record_batch::RecordBatch; use arrow_schema::{DataType, Field, Schema, SchemaRef}; - use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::{ + listing::PartitionedFile, physical_plan::parquet::DefaultParquetFileReaderFactory, + }; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::{ logical_expr::Operator, @@ -573,11 +577,13 @@ mod tests { std::fs::metadata(&parquet_path).unwrap().len(), ); let metrics = ExecutionPlanMetricsSet::new(); - let input = CachedMetaReaderFactory::new( - object_store, - datafusion::execution::object_store::ObjectStoreUrl::parse("test-runtime:///").unwrap(), - ) - .create_liquid_reader(0, partitioned_file, None, &metrics); + let reader_factory = Arc::new(LiquidFileReaderFactory { + factory: Arc::new(DefaultParquetFileReaderFactory::new(object_store)), + partition_index: 0, + partitioned_file, + metadata_size_hint: None, + metrics, + }); let projection = ProjectionMask::roots( reader_metadata.metadata().file_metadata().schema_descr(), [0], @@ -622,7 +628,7 @@ mod tests { fallback: ParquetFallbackConfig { row_group_idx: 0, metadata: Arc::clone(reader_metadata.metadata()), - input, + reader_factory, cache_projection: projection, cache_column_ids: vec![0], cache_batch_size: batch_size, diff --git a/src/datafusion/src/reader/runtime/morsel.rs b/src/datafusion/src/reader/runtime/morsel.rs index 1c6931ff..271f53ff 100644 --- a/src/datafusion/src/reader/runtime/morsel.rs +++ b/src/datafusion/src/reader/runtime/morsel.rs @@ -15,7 +15,7 @@ use parquet::{ use crate::{ cache::{CachedFileRef, CachedRowGroupRef, LiquidCacheParquetRef, RowGroupSnapshots}, - reader::plantime::{LiquidFileMetrics, LiquidRowFilter, ParquetMetadataCacheReader}, + reader::plantime::{LiquidFileMetrics, LiquidFileReaderFactory, LiquidRowFilter}, }; use super::{ @@ -27,7 +27,7 @@ use super::{ pub(crate) struct LiquidRowGroupPlanner { pub(crate) metadata: Arc, - pub(crate) input: ParquetMetadataCacheReader, + pub(crate) reader_factory: Arc, pub(crate) row_filter: Option, pub(crate) cached_file: CachedFileRef, pub(crate) projection: ProjectionMask, @@ -79,7 +79,7 @@ impl LiquidRowGroupPlanner { ParquetFallbackConfig { row_group_idx, metadata: Arc::clone(&self.metadata), - input: self.input.clone(), + reader_factory: Arc::clone(&self.reader_factory), cache_projection: details.cache_projection.clone(), cache_column_ids: details.cache_column_ids.clone(), cache_batch_size,