You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iceberg.apache.org by li...@apache.org on 2025/03/25 04:53:14 UTC
(iceberg-rust) branch main updated: refactor: Split scan module (#1120)
This is an automated email from the ASF dual-hosted git repository.
liurenjie1024 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-rust.git
The following commit(s) were added to refs/heads/main by this push:
new b84e0d28 refactor: Split scan module (#1120)
b84e0d28 is described below
commit b84e0d28697896309ec4131c40e8d192ff00bfd4
Author: Jonathan Chen <ch...@gmail.com>
AuthorDate: Tue Mar 25 00:53:09 2025 -0400
refactor: Split scan module (#1120)
## Which issue does this PR close?
- Closes #.
## What changes are included in this PR?
Split Scan module, i was going through it and all the information for
scan was on one page (too long)
Co-authored-by: Renjie Liu <li...@gmail.com>
---
crates/iceberg/src/scan/cache.rs | 237 ++++++++++++
crates/iceberg/src/scan/context.rs | 275 ++++++++++++++
crates/iceberg/src/{scan.rs => scan/mod.rs} | 564 +---------------------------
crates/iceberg/src/scan/task.rs | 116 ++++++
4 files changed, 639 insertions(+), 553 deletions(-)
diff --git a/crates/iceberg/src/scan/cache.rs b/crates/iceberg/src/scan/cache.rs
new file mode 100644
index 00000000..6f249126
--- /dev/null
+++ b/crates/iceberg/src/scan/cache.rs
@@ -0,0 +1,237 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+use crate::expr::visitors::expression_evaluator::ExpressionEvaluator;
+use crate::expr::visitors::inclusive_projection::InclusiveProjection;
+use crate::expr::visitors::manifest_evaluator::ManifestEvaluator;
+use crate::expr::{Bind, BoundPredicate};
+use crate::spec::{Schema, TableMetadataRef};
+use crate::{Error, ErrorKind, Result};
+
+/// Manages the caching of [`BoundPredicate`] objects
+/// for [`PartitionSpec`]s based on partition spec id.
+#[derive(Debug)]
+pub(crate) struct PartitionFilterCache(RwLock<HashMap<i32, Arc<BoundPredicate>>>);
+
+impl PartitionFilterCache {
+ /// Creates a new [`PartitionFilterCache`]
+ /// with an empty internal HashMap.
+ pub(crate) fn new() -> Self {
+ Self(RwLock::new(HashMap::new()))
+ }
+
+ /// Retrieves a [`BoundPredicate`] from the cache
+ /// or computes it if not present.
+ pub(crate) fn get(
+ &self,
+ spec_id: i32,
+ table_metadata: &TableMetadataRef,
+ schema: &Schema,
+ case_sensitive: bool,
+ filter: BoundPredicate,
+ ) -> Result<Arc<BoundPredicate>> {
+ // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
+ // below, otherwise we hit deadlock
+ {
+ let read = self.0.read().map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "PartitionFilterCache RwLock was poisoned",
+ )
+ })?;
+
+ if read.contains_key(&spec_id) {
+ return Ok(read.get(&spec_id).unwrap().clone());
+ }
+ }
+
+ let partition_spec = table_metadata
+ .partition_spec_by_id(spec_id)
+ .ok_or(Error::new(
+ ErrorKind::Unexpected,
+ format!("Could not find partition spec for id {}", spec_id),
+ ))?;
+
+ let partition_type = partition_spec.partition_type(schema)?;
+ let partition_fields = partition_type.fields().to_owned();
+ let partition_schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(partition_spec.spec_id())
+ .with_fields(partition_fields)
+ .build()?,
+ );
+
+ let mut inclusive_projection = InclusiveProjection::new(partition_spec.clone());
+
+ let partition_filter = inclusive_projection
+ .project(&filter)?
+ .rewrite_not()
+ .bind(partition_schema.clone(), case_sensitive)?;
+
+ self.0
+ .write()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "PartitionFilterCache RwLock was poisoned",
+ )
+ })?
+ .insert(spec_id, Arc::new(partition_filter));
+
+ let read = self.0.read().map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "PartitionFilterCache RwLock was poisoned",
+ )
+ })?;
+
+ Ok(read.get(&spec_id).unwrap().clone())
+ }
+}
+
+/// Manages the caching of [`ManifestEvaluator`] objects
+/// for [`PartitionSpec`]s based on partition spec id.
+#[derive(Debug)]
+pub(crate) struct ManifestEvaluatorCache(RwLock<HashMap<i32, Arc<ManifestEvaluator>>>);
+
+impl ManifestEvaluatorCache {
+ /// Creates a new [`ManifestEvaluatorCache`]
+ /// with an empty internal HashMap.
+ pub(crate) fn new() -> Self {
+ Self(RwLock::new(HashMap::new()))
+ }
+
+ /// Retrieves a [`ManifestEvaluator`] from the cache
+ /// or computes it if not present.
+ pub(crate) fn get(
+ &self,
+ spec_id: i32,
+ partition_filter: Arc<BoundPredicate>,
+ ) -> Arc<ManifestEvaluator> {
+ // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
+ // below, otherwise we hit deadlock
+ {
+ let read = self
+ .0
+ .read()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "ManifestEvaluatorCache RwLock was poisoned",
+ )
+ })
+ .unwrap();
+
+ if read.contains_key(&spec_id) {
+ return read.get(&spec_id).unwrap().clone();
+ }
+ }
+
+ self.0
+ .write()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "ManifestEvaluatorCache RwLock was poisoned",
+ )
+ })
+ .unwrap()
+ .insert(
+ spec_id,
+ Arc::new(ManifestEvaluator::new(partition_filter.as_ref().clone())),
+ );
+
+ let read = self
+ .0
+ .read()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "ManifestEvaluatorCache RwLock was poisoned",
+ )
+ })
+ .unwrap();
+
+ read.get(&spec_id).unwrap().clone()
+ }
+}
+
+/// Manages the caching of [`ExpressionEvaluator`] objects
+/// for [`PartitionSpec`]s based on partition spec id.
+#[derive(Debug)]
+pub(crate) struct ExpressionEvaluatorCache(RwLock<HashMap<i32, Arc<ExpressionEvaluator>>>);
+
+impl ExpressionEvaluatorCache {
+ /// Creates a new [`ExpressionEvaluatorCache`]
+ /// with an empty internal HashMap.
+ pub(crate) fn new() -> Self {
+ Self(RwLock::new(HashMap::new()))
+ }
+
+ /// Retrieves a [`ExpressionEvaluator`] from the cache
+ /// or computes it if not present.
+ pub(crate) fn get(
+ &self,
+ spec_id: i32,
+ partition_filter: &BoundPredicate,
+ ) -> Result<Arc<ExpressionEvaluator>> {
+ // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
+ // below, otherwise we hit deadlock
+ {
+ let read = self.0.read().map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "PartitionFilterCache RwLock was poisoned",
+ )
+ })?;
+
+ if read.contains_key(&spec_id) {
+ return Ok(read.get(&spec_id).unwrap().clone());
+ }
+ }
+
+ self.0
+ .write()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "ManifestEvaluatorCache RwLock was poisoned",
+ )
+ })
+ .unwrap()
+ .insert(
+ spec_id,
+ Arc::new(ExpressionEvaluator::new(partition_filter.clone())),
+ );
+
+ let read = self
+ .0
+ .read()
+ .map_err(|_| {
+ Error::new(
+ ErrorKind::Unexpected,
+ "ManifestEvaluatorCache RwLock was poisoned",
+ )
+ })
+ .unwrap();
+
+ Ok(read.get(&spec_id).unwrap().clone())
+ }
+}
diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs
new file mode 100644
index 00000000..6bfb12b2
--- /dev/null
+++ b/crates/iceberg/src/scan/context.rs
@@ -0,0 +1,275 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use futures::channel::mpsc::Sender;
+use futures::{SinkExt, TryFutureExt};
+
+use crate::delete_file_index::DeleteFileIndex;
+use crate::expr::{Bind, BoundPredicate, Predicate};
+use crate::io::object_cache::ObjectCache;
+use crate::scan::{
+ BoundPredicates, ExpressionEvaluatorCache, FileScanTask, ManifestEvaluatorCache,
+ PartitionFilterCache,
+};
+use crate::spec::{
+ ManifestContentType, ManifestEntryRef, ManifestFile, ManifestList, SchemaRef, SnapshotRef,
+ TableMetadataRef,
+};
+use crate::{Error, ErrorKind, Result};
+
+/// Wraps a [`ManifestFile`] alongside the objects that are needed
+/// to process it in a thread-safe manner
+pub(crate) struct ManifestFileContext {
+ manifest_file: ManifestFile,
+
+ sender: Sender<ManifestEntryContext>,
+
+ field_ids: Arc<Vec<i32>>,
+ bound_predicates: Option<Arc<BoundPredicates>>,
+ object_cache: Arc<ObjectCache>,
+ snapshot_schema: SchemaRef,
+ expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
+ delete_file_index: Option<DeleteFileIndex>,
+}
+
+/// Wraps a [`ManifestEntryRef`] alongside the objects that are needed
+/// to process it in a thread-safe manner
+pub(crate) struct ManifestEntryContext {
+ pub manifest_entry: ManifestEntryRef,
+
+ pub expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
+ pub field_ids: Arc<Vec<i32>>,
+ pub bound_predicates: Option<Arc<BoundPredicates>>,
+ pub partition_spec_id: i32,
+ pub snapshot_schema: SchemaRef,
+ pub delete_file_index: Option<DeleteFileIndex>,
+}
+
+impl ManifestFileContext {
+ /// Consumes this [`ManifestFileContext`], fetching its Manifest from FileIO and then
+ /// streaming its constituent [`ManifestEntries`] to the channel provided in the context
+ pub(crate) async fn fetch_manifest_and_stream_manifest_entries(self) -> Result<()> {
+ let ManifestFileContext {
+ object_cache,
+ manifest_file,
+ bound_predicates,
+ snapshot_schema,
+ field_ids,
+ mut sender,
+ expression_evaluator_cache,
+ delete_file_index,
+ ..
+ } = self;
+
+ let manifest = object_cache.get_manifest(&manifest_file).await?;
+
+ for manifest_entry in manifest.entries() {
+ let manifest_entry_context = ManifestEntryContext {
+ // TODO: refactor to avoid the expensive ManifestEntry clone
+ manifest_entry: manifest_entry.clone(),
+ expression_evaluator_cache: expression_evaluator_cache.clone(),
+ field_ids: field_ids.clone(),
+ partition_spec_id: manifest_file.partition_spec_id,
+ bound_predicates: bound_predicates.clone(),
+ snapshot_schema: snapshot_schema.clone(),
+ delete_file_index: delete_file_index.clone(),
+ };
+
+ sender
+ .send(manifest_entry_context)
+ .map_err(|_| Error::new(ErrorKind::Unexpected, "mpsc channel SendError"))
+ .await?;
+ }
+
+ Ok(())
+ }
+}
+
+impl ManifestEntryContext {
+ /// consume this `ManifestEntryContext`, returning a `FileScanTask`
+ /// created from it
+ pub(crate) async fn into_file_scan_task(self) -> Result<FileScanTask> {
+ let deletes = if let Some(delete_file_index) = self.delete_file_index {
+ delete_file_index
+ .get_deletes_for_data_file(
+ self.manifest_entry.data_file(),
+ self.manifest_entry.sequence_number(),
+ )
+ .await?
+ } else {
+ vec![]
+ };
+
+ Ok(FileScanTask {
+ start: 0,
+ length: self.manifest_entry.file_size_in_bytes(),
+ record_count: Some(self.manifest_entry.record_count()),
+
+ data_file_path: self.manifest_entry.file_path().to_string(),
+ data_file_content: self.manifest_entry.content_type(),
+ data_file_format: self.manifest_entry.file_format(),
+
+ schema: self.snapshot_schema,
+ project_field_ids: self.field_ids.to_vec(),
+ predicate: self
+ .bound_predicates
+ .map(|x| x.as_ref().snapshot_bound_predicate.clone()),
+
+ deletes,
+ })
+ }
+}
+
+/// PlanContext wraps a [`SnapshotRef`] alongside all the other
+/// objects that are required to perform a scan file plan.
+#[derive(Debug)]
+pub(crate) struct PlanContext {
+ pub snapshot: SnapshotRef,
+
+ pub table_metadata: TableMetadataRef,
+ pub snapshot_schema: SchemaRef,
+ pub case_sensitive: bool,
+ pub predicate: Option<Arc<Predicate>>,
+ pub snapshot_bound_predicate: Option<Arc<BoundPredicate>>,
+ pub object_cache: Arc<ObjectCache>,
+ pub field_ids: Arc<Vec<i32>>,
+
+ pub partition_filter_cache: Arc<PartitionFilterCache>,
+ pub manifest_evaluator_cache: Arc<ManifestEvaluatorCache>,
+ pub expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
+}
+
+impl PlanContext {
+ pub(crate) async fn get_manifest_list(&self) -> Result<Arc<ManifestList>> {
+ self.object_cache
+ .as_ref()
+ .get_manifest_list(&self.snapshot, &self.table_metadata)
+ .await
+ }
+
+ fn get_partition_filter(&self, manifest_file: &ManifestFile) -> Result<Arc<BoundPredicate>> {
+ let partition_spec_id = manifest_file.partition_spec_id;
+
+ let partition_filter = self.partition_filter_cache.get(
+ partition_spec_id,
+ &self.table_metadata,
+ &self.snapshot_schema,
+ self.case_sensitive,
+ self.predicate
+ .as_ref()
+ .ok_or(Error::new(
+ ErrorKind::Unexpected,
+ "Expected a predicate but none present",
+ ))?
+ .as_ref()
+ .bind(self.snapshot_schema.clone(), self.case_sensitive)?,
+ )?;
+
+ Ok(partition_filter)
+ }
+
+ pub(crate) fn build_manifest_file_contexts(
+ &self,
+ manifest_list: Arc<ManifestList>,
+ tx_data: Sender<ManifestEntryContext>,
+ delete_file_idx_and_tx: Option<(DeleteFileIndex, Sender<ManifestEntryContext>)>,
+ ) -> Result<Box<impl Iterator<Item = Result<ManifestFileContext>> + 'static>> {
+ let manifest_files = manifest_list.entries().iter();
+
+ // TODO: Ideally we could ditch this intermediate Vec as we return an iterator.
+ let mut filtered_mfcs = vec![];
+
+ for manifest_file in manifest_files {
+ let (delete_file_idx, tx) = if manifest_file.content == ManifestContentType::Deletes {
+ let Some((delete_file_idx, tx)) = delete_file_idx_and_tx.as_ref() else {
+ continue;
+ };
+ (Some(delete_file_idx.clone()), tx.clone())
+ } else {
+ (
+ delete_file_idx_and_tx.as_ref().map(|x| x.0.clone()),
+ tx_data.clone(),
+ )
+ };
+
+ let partition_bound_predicate = if self.predicate.is_some() {
+ let partition_bound_predicate = self.get_partition_filter(manifest_file)?;
+
+ // evaluate the ManifestFile against the partition filter. Skip
+ // if it cannot contain any matching rows
+ if !self
+ .manifest_evaluator_cache
+ .get(
+ manifest_file.partition_spec_id,
+ partition_bound_predicate.clone(),
+ )
+ .eval(manifest_file)?
+ {
+ continue;
+ }
+
+ Some(partition_bound_predicate)
+ } else {
+ None
+ };
+
+ let mfc = self.create_manifest_file_context(
+ manifest_file,
+ partition_bound_predicate,
+ tx,
+ delete_file_idx,
+ );
+
+ filtered_mfcs.push(Ok(mfc));
+ }
+
+ Ok(Box::new(filtered_mfcs.into_iter()))
+ }
+
+ fn create_manifest_file_context(
+ &self,
+ manifest_file: &ManifestFile,
+ partition_filter: Option<Arc<BoundPredicate>>,
+ sender: Sender<ManifestEntryContext>,
+ delete_file_index: Option<DeleteFileIndex>,
+ ) -> ManifestFileContext {
+ let bound_predicates =
+ if let (Some(ref partition_bound_predicate), Some(snapshot_bound_predicate)) =
+ (partition_filter, &self.snapshot_bound_predicate)
+ {
+ Some(Arc::new(BoundPredicates {
+ partition_bound_predicate: partition_bound_predicate.as_ref().clone(),
+ snapshot_bound_predicate: snapshot_bound_predicate.as_ref().clone(),
+ }))
+ } else {
+ None
+ };
+
+ ManifestFileContext {
+ manifest_file: manifest_file.clone(),
+ bound_predicates,
+ sender,
+ object_cache: self.object_cache.clone(),
+ snapshot_schema: self.snapshot_schema.clone(),
+ field_ids: self.field_ids.clone(),
+ expression_evaluator_cache: self.expression_evaluator_cache.clone(),
+ delete_file_index,
+ }
+ }
+}
diff --git a/crates/iceberg/src/scan.rs b/crates/iceberg/src/scan/mod.rs
similarity index 79%
rename from crates/iceberg/src/scan.rs
rename to crates/iceberg/src/scan/mod.rs
index 603e1aa8..7cddcfeb 100644
--- a/crates/iceberg/src/scan.rs
+++ b/crates/iceberg/src/scan/mod.rs
@@ -17,35 +17,31 @@
//! Table scan api.
-use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
+mod cache;
+use cache::*;
+mod context;
+use context::*;
+mod task;
+
+use std::sync::Arc;
use arrow_array::RecordBatch;
use futures::channel::mpsc::{channel, Sender};
use futures::stream::BoxStream;
-use futures::{SinkExt, StreamExt, TryFutureExt, TryStreamExt};
-use serde::{Deserialize, Serialize};
+use futures::{SinkExt, StreamExt, TryStreamExt};
+pub use task::*;
use crate::arrow::ArrowReaderBuilder;
use crate::delete_file_index::DeleteFileIndex;
-use crate::expr::visitors::expression_evaluator::ExpressionEvaluator;
use crate::expr::visitors::inclusive_metrics_evaluator::InclusiveMetricsEvaluator;
-use crate::expr::visitors::inclusive_projection::InclusiveProjection;
-use crate::expr::visitors::manifest_evaluator::ManifestEvaluator;
use crate::expr::{Bind, BoundPredicate, Predicate};
-use crate::io::object_cache::ObjectCache;
use crate::io::FileIO;
use crate::runtime::spawn;
-use crate::spec::{
- DataContentType, DataFileFormat, ManifestContentType, ManifestEntryRef, ManifestFile,
- ManifestList, Schema, SchemaRef, SnapshotRef, TableMetadataRef,
-};
+use crate::spec::{DataContentType, SnapshotRef};
use crate::table::Table;
use crate::utils::available_parallelism;
use crate::{Error, ErrorKind, Result};
-/// A stream of [`FileScanTask`].
-pub type FileScanTaskStream = BoxStream<'static, Result<FileScanTask>>;
/// A stream of arrow [`RecordBatch`]es.
pub type ArrowRecordBatchStream = BoxStream<'static, Result<RecordBatch>>;
@@ -341,25 +337,6 @@ pub struct TableScan {
delete_file_processing_enabled: bool,
}
-/// PlanContext wraps a [`SnapshotRef`] alongside all the other
-/// objects that are required to perform a scan file plan.
-#[derive(Debug)]
-struct PlanContext {
- snapshot: SnapshotRef,
-
- table_metadata: TableMetadataRef,
- snapshot_schema: SchemaRef,
- case_sensitive: bool,
- predicate: Option<Arc<Predicate>>,
- snapshot_bound_predicate: Option<Arc<BoundPredicate>>,
- object_cache: Arc<ObjectCache>,
- field_ids: Arc<Vec<i32>>,
-
- partition_filter_cache: Arc<PartitionFilterCache>,
- manifest_evaluator_cache: Arc<ManifestEvaluatorCache>,
- expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
-}
-
impl TableScan {
/// Returns a stream of [`FileScanTask`]s.
pub async fn plan_files(&self) -> Result<FileScanTaskStream> {
@@ -591,530 +568,11 @@ impl TableScan {
}
}
-struct BoundPredicates {
+pub(crate) struct BoundPredicates {
partition_bound_predicate: BoundPredicate,
snapshot_bound_predicate: BoundPredicate,
}
-/// Wraps a [`ManifestFile`] alongside the objects that are needed
-/// to process it in a thread-safe manner
-struct ManifestFileContext {
- manifest_file: ManifestFile,
-
- sender: Sender<ManifestEntryContext>,
-
- field_ids: Arc<Vec<i32>>,
- bound_predicates: Option<Arc<BoundPredicates>>,
- object_cache: Arc<ObjectCache>,
- snapshot_schema: SchemaRef,
- expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
- delete_file_index: Option<DeleteFileIndex>,
-}
-
-/// Wraps a [`ManifestEntryRef`] alongside the objects that are needed
-/// to process it in a thread-safe manner
-struct ManifestEntryContext {
- manifest_entry: ManifestEntryRef,
-
- expression_evaluator_cache: Arc<ExpressionEvaluatorCache>,
- field_ids: Arc<Vec<i32>>,
- bound_predicates: Option<Arc<BoundPredicates>>,
- partition_spec_id: i32,
- snapshot_schema: SchemaRef,
- delete_file_index: Option<DeleteFileIndex>,
-}
-
-impl ManifestFileContext {
- /// Consumes this [`ManifestFileContext`], fetching its Manifest from FileIO and then
- /// streaming its constituent [`ManifestEntries`] to the channel provided in the context
- async fn fetch_manifest_and_stream_manifest_entries(self) -> Result<()> {
- let ManifestFileContext {
- object_cache,
- manifest_file,
- bound_predicates,
- snapshot_schema,
- field_ids,
- mut sender,
- expression_evaluator_cache,
- delete_file_index,
- ..
- } = self;
-
- let manifest = object_cache.get_manifest(&manifest_file).await?;
-
- for manifest_entry in manifest.entries() {
- let manifest_entry_context = ManifestEntryContext {
- // TODO: refactor to avoid the expensive ManifestEntry clone
- manifest_entry: manifest_entry.clone(),
- expression_evaluator_cache: expression_evaluator_cache.clone(),
- field_ids: field_ids.clone(),
- partition_spec_id: manifest_file.partition_spec_id,
- bound_predicates: bound_predicates.clone(),
- snapshot_schema: snapshot_schema.clone(),
- delete_file_index: delete_file_index.clone(),
- };
-
- sender
- .send(manifest_entry_context)
- .map_err(|_| Error::new(ErrorKind::Unexpected, "mpsc channel SendError"))
- .await?;
- }
-
- Ok(())
- }
-}
-
-impl ManifestEntryContext {
- /// consume this `ManifestEntryContext`, returning a `FileScanTask`
- /// created from it
- async fn into_file_scan_task(self) -> Result<FileScanTask> {
- let deletes = if let Some(delete_file_index) = self.delete_file_index {
- delete_file_index
- .get_deletes_for_data_file(
- self.manifest_entry.data_file(),
- self.manifest_entry.sequence_number(),
- )
- .await?
- } else {
- vec![]
- };
-
- Ok(FileScanTask {
- start: 0,
- length: self.manifest_entry.file_size_in_bytes(),
- record_count: Some(self.manifest_entry.record_count()),
-
- data_file_path: self.manifest_entry.file_path().to_string(),
- data_file_content: self.manifest_entry.content_type(),
- data_file_format: self.manifest_entry.file_format(),
-
- schema: self.snapshot_schema,
- project_field_ids: self.field_ids.to_vec(),
- predicate: self
- .bound_predicates
- .map(|x| x.as_ref().snapshot_bound_predicate.clone()),
-
- deletes,
- })
- }
-}
-
-impl PlanContext {
- async fn get_manifest_list(&self) -> Result<Arc<ManifestList>> {
- self.object_cache
- .as_ref()
- .get_manifest_list(&self.snapshot, &self.table_metadata)
- .await
- }
-
- fn get_partition_filter(&self, manifest_file: &ManifestFile) -> Result<Arc<BoundPredicate>> {
- let partition_spec_id = manifest_file.partition_spec_id;
-
- let partition_filter = self.partition_filter_cache.get(
- partition_spec_id,
- &self.table_metadata,
- &self.snapshot_schema,
- self.case_sensitive,
- self.predicate
- .as_ref()
- .ok_or(Error::new(
- ErrorKind::Unexpected,
- "Expected a predicate but none present",
- ))?
- .as_ref()
- .bind(self.snapshot_schema.clone(), self.case_sensitive)?,
- )?;
-
- Ok(partition_filter)
- }
-
- fn build_manifest_file_contexts(
- &self,
- manifest_list: Arc<ManifestList>,
- tx_data: Sender<ManifestEntryContext>,
- delete_file_idx_and_tx: Option<(DeleteFileIndex, Sender<ManifestEntryContext>)>,
- ) -> Result<Box<impl Iterator<Item = Result<ManifestFileContext>> + 'static>> {
- let manifest_files = manifest_list.entries().iter();
-
- // TODO: Ideally we could ditch this intermediate Vec as we return an iterator.
- let mut filtered_mfcs = vec![];
-
- for manifest_file in manifest_files {
- let (delete_file_idx, tx) = if manifest_file.content == ManifestContentType::Deletes {
- let Some((delete_file_idx, tx)) = delete_file_idx_and_tx.as_ref() else {
- continue;
- };
- (Some(delete_file_idx.clone()), tx.clone())
- } else {
- (
- delete_file_idx_and_tx.as_ref().map(|x| x.0.clone()),
- tx_data.clone(),
- )
- };
-
- let partition_bound_predicate = if self.predicate.is_some() {
- let partition_bound_predicate = self.get_partition_filter(manifest_file)?;
-
- // evaluate the ManifestFile against the partition filter. Skip
- // if it cannot contain any matching rows
- if !self
- .manifest_evaluator_cache
- .get(
- manifest_file.partition_spec_id,
- partition_bound_predicate.clone(),
- )
- .eval(manifest_file)?
- {
- continue;
- }
-
- Some(partition_bound_predicate)
- } else {
- None
- };
-
- let mfc = self.create_manifest_file_context(
- manifest_file,
- partition_bound_predicate,
- tx,
- delete_file_idx,
- );
-
- filtered_mfcs.push(Ok(mfc));
- }
-
- Ok(Box::new(filtered_mfcs.into_iter()))
- }
-
- fn create_manifest_file_context(
- &self,
- manifest_file: &ManifestFile,
- partition_filter: Option<Arc<BoundPredicate>>,
- sender: Sender<ManifestEntryContext>,
- delete_file_index: Option<DeleteFileIndex>,
- ) -> ManifestFileContext {
- let bound_predicates =
- if let (Some(ref partition_bound_predicate), Some(snapshot_bound_predicate)) =
- (partition_filter, &self.snapshot_bound_predicate)
- {
- Some(Arc::new(BoundPredicates {
- partition_bound_predicate: partition_bound_predicate.as_ref().clone(),
- snapshot_bound_predicate: snapshot_bound_predicate.as_ref().clone(),
- }))
- } else {
- None
- };
-
- ManifestFileContext {
- manifest_file: manifest_file.clone(),
- bound_predicates,
- sender,
- object_cache: self.object_cache.clone(),
- snapshot_schema: self.snapshot_schema.clone(),
- field_ids: self.field_ids.clone(),
- expression_evaluator_cache: self.expression_evaluator_cache.clone(),
- delete_file_index,
- }
- }
-}
-
-/// Manages the caching of [`BoundPredicate`] objects
-/// for [`PartitionSpec`]s based on partition spec id.
-#[derive(Debug)]
-struct PartitionFilterCache(RwLock<HashMap<i32, Arc<BoundPredicate>>>);
-
-impl PartitionFilterCache {
- /// Creates a new [`PartitionFilterCache`]
- /// with an empty internal HashMap.
- fn new() -> Self {
- Self(RwLock::new(HashMap::new()))
- }
-
- /// Retrieves a [`BoundPredicate`] from the cache
- /// or computes it if not present.
- fn get(
- &self,
- spec_id: i32,
- table_metadata: &TableMetadataRef,
- schema: &Schema,
- case_sensitive: bool,
- filter: BoundPredicate,
- ) -> Result<Arc<BoundPredicate>> {
- // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
- // below, otherwise we hit deadlock
- {
- let read = self.0.read().map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "PartitionFilterCache RwLock was poisoned",
- )
- })?;
-
- if read.contains_key(&spec_id) {
- return Ok(read.get(&spec_id).unwrap().clone());
- }
- }
-
- let partition_spec = table_metadata
- .partition_spec_by_id(spec_id)
- .ok_or(Error::new(
- ErrorKind::Unexpected,
- format!("Could not find partition spec for id {}", spec_id),
- ))?;
-
- let partition_type = partition_spec.partition_type(schema)?;
- let partition_fields = partition_type.fields().to_owned();
- let partition_schema = Arc::new(
- Schema::builder()
- .with_schema_id(partition_spec.spec_id())
- .with_fields(partition_fields)
- .build()?,
- );
-
- let mut inclusive_projection = InclusiveProjection::new(partition_spec.clone());
-
- let partition_filter = inclusive_projection
- .project(&filter)?
- .rewrite_not()
- .bind(partition_schema.clone(), case_sensitive)?;
-
- self.0
- .write()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "PartitionFilterCache RwLock was poisoned",
- )
- })?
- .insert(spec_id, Arc::new(partition_filter));
-
- let read = self.0.read().map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "PartitionFilterCache RwLock was poisoned",
- )
- })?;
-
- Ok(read.get(&spec_id).unwrap().clone())
- }
-}
-
-/// Manages the caching of [`ManifestEvaluator`] objects
-/// for [`PartitionSpec`]s based on partition spec id.
-#[derive(Debug)]
-struct ManifestEvaluatorCache(RwLock<HashMap<i32, Arc<ManifestEvaluator>>>);
-
-impl ManifestEvaluatorCache {
- /// Creates a new [`ManifestEvaluatorCache`]
- /// with an empty internal HashMap.
- fn new() -> Self {
- Self(RwLock::new(HashMap::new()))
- }
-
- /// Retrieves a [`ManifestEvaluator`] from the cache
- /// or computes it if not present.
- fn get(&self, spec_id: i32, partition_filter: Arc<BoundPredicate>) -> Arc<ManifestEvaluator> {
- // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
- // below, otherwise we hit deadlock
- {
- let read = self
- .0
- .read()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "ManifestEvaluatorCache RwLock was poisoned",
- )
- })
- .unwrap();
-
- if read.contains_key(&spec_id) {
- return read.get(&spec_id).unwrap().clone();
- }
- }
-
- self.0
- .write()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "ManifestEvaluatorCache RwLock was poisoned",
- )
- })
- .unwrap()
- .insert(
- spec_id,
- Arc::new(ManifestEvaluator::new(partition_filter.as_ref().clone())),
- );
-
- let read = self
- .0
- .read()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "ManifestEvaluatorCache RwLock was poisoned",
- )
- })
- .unwrap();
-
- read.get(&spec_id).unwrap().clone()
- }
-}
-
-/// Manages the caching of [`ExpressionEvaluator`] objects
-/// for [`PartitionSpec`]s based on partition spec id.
-#[derive(Debug)]
-struct ExpressionEvaluatorCache(RwLock<HashMap<i32, Arc<ExpressionEvaluator>>>);
-
-impl ExpressionEvaluatorCache {
- /// Creates a new [`ExpressionEvaluatorCache`]
- /// with an empty internal HashMap.
- fn new() -> Self {
- Self(RwLock::new(HashMap::new()))
- }
-
- /// Retrieves a [`ExpressionEvaluator`] from the cache
- /// or computes it if not present.
- fn get(
- &self,
- spec_id: i32,
- partition_filter: &BoundPredicate,
- ) -> Result<Arc<ExpressionEvaluator>> {
- // we need a block here to ensure that the `read()` gets dropped before we hit the `write()`
- // below, otherwise we hit deadlock
- {
- let read = self.0.read().map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "PartitionFilterCache RwLock was poisoned",
- )
- })?;
-
- if read.contains_key(&spec_id) {
- return Ok(read.get(&spec_id).unwrap().clone());
- }
- }
-
- self.0
- .write()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "ManifestEvaluatorCache RwLock was poisoned",
- )
- })
- .unwrap()
- .insert(
- spec_id,
- Arc::new(ExpressionEvaluator::new(partition_filter.clone())),
- );
-
- let read = self
- .0
- .read()
- .map_err(|_| {
- Error::new(
- ErrorKind::Unexpected,
- "ManifestEvaluatorCache RwLock was poisoned",
- )
- })
- .unwrap();
-
- Ok(read.get(&spec_id).unwrap().clone())
- }
-}
-
-/// A task to scan part of file.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-pub struct FileScanTask {
- /// The start offset of the file to scan.
- pub start: u64,
- /// The length of the file to scan.
- pub length: u64,
- /// The number of records in the file to scan.
- ///
- /// This is an optional field, and only available if we are
- /// reading the entire data file.
- pub record_count: Option<u64>,
-
- /// The data file path corresponding to the task.
- pub data_file_path: String,
-
- /// The content type of the file to scan.
- pub data_file_content: DataContentType,
-
- /// The format of the file to scan.
- pub data_file_format: DataFileFormat,
-
- /// The schema of the file to scan.
- pub schema: SchemaRef,
- /// The field ids to project.
- pub project_field_ids: Vec<i32>,
- /// The predicate to filter.
- #[serde(skip_serializing_if = "Option::is_none")]
- pub predicate: Option<BoundPredicate>,
-
- /// The list of delete files that may need to be applied to this data file
- pub deletes: Vec<FileScanTaskDeleteFile>,
-}
-
-/// A task to scan part of file.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-pub struct FileScanTaskDeleteFile {
- /// The delete file path
- pub file_path: String,
-
- /// delete file type
- pub file_type: DataContentType,
-
- /// partition id
- pub partition_spec_id: i32,
-}
-
-#[derive(Debug)]
-pub(crate) struct DeleteFileContext {
- pub(crate) manifest_entry: ManifestEntryRef,
- pub(crate) partition_spec_id: i32,
-}
-
-impl From<&DeleteFileContext> for FileScanTaskDeleteFile {
- fn from(ctx: &DeleteFileContext) -> Self {
- FileScanTaskDeleteFile {
- file_path: ctx.manifest_entry.file_path().to_string(),
- file_type: ctx.manifest_entry.content_type(),
- partition_spec_id: ctx.partition_spec_id,
- }
- }
-}
-
-impl FileScanTask {
- /// Returns the data file path of this file scan task.
- pub fn data_file_path(&self) -> &str {
- &self.data_file_path
- }
-
- /// Returns the project field id of this file scan task.
- pub fn project_field_ids(&self) -> &[i32] {
- &self.project_field_ids
- }
-
- /// Returns the predicate of this file scan task.
- pub fn predicate(&self) -> Option<&BoundPredicate> {
- self.predicate.as_ref()
- }
-
- /// Returns the schema of this file scan task as a reference
- pub fn schema(&self) -> &Schema {
- &self.schema
- }
-
- /// Returns the schema of this file scan task as a SchemaRef
- pub fn schema_ref(&self) -> SchemaRef {
- self.schema.clone()
- }
-}
-
#[cfg(test)]
pub mod tests {
use std::collections::HashMap;
diff --git a/crates/iceberg/src/scan/task.rs b/crates/iceberg/src/scan/task.rs
new file mode 100644
index 00000000..a5f4e462
--- /dev/null
+++ b/crates/iceberg/src/scan/task.rs
@@ -0,0 +1,116 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use futures::stream::BoxStream;
+use serde::{Deserialize, Serialize};
+
+use crate::expr::BoundPredicate;
+use crate::spec::{DataContentType, DataFileFormat, ManifestEntryRef, Schema, SchemaRef};
+use crate::Result;
+
+/// A stream of [`FileScanTask`].
+pub type FileScanTaskStream = BoxStream<'static, Result<FileScanTask>>;
+
+/// A task to scan part of file.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct FileScanTask {
+ /// The start offset of the file to scan.
+ pub start: u64,
+ /// The length of the file to scan.
+ pub length: u64,
+ /// The number of records in the file to scan.
+ ///
+ /// This is an optional field, and only available if we are
+ /// reading the entire data file.
+ pub record_count: Option<u64>,
+
+ /// The data file path corresponding to the task.
+ pub data_file_path: String,
+
+ /// The content type of the file to scan.
+ pub data_file_content: DataContentType,
+
+ /// The format of the file to scan.
+ pub data_file_format: DataFileFormat,
+
+ /// The schema of the file to scan.
+ pub schema: SchemaRef,
+ /// The field ids to project.
+ pub project_field_ids: Vec<i32>,
+ /// The predicate to filter.
+ #[serde(skip_serializing_if = "Option::is_none")]
+ pub predicate: Option<BoundPredicate>,
+
+ /// The list of delete files that may need to be applied to this data file
+ pub deletes: Vec<FileScanTaskDeleteFile>,
+}
+
+impl FileScanTask {
+ /// Returns the data file path of this file scan task.
+ pub fn data_file_path(&self) -> &str {
+ &self.data_file_path
+ }
+
+ /// Returns the project field id of this file scan task.
+ pub fn project_field_ids(&self) -> &[i32] {
+ &self.project_field_ids
+ }
+
+ /// Returns the predicate of this file scan task.
+ pub fn predicate(&self) -> Option<&BoundPredicate> {
+ self.predicate.as_ref()
+ }
+
+ /// Returns the schema of this file scan task as a reference
+ pub fn schema(&self) -> &Schema {
+ &self.schema
+ }
+
+ /// Returns the schema of this file scan task as a SchemaRef
+ pub fn schema_ref(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+}
+
+#[derive(Debug)]
+pub(crate) struct DeleteFileContext {
+ pub(crate) manifest_entry: ManifestEntryRef,
+ pub(crate) partition_spec_id: i32,
+}
+
+impl From<&DeleteFileContext> for FileScanTaskDeleteFile {
+ fn from(ctx: &DeleteFileContext) -> Self {
+ FileScanTaskDeleteFile {
+ file_path: ctx.manifest_entry.file_path().to_string(),
+ file_type: ctx.manifest_entry.content_type(),
+ partition_spec_id: ctx.partition_spec_id,
+ }
+ }
+}
+
+/// A task to scan part of file.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct FileScanTaskDeleteFile {
+ /// The delete file path
+ pub file_path: String,
+
+ /// delete file type
+ pub file_type: DataContentType,
+
+ /// partition id
+ pub partition_spec_id: i32,
+}