diff --git a/Cargo.lock b/Cargo.lock index 1feac7435c..37b0e10199 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2910,6 +2910,7 @@ dependencies = [ "regex", "reqwest 0.12.28", "roaring", + "rstest", "serde", "serde_arrow", "serde_bytes", diff --git a/crates/catalog/s3tables/src/catalog.rs b/crates/catalog/s3tables/src/catalog.rs index 254f20ae32..c96d06abd4 100644 --- a/crates/catalog/s3tables/src/catalog.rs +++ b/crates/catalog/s3tables/src/catalog.rs @@ -1345,7 +1345,7 @@ mod tests { .select_all() .build() .expect("scan to be valid (snapshot exists, schema is OK)") - .to_arrow() + .to_arrow(Default::default()) .await .expect("scan tasks should be OK") .try_collect() diff --git a/crates/examples/src/oss_backend.rs b/crates/examples/src/oss_backend.rs index 2d56cceb7f..5f56e4a301 100644 --- a/crates/examples/src/oss_backend.rs +++ b/crates/examples/src/oss_backend.rs @@ -81,7 +81,7 @@ async fn main() { println!("Table {TABLE_NAME} loaded!"); let scan = table.scan().select_all().build().unwrap(); - let reader = scan.to_arrow().await.unwrap(); + let reader = scan.to_arrow(Default::default()).await.unwrap(); let buf = reader.collect::>().await; println!("Table {TABLE_NAME} has {} batches.", buf.len()); diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index fa5e6cce73..a6f1bb7d35 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -99,6 +99,7 @@ mockall = { workspace = true } pretty_assertions = { workspace = true } rand = { workspace = true } regex = { workspace = true } +rstest = { workspace = true } serde_arrow = { version = "0.14", features = ["arrow-59"] } tempfile = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/crates/iceberg/README.md b/crates/iceberg/README.md index 08ce82bd2f..bb22750519 100644 --- a/crates/iceberg/README.md +++ b/crates/iceberg/README.md @@ -56,7 +56,7 @@ async fn main() -> Result<()> { .scan() .select(["name", "id"]) .build()? - .to_arrow() + .to_arrow(Default::default()) .await?; // Consume this stream like arrow record batch stream. diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 92267f0a99..d879e2c8a1 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -71,14 +71,54 @@ pub struct iceberg::arrow::ArrowReaderBuilder impl iceberg::arrow::ArrowReaderBuilder pub fn iceberg::arrow::ArrowReaderBuilder::build(self) -> iceberg::arrow::ArrowReader pub fn iceberg::arrow::ArrowReaderBuilder::new(file_io: iceberg::io::FileIO, runtime: iceberg::Runtime) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_batch_size(self, batch_size: usize) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_bloom_filter_enabled(self, bloom_filter_enabled: bool) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_data_file_concurrency_limit(self, val: usize) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_metadata_size_hint(self, metadata_size_hint: usize) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_range_coalesce_bytes(self, range_coalesce_bytes: u64) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_range_fetch_concurrency(self, range_fetch_concurrency: usize) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_row_group_filtering_enabled(self, row_group_filtering_enabled: bool) -> Self -pub fn iceberg::arrow::ArrowReaderBuilder::with_row_selection_enabled(self, row_selection_enabled: bool) -> Self +pub fn iceberg::arrow::ArrowReaderBuilder::with_options(self, options: iceberg::arrow::ArrowReaderOptions) -> Self +pub fn iceberg::arrow::ArrowReaderBuilder::with_parquet_options(self, options: iceberg::arrow::ParquetReadOptions) -> Self +pub struct iceberg::arrow::ArrowReaderOptions +impl iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::batch_size(&self) -> core::option::Option +pub fn iceberg::arrow::ArrowReaderOptions::bloom_filter_enabled(&self) -> bool +pub fn iceberg::arrow::ArrowReaderOptions::data_file_concurrency_limit(&self) -> core::option::Option +pub fn iceberg::arrow::ArrowReaderOptions::parquet_options(&self) -> &iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ArrowReaderOptions::row_group_filtering_enabled(&self) -> bool +pub fn iceberg::arrow::ArrowReaderOptions::row_selection_enabled(&self) -> bool +impl core::clone::Clone for iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::clone(&self) -> iceberg::arrow::ArrowReaderOptions +impl core::cmp::PartialEq for iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::eq(&self, other: &iceberg::arrow::ArrowReaderOptions) -> bool +impl core::default::Default for iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::default() -> Self +impl core::fmt::Debug for iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg::arrow::ArrowReaderOptions +impl iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::builder() -> ArrowReaderOptionsBuilder<((), (), (), (), (), ())> +impl serde_core::ser::Serialize for iceberg::arrow::ArrowReaderOptions +pub fn iceberg::arrow::ArrowReaderOptions::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg::arrow::ArrowReaderOptions where iceberg::arrow::ArrowReaderOptions: core::default::Default +pub fn iceberg::arrow::ArrowReaderOptions::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg::arrow::ParquetReadOptions +impl iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::metadata_size_hint(&self) -> usize +pub fn iceberg::arrow::ParquetReadOptions::preload_column_index(&self) -> bool +pub fn iceberg::arrow::ParquetReadOptions::preload_offset_index(&self) -> bool +pub fn iceberg::arrow::ParquetReadOptions::preload_page_index(&self) -> bool +pub fn iceberg::arrow::ParquetReadOptions::range_coalesce_bytes(&self) -> u64 +pub fn iceberg::arrow::ParquetReadOptions::range_fetch_concurrency(&self) -> usize +impl core::clone::Clone for iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::clone(&self) -> iceberg::arrow::ParquetReadOptions +impl core::cmp::PartialEq for iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::eq(&self, other: &iceberg::arrow::ParquetReadOptions) -> bool +impl core::default::Default for iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::default() -> Self +impl core::fmt::Debug for iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg::arrow::ParquetReadOptions +impl iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::builder() -> ParquetReadOptionsBuilder<((), (), (), (), (), ())> +impl serde_core::ser::Serialize for iceberg::arrow::ParquetReadOptions +pub fn iceberg::arrow::ParquetReadOptions::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg::arrow::ParquetReadOptions where iceberg::arrow::ParquetReadOptions: core::default::Default +pub fn iceberg::arrow::ParquetReadOptions::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg::arrow::PartitionValueCalculator impl iceberg::arrow::PartitionValueCalculator pub fn iceberg::arrow::PartitionValueCalculator::calculate(&self, batch: &arrow_array::record_batch::RecordBatch) -> iceberg::Result @@ -1437,7 +1477,7 @@ impl iceberg::scan::TableScan pub fn iceberg::scan::TableScan::column_names(&self) -> core::option::Option<&[alloc::string::String]> pub async fn iceberg::scan::TableScan::plan_files(&self) -> iceberg::Result pub fn iceberg::scan::TableScan::snapshot(&self) -> core::option::Option<&iceberg::spec::SnapshotRef> -pub async fn iceberg::scan::TableScan::to_arrow(&self) -> iceberg::Result +pub async fn iceberg::scan::TableScan::to_arrow(&self, options: iceberg::arrow::ArrowReaderOptions) -> iceberg::Result impl core::fmt::Debug for iceberg::scan::TableScan pub fn iceberg::scan::TableScan::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::scan::TableScanBuilder<'a> @@ -1448,15 +1488,10 @@ pub fn iceberg::scan::TableScanBuilder<'a>::select(self, column_names: impl core pub fn iceberg::scan::TableScanBuilder<'a>::select_all(self) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::select_empty(self) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::snapshot_id(self, snapshot_id: i64) -> Self -pub fn iceberg::scan::TableScanBuilder<'a>::with_batch_size(self, batch_size: core::option::Option) -> Self -pub fn iceberg::scan::TableScanBuilder<'a>::with_bloom_filter_enabled(self, bloom_filter_enabled: bool) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_case_sensitive(self, case_sensitive: bool) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_concurrency_limit(self, limit: usize) -> Self -pub fn iceberg::scan::TableScanBuilder<'a>::with_data_file_concurrency_limit(self, limit: usize) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_filter(self, predicate: iceberg::expr::Predicate) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_manifest_entry_concurrency_limit(self, limit: usize) -> Self -pub fn iceberg::scan::TableScanBuilder<'a>::with_row_group_filtering_enabled(self, row_group_filtering_enabled: bool) -> Self -pub fn iceberg::scan::TableScanBuilder<'a>::with_row_selection_enabled(self, row_selection_enabled: bool) -> Self pub type iceberg::scan::ArrowRecordBatchStream = futures_core::stream::BoxStream<'static, iceberg::Result> pub type iceberg::scan::FileScanTaskStream = futures_core::stream::BoxStream<'static, iceberg::Result> pub mod iceberg::sensitive diff --git a/crates/iceberg/src/arrow/reader/file_reader.rs b/crates/iceberg/src/arrow/reader/file_reader.rs index 63b7f06758..22355737a1 100644 --- a/crates/iceberg/src/arrow/reader/file_reader.rs +++ b/crates/iceberg/src/arrow/reader/file_reader.rs @@ -118,7 +118,7 @@ impl AsyncFileReader for ArrowFileReader { async move { let reader = ParquetMetaDataReader::new() - .with_prefetch_hint(self.parquet_read_options.metadata_size_hint()) + .with_prefetch_hint(Some(self.parquet_read_options.metadata_size_hint())) // Set the page policy first because it updates both column and offset policies. .with_page_index_policy(PageIndexPolicy::from( self.parquet_read_options.preload_page_index(), diff --git a/crates/iceberg/src/arrow/reader/mod.rs b/crates/iceberg/src/arrow/reader/mod.rs index cda1909846..bb0c0e8642 100644 --- a/crates/iceberg/src/arrow/reader/mod.rs +++ b/crates/iceberg/src/arrow/reader/mod.rs @@ -43,7 +43,7 @@ mod projection; mod row_filter; mod row_lineage; pub use file_reader::ArrowFileReader; -pub(crate) use options::ParquetReadOptions; +pub use options::{ArrowReaderOptions, ParquetReadOptions}; use predicate_visitor::{CollectFieldIdVisitor, PredicateConverter}; use projection::{ add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, @@ -52,114 +52,50 @@ use projection::{ /// Builder to create ArrowReader pub struct ArrowReaderBuilder { - batch_size: Option, + options: ArrowReaderOptions, file_io: FileIO, - concurrency_limit_data_files: usize, - row_group_filtering_enabled: bool, - row_selection_enabled: bool, - bloom_filter_enabled: bool, - parquet_read_options: ParquetReadOptions, runtime: Runtime, } impl ArrowReaderBuilder { /// Create a new ArrowReaderBuilder pub fn new(file_io: FileIO, runtime: Runtime) -> Self { - let num_cpus = available_parallelism().get(); - ArrowReaderBuilder { - batch_size: None, + options: ArrowReaderOptions::default(), file_io, - concurrency_limit_data_files: num_cpus, - row_group_filtering_enabled: true, - row_selection_enabled: false, - bloom_filter_enabled: false, - parquet_read_options: ParquetReadOptions::builder().build(), runtime, } } - /// Sets the max number of in flight data files that are being fetched - pub fn with_data_file_concurrency_limit(mut self, val: usize) -> Self { - self.concurrency_limit_data_files = val; - self - } - - /// Sets the desired size of batches in the response - /// to something other than the default - pub fn with_batch_size(mut self, batch_size: usize) -> Self { - self.batch_size = Some(batch_size); - self - } - - /// Determines whether to enable row group filtering. - pub fn with_row_group_filtering_enabled(mut self, row_group_filtering_enabled: bool) -> Self { - self.row_group_filtering_enabled = row_group_filtering_enabled; + /// Replaces all reader settings while preserving local I/O and runtime resources. + pub fn with_options(mut self, options: ArrowReaderOptions) -> Self { + self.options = options; self } - /// Determines whether to enable row selection. - pub fn with_row_selection_enabled(mut self, row_selection_enabled: bool) -> Self { - self.row_selection_enabled = row_selection_enabled; - self - } - - /// Determines whether to enable bloom filter-based row group filtering. - /// - /// When enabled, if a read is performed with an equality or IN predicate, - /// the bloom filter for relevant columns in each row group is read and - /// checked. Row groups where the bloom filter proves the value is absent - /// are skipped entirely. - /// - /// Defaults to disabled. Each bloom filter is a separate read, and they are - /// issued serially — one round trip per relevant column per row group, before - /// any data is read. TODO(#3191) - pub fn with_bloom_filter_enabled(mut self, bloom_filter_enabled: bool) -> Self { - self.bloom_filter_enabled = bloom_filter_enabled; - self - } - - /// Provide a hint as to the number of bytes to prefetch for parsing the Parquet metadata - /// - /// This hint can help reduce the number of fetch requests. For more details see the - /// [ParquetMetaDataReader documentation](https://docs.rs/parquet/latest/parquet/file/metadata/struct.ParquetMetaDataReader.html#method.with_prefetch_hint). - pub fn with_metadata_size_hint(mut self, metadata_size_hint: usize) -> Self { - self.parquet_read_options.metadata_size_hint = Some(metadata_size_hint); - self - } - - /// Sets the gap threshold for merging nearby byte ranges into a single request. - /// Ranges with gaps smaller than this value will be coalesced. - /// - /// Defaults to 1 MiB, matching object_store's OBJECT_STORE_COALESCE_DEFAULT. - pub fn with_range_coalesce_bytes(mut self, range_coalesce_bytes: u64) -> Self { - self.parquet_read_options.range_coalesce_bytes = range_coalesce_bytes; - self - } - - /// Sets the maximum number of merged byte ranges to fetch concurrently. - /// - /// Defaults to 10, matching object_store's OBJECT_STORE_COALESCE_PARALLEL. - pub fn with_range_fetch_concurrency(mut self, range_fetch_concurrency: usize) -> Self { - self.parquet_read_options.range_fetch_concurrency = range_fetch_concurrency; + /// Replaces Parquet I/O settings while preserving other reader settings. + pub fn with_parquet_options(mut self, options: ParquetReadOptions) -> Self { + self.options = self.options.with_parquet_options(options); self } /// Build the ArrowReader. pub fn build(self) -> ArrowReader { + let options = self + .options + .with_data_file_concurrency_limit_if_unset(available_parallelism().get().max(1)); + let concurrency_limit_data_files = options + .data_file_concurrency_limit() + .expect("reader concurrency was set above") + .max(1); ArrowReader { - batch_size: self.batch_size, + options, file_io: self.file_io.clone(), delete_file_loader: CachingDeleteFileLoader::new( self.file_io.clone(), - self.concurrency_limit_data_files, + concurrency_limit_data_files, self.runtime.clone(), ), - concurrency_limit_data_files: self.concurrency_limit_data_files, - row_group_filtering_enabled: self.row_group_filtering_enabled, - row_selection_enabled: self.row_selection_enabled, - bloom_filter_enabled: self.bloom_filter_enabled, - parquet_read_options: self.parquet_read_options, runtime: self.runtime.clone(), } } @@ -168,16 +104,90 @@ impl ArrowReaderBuilder { /// Reads data from Parquet files #[derive(Clone)] pub struct ArrowReader { - batch_size: Option, + options: ArrowReaderOptions, file_io: FileIO, delete_file_loader: CachingDeleteFileLoader, - /// the maximum number of data files that can be fetched at the same time - concurrency_limit_data_files: usize, - - row_group_filtering_enabled: bool, - row_selection_enabled: bool, - bloom_filter_enabled: bool, - parquet_read_options: ParquetReadOptions, runtime: Runtime, } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use futures::{TryStreamExt, stream}; + + use super::{ArrowReaderBuilder, ArrowReaderOptions, ParquetReadOptions}; + use crate::io::FileIO; + use crate::test_utils::test_runtime; + use crate::util::available_parallelism; + + #[tokio::test] + async fn deserialized_zero_concurrency_completes_reads() { + let options: ArrowReaderOptions = + serde_json::from_str(r#"{"data_file_concurrency_limit":0}"#).unwrap(); + let reader = ArrowReaderBuilder::new(FileIO::new_with_memory(), test_runtime()) + .with_options(options) + .build(); + let batches = tokio::time::timeout( + Duration::from_secs(1), + reader + .read(Box::pin(stream::empty())) + .unwrap() + .stream() + .try_collect::>(), + ) + .await + .expect("zero concurrency must not stall the task stream") + .unwrap(); + assert!(batches.is_empty()); + } + + #[test] + fn default_options_restore_local_reader_defaults() { + let reader = ArrowReaderBuilder::new(FileIO::new_with_memory(), test_runtime()) + .with_options( + ArrowReaderOptions::builder() + .with_batch_size(1024) + .with_data_file_concurrency_limit(3) + .with_row_group_filtering_enabled(false) + .with_row_selection_enabled(true) + .with_bloom_filter_enabled(true) + .with_parquet_options( + ParquetReadOptions::builder() + .with_range_fetch_concurrency(99) + .build(), + ) + .build(), + ) + .with_options(ArrowReaderOptions::default()) + .build(); + assert_eq!( + reader.options, + ArrowReaderOptions::builder() + .with_data_file_concurrency_limit(available_parallelism().get()) + .build() + ); + } + + #[test] + fn explicit_options_configure_reader_settings() { + let parquet = ParquetReadOptions::builder() + .with_metadata_size_hint(4096) + .with_range_coalesce_bytes(8192) + .with_range_fetch_concurrency(2) + .build(); + let options = ArrowReaderOptions::builder() + .with_batch_size(17) + .with_data_file_concurrency_limit(3) + .with_row_group_filtering_enabled(false) + .with_row_selection_enabled(true) + .with_bloom_filter_enabled(true) + .with_parquet_options(parquet) + .build(); + let reader = ArrowReaderBuilder::new(FileIO::new_with_memory(), test_runtime()) + .with_options(options.clone()) + .build(); + assert_eq!(reader.options, options); + } +} diff --git a/crates/iceberg/src/arrow/reader/options.rs b/crates/iceberg/src/arrow/reader/options.rs index ae6a3ed18e..ce69c15491 100644 --- a/crates/iceberg/src/arrow/reader/options.rs +++ b/crates/iceberg/src/arrow/reader/options.rs @@ -15,70 +15,286 @@ // specific language governing permissions and limitations // under the License. -//! Tunables for Parquet file I/O used by `ArrowReader`. - +use serde::{Deserialize, Serialize}; use typed_builder::TypedBuilder; use super::{ DEFAULT_METADATA_SIZE_HINT, DEFAULT_RANGE_COALESCE_BYTES, DEFAULT_RANGE_FETCH_CONCURRENCY, }; +/// Configures Arrow reads independently of file planning and local I/O resources. +/// +/// Serializable alongside planned file tasks so readers can be constructed later. +/// In self-describing map formats such as JSON, field names are part of the +/// serialized format: missing fields use defaults and unknown fields are ignored. +/// Compatibility across different crate versions is not otherwise guaranteed. +/// +/// ```no_run +/// # use iceberg::{Result, arrow::{ArrowReaderOptions, ParquetReadOptions}, table::Table}; +/// # async fn read(table: &Table) -> Result<()> { +/// let options = ArrowReaderOptions::builder() +/// .with_batch_size(1024) +/// .with_parquet_options( +/// ParquetReadOptions::builder() +/// .with_metadata_size_hint(512 * 1024) +/// .build(), +/// ) +/// .build(); +/// let stream = table.scan().build()?.to_arrow(options).await?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[serde(default)] +#[builder(field_defaults(default, setter(prefix = "with_")))] +pub struct ArrowReaderOptions { + /// Desired batch size, or the reader default when unset. + #[builder(default = None, setter(strip_option))] + batch_size: Option, + /// Maximum in-flight data files, or the construction path's default when unset. + #[builder(default = None, setter(strip_option))] + data_file_concurrency_limit: Option, + /// Whether to prune row groups using their statistics. + #[builder(default = true)] + row_group_filtering_enabled: bool, + /// Whether to use page indexes for row selection. + row_selection_enabled: bool, + /// Whether to use bloom filters for row-group pruning. + bloom_filter_enabled: bool, + /// I/O tuning applied when reading Parquet files. + parquet_options: ParquetReadOptions, +} + +impl Default for ArrowReaderOptions { + fn default() -> Self { + Self::builder().build() + } +} + +impl ArrowReaderOptions { + /// Supplies a fallback limit without replacing an explicit reader limit. + pub(super) fn with_data_file_concurrency_limit_if_unset(mut self, limit: usize) -> Self { + self.data_file_concurrency_limit.get_or_insert(limit); + self + } + + /// Returns the configured batch size. + pub fn batch_size(&self) -> Option { + self.batch_size + } + + /// Returns the configured data file concurrency limit. + /// + /// When unset, reader construction uses execution-host parallelism. + /// Zero is treated as one. + pub fn data_file_concurrency_limit(&self) -> Option { + self.data_file_concurrency_limit + } + + /// Returns whether row-group statistics filtering is enabled. + /// + /// Enabled by default: statistics can skip groups that cannot match the + /// predicate, generally improving performance with little overhead. + pub fn row_group_filtering_enabled(&self) -> bool { + self.row_group_filtering_enabled + } + + /// Returns whether page-index row selection is enabled. + /// + /// Disabled by default: parsing page indexes can cost more than the reads + /// it avoids. Evaluate it against the table's sorting, partitioning, row-group + /// size, and page size before enabling it. + pub fn row_selection_enabled(&self) -> bool { + self.row_selection_enabled + } + + /// Returns whether bloom-filter pruning is enabled. + /// + /// Disabled by default: bloom filters require extra reads per relevant + /// column and row group. They can eliminate groups for equality and IN predicates. + pub fn bloom_filter_enabled(&self) -> bool { + self.bloom_filter_enabled + } + + /// Returns the configured parquet options. + pub fn parquet_options(&self) -> &ParquetReadOptions { + &self.parquet_options + } + + pub(super) fn with_parquet_options(mut self, value: ParquetReadOptions) -> Self { + self.parquet_options = value; + self + } +} + /// Options for tuning Parquet file I/O. -#[derive(Clone, Copy, Debug, TypedBuilder)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[serde(default)] #[builder(field_defaults(setter(prefix = "with_")))] -pub(crate) struct ParquetReadOptions { +pub struct ParquetReadOptions { /// Number of bytes to prefetch for parsing the Parquet metadata. /// /// This hint can help reduce the number of fetch requests. For more details see the /// [ParquetMetaDataReader documentation](https://docs.rs/parquet/latest/parquet/file/metadata/struct.ParquetMetaDataReader.html#method.with_prefetch_hint). /// /// Defaults to 512 KiB, matching DataFusion's default `ParquetOptions::metadata_size_hint`. - #[builder(default = Some(DEFAULT_METADATA_SIZE_HINT))] - pub(crate) metadata_size_hint: Option, + #[builder(default = DEFAULT_METADATA_SIZE_HINT)] + metadata_size_hint: usize, /// Gap threshold for merging nearby byte ranges into a single request. /// Ranges with gaps smaller than this value will be coalesced. /// /// Defaults to 1 MiB, matching object_store's `OBJECT_STORE_COALESCE_DEFAULT`. #[builder(default = DEFAULT_RANGE_COALESCE_BYTES)] - pub(crate) range_coalesce_bytes: u64, + range_coalesce_bytes: u64, /// Maximum number of merged byte ranges to fetch concurrently. /// /// Defaults to 10, matching object_store's `OBJECT_STORE_COALESCE_PARALLEL`. #[builder(default = DEFAULT_RANGE_FETCH_CONCURRENCY)] - pub(crate) range_fetch_concurrency: usize, - /// Whether to preload the column index when reading Parquet metadata. + range_fetch_concurrency: usize, #[builder(default = true)] - pub(crate) preload_column_index: bool, - /// Whether to preload the offset index when reading Parquet metadata. + preload_column_index: bool, #[builder(default = true)] - pub(crate) preload_offset_index: bool, - /// Whether to preload the page index when reading Parquet metadata. + preload_offset_index: bool, + // Execution derives page-index preloading per file from predicates and deletes. #[builder(default = false)] - pub(crate) preload_page_index: bool, + preload_page_index: bool, +} + +impl Default for ParquetReadOptions { + fn default() -> Self { + Self::builder().build() + } } impl ParquetReadOptions { - pub(crate) fn metadata_size_hint(&self) -> Option { + pub(super) fn with_preload_page_index(mut self, enabled: bool) -> Self { + self.preload_page_index = enabled; + self + } + + /// Returns whether to preload the column index when reading Parquet metadata. + pub fn preload_column_index(&self) -> bool { + self.preload_column_index + } + + /// Returns whether to preload the offset index when reading Parquet metadata. + pub fn preload_offset_index(&self) -> bool { + self.preload_offset_index + } + + /// Returns whether to preload the page index when reading Parquet metadata. + pub fn preload_page_index(&self) -> bool { + self.preload_page_index + } + + /// Returns the configured metadata size hint. + pub fn metadata_size_hint(&self) -> usize { self.metadata_size_hint } - pub(crate) fn range_coalesce_bytes(&self) -> u64 { + /// Returns the byte-gap threshold for coalescing read requests. + pub fn range_coalesce_bytes(&self) -> u64 { self.range_coalesce_bytes } - pub(crate) fn range_fetch_concurrency(&self) -> usize { + /// Returns the maximum number of concurrent range requests. + pub fn range_fetch_concurrency(&self) -> usize { self.range_fetch_concurrency } +} - pub(crate) fn preload_column_index(&self) -> bool { - self.preload_column_index +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::unset(None, 7)] + #[case::explicit(Some(3), 3)] + #[case::zero(Some(0), 0)] + fn default_data_file_concurrency_preserves_explicit_options( + #[case] configured: Option, + #[case] expected: usize, + ) { + // Include unrelated settings to check that resolving concurrency preserves them. + let options: ArrowReaderOptions = serde_json::from_value(serde_json::json!({ + "batch_size": 17, + "data_file_concurrency_limit": configured, + "row_selection_enabled": true, + })) + .unwrap(); + // Resolve the reader fallback; explicit zero is normalized later by the reader. + let options = options.with_data_file_concurrency_limit_if_unset(7); + let expected = ArrowReaderOptions::builder() + .with_batch_size(17) + .with_data_file_concurrency_limit(expected) + .with_row_selection_enabled(true) + .build(); + assert_eq!(options, expected); } - pub(crate) fn preload_offset_index(&self) -> bool { - self.preload_offset_index + #[test] + fn missing_options_use_defaults() { + let decoded: ArrowReaderOptions = serde_json::from_str(r#"{"batch_size":17}"#).unwrap(); + assert_eq!( + decoded, + ArrowReaderOptions::builder().with_batch_size(17).build() + ); } - pub(crate) fn preload_page_index(&self) -> bool { - self.preload_page_index + #[test] + fn unknown_options_are_ignored() { + let decoded: ArrowReaderOptions = + serde_json::from_str(r#"{"future_option":true}"#).unwrap(); + assert_eq!(decoded, ArrowReaderOptions::default()); + } + + #[test] + fn unset_concurrency_serializes_as_null() { + assert_eq!( + serde_json::to_value(ArrowReaderOptions::default()).unwrap()["data_file_concurrency_limit"], + serde_json::Value::Null, + ); + } + + #[test] + fn parquet_serialization_preserves_preload_settings() { + let options = ParquetReadOptions { + preload_column_index: false, + preload_offset_index: false, + preload_page_index: true, + ..ParquetReadOptions::default() + }; + let encoded = serde_json::to_value(&options).unwrap(); + assert_eq!( + encoded, + serde_json::json!({ + "metadata_size_hint": DEFAULT_METADATA_SIZE_HINT, + "range_coalesce_bytes": DEFAULT_RANGE_COALESCE_BYTES, + "range_fetch_concurrency": DEFAULT_RANGE_FETCH_CONCURRENCY, + "preload_column_index": false, + "preload_offset_index": false, + "preload_page_index": true, + }), + ); + let decoded: ParquetReadOptions = serde_json::from_value(encoded).unwrap(); + assert_eq!(decoded, options); + } + + #[test] + fn partial_parquet_options_preserve_defaults() { + let options: ArrowReaderOptions = + serde_json::from_str(r#"{"parquet_options":{"range_fetch_concurrency":2}}"#).unwrap(); + assert_eq!( + options, + ArrowReaderOptions::builder() + .with_parquet_options( + ParquetReadOptions::builder() + .with_range_fetch_concurrency(2) + .build(), + ) + .build() + ); } } diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 1ddee493c4..30864a02f8 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -65,19 +65,23 @@ impl ArrowReader { /// Take a stream of FileScanTasks and reads all the files. /// Returns a [`ScanResult`] containing the record batch stream and scan metrics. pub fn read(self, tasks: FileScanTaskStream) -> Result { - let concurrency_limit_data_files = self.concurrency_limit_data_files; + let concurrency_limit_data_files = self + .options + .data_file_concurrency_limit() + .expect("reader construction sets concurrency") + .max(1); let scan_metrics = ScanMetrics::new(); let task_reader = FileScanTaskReader { - batch_size: self.batch_size, + batch_size: self.options.batch_size(), file_io: self.file_io, delete_file_loader: self .delete_file_loader .with_scan_metrics(scan_metrics.clone()), - row_group_filtering_enabled: self.row_group_filtering_enabled, - row_selection_enabled: self.row_selection_enabled, - bloom_filter_enabled: self.bloom_filter_enabled, - parquet_read_options: self.parquet_read_options, + row_group_filtering_enabled: self.options.row_group_filtering_enabled(), + row_selection_enabled: self.options.row_selection_enabled(), + bloom_filter_enabled: self.options.bloom_filter_enabled(), + parquet_read_options: self.options.parquet_options().clone(), scan_metrics: scan_metrics.clone(), runtime: self.runtime, }; @@ -153,8 +157,6 @@ impl FileScanTaskReader { async fn process(self, task: FileScanTask) -> Result { let should_load_page_index = (self.row_selection_enabled && task.predicate().is_some()) || !task.deletes().is_empty(); - let mut parquet_read_options = self.parquet_read_options; - parquet_read_options.preload_page_index = should_load_page_index; let delete_filter_rx = self .delete_file_loader @@ -169,6 +171,10 @@ impl FileScanTaskReader { let file_size_in_bytes = task.file_size_in_bytes(); let bytes_read = Arc::clone(self.scan_metrics.bytes_read_counter()); let key_metadata = task.key_metadata().map(|km| km.to_vec()); + let parquet_read_options = self + .parquet_read_options + .clone() + .with_preload_page_index(should_load_page_index); let opening = self .runtime .io() @@ -903,8 +909,8 @@ mod tests { use tempfile::TempDir; use crate::Runtime; - use crate::arrow::ArrowReaderBuilder; use crate::arrow::test_utils::write_encrypted_parquet; + use crate::arrow::{ArrowReaderBuilder, ArrowReaderOptions as ReaderOptions}; use crate::io::FileIO; use crate::metadata_columns::{ RESERVED_COL_NAME_POS, RESERVED_COL_NAME_ROW_ID, RESERVED_FIELD_ID_FILE, @@ -2330,7 +2336,11 @@ mod tests { // Row selection must be enabled for the predicate to produce a RowSelection. let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; let batches: Vec = reader @@ -2371,7 +2381,11 @@ mod tests { ); let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; let batches: Vec = reader @@ -2722,7 +2736,11 @@ mod tests { // Read with concurrency=1 (fast-path) let reader = ArrowReaderBuilder::new(file_io, Runtime::current()) - .with_data_file_concurrency_limit(1) + .with_options( + ReaderOptions::builder() + .with_data_file_concurrency_limit(1) + .build(), + ) .build(); // Create tasks in a specific order: file_0, file_1, file_2 @@ -3554,7 +3572,11 @@ mod tests { // reads `id` for its own evaluation even though `id` is not projected; the surviving // rows must keep their ABSOLUTE positions (0 and 2), not renumbered (0 and 1). let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; let batches: Vec = reader @@ -3957,6 +3979,7 @@ mod tests { use super::*; use crate::Result; + use crate::arrow::ParquetReadOptions; use crate::arrow::scan_metrics::ScanMetrics; use crate::expr::{Bind, Predicate, Reference}; use crate::io::{ @@ -4141,9 +4164,17 @@ mod tests { let file_io = FileIOBuilder::new(Arc::new(DeferredStorageFactory { storage })).build(); let reader = ArrowReaderBuilder::new(file_io, Runtime::current()) - .with_data_file_concurrency_limit(CONCURRENCY) - .with_metadata_size_hint(METADATA_SIZE_HINT) - .with_bloom_filter_enabled(bound.is_some()) + .with_options( + ReaderOptions::builder() + .with_data_file_concurrency_limit(CONCURRENCY) + .with_bloom_filter_enabled(bound.is_some()) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(METADATA_SIZE_HINT) + .build(), + ) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as FileScanTaskStream; diff --git a/crates/iceberg/src/arrow/reader/projection.rs b/crates/iceberg/src/arrow/reader/projection.rs index 3474b21624..054489ec6b 100644 --- a/crates/iceberg/src/arrow/reader/projection.rs +++ b/crates/iceberg/src/arrow/reader/projection.rs @@ -588,7 +588,7 @@ mod tests { use parquet::schema::types::SchemaDescriptor; use tempfile::TempDir; - use crate::arrow::{ArrowReader, ArrowReaderBuilder}; + use crate::arrow::{ArrowReader, ArrowReaderBuilder, ArrowReaderOptions as ReaderOptions}; use crate::expr::{Bind, Reference}; use crate::io::FileIO; use crate::scan::{FileScanTask, FileScanTaskStream}; @@ -1127,8 +1127,12 @@ message schema { let predicate = Reference::new("name").equal_to(Datum::string("Alice")); let reader = ArrowReaderBuilder::new(file_io, Runtime::current()) - .with_row_group_filtering_enabled(true) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(true) + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter( @@ -1965,8 +1969,12 @@ message schema { // Enable both row_group_filtering and row_selection - triggered the panic let reader = ArrowReaderBuilder::new(file_io, Runtime::current()) - .with_row_group_filtering_enabled(true) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(true) + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter( @@ -2320,8 +2328,12 @@ message schema { let predicate = Reference::new("id").greater_than(Datum::int(1)); let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_group_filtering_enabled(true) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(true) + .with_row_selection_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter( diff --git a/crates/iceberg/src/arrow/reader/row_filter.rs b/crates/iceberg/src/arrow/reader/row_filter.rs index 331e7f9656..7e9746adbc 100644 --- a/crates/iceberg/src/arrow/reader/row_filter.rs +++ b/crates/iceberg/src/arrow/reader/row_filter.rs @@ -242,7 +242,9 @@ mod tests { use tempfile::TempDir; use crate::Runtime; - use crate::arrow::{ArrowReader, ArrowReaderBuilder}; + use crate::arrow::{ + ArrowReader, ArrowReaderBuilder, ArrowReaderOptions as ReaderOptions, ParquetReadOptions, + }; use crate::expr::{Bind, BoundPredicate, Predicate, Reference}; use crate::io::FileIO; use crate::scan::{FileScanTask, FileScanTaskDeleteFile, FileScanTaskStream}; @@ -686,7 +688,11 @@ mod tests { .unwrap(); let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_group_filtering_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(true) + .build(), + ) .build(); let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream; let result = reader @@ -1013,7 +1019,11 @@ mod tests { ); let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_group_filtering_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(true) + .build(), + ) .build(); let task = FileScanTask::builder() @@ -1093,8 +1103,12 @@ mod tests { assert!(matches!(predicate, BoundPredicate::Not(_))); let reader = ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::current()) - .with_row_group_filtering_enabled(false) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_group_filtering_enabled(false) + .with_row_selection_enabled(true) + .build(), + ) .build(); let task = FileScanTask::builder() @@ -1380,7 +1394,11 @@ mod tests { let file_io = FileIO::new_with_fs(); let reader = ArrowReaderBuilder::new(file_io.clone(), Runtime::current()) // Enabling row selection () - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .build(); let file_size = std::fs::metadata(&file_path).unwrap().len(); @@ -1468,7 +1486,11 @@ mod tests { .unwrap(); let reader_sub2 = ArrowReaderBuilder::new(file_io.clone(), Runtime::current()) - .with_row_selection_enabled(true) + .with_options( + ReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .build(); let task_sub2 = FileScanTask::builder() @@ -1582,10 +1604,18 @@ mod tests { ) -> (Vec, u64) { let file_io = FileIO::new_with_fs(); let reader = ArrowReaderBuilder::new(file_io, Runtime::current()) - .with_bloom_filter_enabled(bloom_enabled) + .with_options( + ReaderOptions::builder() + .with_bloom_filter_enabled(bloom_enabled) + .build(), + ) // Keep the fixed footer prefetch small so the byte measurement reflects // row group I/O rather than a constant metadata read. - .with_metadata_size_hint(8 * 1024) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(8 * 1024) + .build(), + ) .build(); let task = FileScanTask::builder() diff --git a/crates/iceberg/src/cow_rewrite/mod.rs b/crates/iceberg/src/cow_rewrite/mod.rs index 25d5fb4b8f..fa567834bd 100644 --- a/crates/iceberg/src/cow_rewrite/mod.rs +++ b/crates/iceberg/src/cow_rewrite/mod.rs @@ -80,6 +80,7 @@ use futures::TryStreamExt; pub use plan::CowRewriteFile; pub use rewriter::{CowBatchRewrite, CowBatchRewriter}; +use crate::arrow::ArrowReaderOptions; use crate::expr::Predicate; use crate::scan::FileScanTaskStream; use crate::spec::{DataFile, PartitionKey}; @@ -273,7 +274,11 @@ impl<'a> CowRewriteBuilder<'a> { // candidates are fetched once per file. let mut reader_builder = self.table.reader_builder(); if let Some(batch_size) = self.batch_size { - reader_builder = reader_builder.with_batch_size(batch_size); + reader_builder = reader_builder.with_options( + ArrowReaderOptions::builder() + .with_batch_size(batch_size) + .build(), + ); } let mut batches = reader_builder.build().read(tasks)?.stream(); diff --git a/crates/iceberg/src/lib.rs b/crates/iceberg/src/lib.rs index 72c2acd3e6..17fbfb0e38 100644 --- a/crates/iceberg/src/lib.rs +++ b/crates/iceberg/src/lib.rs @@ -51,7 +51,7 @@ //! .scan() //! .select(["name", "id"]) //! .build()? -//! .to_arrow() +//! .to_arrow(Default::default()) //! .await?; //! //! // Consume this stream like arrow record batch stream. diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 005b5ea577..bc103f11c9 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -33,7 +33,7 @@ use futures::stream::BoxStream; use futures::{SinkExt, StreamExt, TryStreamExt}; pub use task::*; -use crate::arrow::ArrowReaderBuilder; +use crate::arrow::{ArrowReaderBuilder, ArrowReaderOptions}; pub use crate::arrow::{ScanMetrics, ScanResult}; use crate::delete_file_index::DeleteFileIndex; use crate::error::invalid_data; @@ -143,15 +143,10 @@ pub struct TableScanBuilder<'a> { // Defaults to none which means select all columns column_names: Option>, snapshot_selection: Option, - batch_size: Option, case_sensitive: bool, filter: Option, - concurrency_limit_data_files: usize, concurrency_limit_manifest_entries: usize, concurrency_limit_manifest_files: usize, - row_group_filtering_enabled: bool, - row_selection_enabled: bool, - bloom_filter_enabled: bool, } impl<'a> TableScanBuilder<'a> { @@ -162,25 +157,13 @@ impl<'a> TableScanBuilder<'a> { table, column_names: None, snapshot_selection: None, - batch_size: None, case_sensitive: true, filter: None, - concurrency_limit_data_files: num_cpus, concurrency_limit_manifest_entries: num_cpus, concurrency_limit_manifest_files: num_cpus, - row_group_filtering_enabled: true, - row_selection_enabled: false, - bloom_filter_enabled: false, } } - /// Sets the desired size of batches in the response - /// to something other than the default - pub fn with_batch_size(mut self, batch_size: Option) -> Self { - self.batch_size = batch_size; - self - } - /// Sets the scan's case sensitivity pub fn with_case_sensitive(mut self, case_sensitive: bool) -> Self { self.case_sensitive = case_sensitive; @@ -254,18 +237,10 @@ impl<'a> TableScanBuilder<'a> { self } - /// Sets the concurrency limit for both manifest files and manifest - /// entries for this scan + /// Sets manifest file and manifest entry concurrency for file planning. pub fn with_concurrency_limit(mut self, limit: usize) -> Self { self.concurrency_limit_manifest_files = limit; self.concurrency_limit_manifest_entries = limit; - self.concurrency_limit_data_files = limit; - self - } - - /// Sets the data file concurrency limit for this scan - pub fn with_data_file_concurrency_limit(mut self, limit: usize) -> Self { - self.concurrency_limit_data_files = limit; self } @@ -275,52 +250,6 @@ impl<'a> TableScanBuilder<'a> { self } - /// Determines whether to enable row group filtering. - /// When enabled, if a read is performed with a filter predicate, - /// then the metadata for each row group in the parquet file is - /// evaluated against the filter predicate and row groups - /// that cant contain matching rows will be skipped entirely. - /// - /// Defaults to enabled, as it generally improves performance or - /// keeps it the same, with performance degradation unlikely. - pub fn with_row_group_filtering_enabled(mut self, row_group_filtering_enabled: bool) -> Self { - self.row_group_filtering_enabled = row_group_filtering_enabled; - self - } - - /// Determines whether to enable row selection. - /// When enabled, if a read is performed with a filter predicate, - /// then (for row groups that have not been skipped) the page index - /// for each row group in a parquet file is parsed and evaluated - /// against the filter predicate to determine if ranges of rows - /// within a row group can be skipped, based upon the page-level - /// statistics for each column. - /// - /// Defaults to being disabled. Enabling requires parsing the parquet page - /// index, which can be slow enough that parsing the page index outweighs any - /// gains from the reduced number of rows that need scanning. - /// It is recommended to experiment with partitioning, sorting, row group size, - /// page size, and page row limit Iceberg settings on the table being scanned in - /// order to get the best performance from using row selection. - pub fn with_row_selection_enabled(mut self, row_selection_enabled: bool) -> Self { - self.row_selection_enabled = row_selection_enabled; - self - } - - /// Determines whether to enable bloom filter-based row group filtering. - /// - /// When enabled, if a read is performed with an equality or IN predicate, - /// the bloom filter for relevant columns in each row group is read and - /// checked. Row groups where the bloom filter proves the value is absent - /// are skipped entirely. - /// - /// Defaults to disabled, as reading bloom filters requires additional I/O - /// per column per row group. - pub fn with_bloom_filter_enabled(mut self, bloom_filter_enabled: bool) -> Self { - self.bloom_filter_enabled = bloom_filter_enabled; - self - } - /// Build the table scan. pub fn build(self) -> Result { let selected_snapshot = match self.snapshot_selection { @@ -350,16 +279,11 @@ impl<'a> TableScanBuilder<'a> { None => { let Some(current_snapshot_id) = self.table.metadata().current_snapshot() else { return Ok(TableScan { - batch_size: self.batch_size, column_names: self.column_names, file_io: self.table.file_io().clone(), plan_context: None, - concurrency_limit_data_files: self.concurrency_limit_data_files, concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries, concurrency_limit_manifest_files: self.concurrency_limit_manifest_files, - row_group_filtering_enabled: self.row_group_filtering_enabled, - row_selection_enabled: self.row_selection_enabled, - bloom_filter_enabled: self.bloom_filter_enabled, runtime: self.table.runtime().clone(), }); }; @@ -414,16 +338,11 @@ impl<'a> TableScanBuilder<'a> { }; Ok(TableScan { - batch_size: self.batch_size, column_names: self.column_names, file_io: self.table.file_io().clone(), plan_context: Some(plan_context), - concurrency_limit_data_files: self.concurrency_limit_data_files, concurrency_limit_manifest_entries: self.concurrency_limit_manifest_entries, concurrency_limit_manifest_files: self.concurrency_limit_manifest_files, - row_group_filtering_enabled: self.row_group_filtering_enabled, - row_selection_enabled: self.row_selection_enabled, - bloom_filter_enabled: self.bloom_filter_enabled, runtime: self.table.runtime().clone(), }) } @@ -436,7 +355,6 @@ pub struct TableScan { /// /// If this is None, then the scan contains no rows. plan_context: Option, - batch_size: Option, file_io: FileIO, column_names: Option>, /// The maximum number of manifest files that will be @@ -447,14 +365,6 @@ pub struct TableScan { /// be processed in parallel concurrency_limit_manifest_entries: usize, - /// The maximum number of [`ManifestEntry`]s that will - /// be processed in parallel - concurrency_limit_data_files: usize, - - row_group_filtering_enabled: bool, - row_selection_enabled: bool, - bloom_filter_enabled: bool, - runtime: Runtime, } @@ -600,20 +510,44 @@ impl TableScan { Ok(planned_file_rx.boxed()) } - /// Returns an [`ArrowRecordBatchStream`]. - pub async fn to_arrow(&self) -> Result { - let mut arrow_reader_builder = - ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone()) - .with_data_file_concurrency_limit(self.concurrency_limit_data_files) - .with_row_group_filtering_enabled(self.row_group_filtering_enabled) - .with_row_selection_enabled(self.row_selection_enabled) - .with_bloom_filter_enabled(self.bloom_filter_enabled); - - if let Some(batch_size) = self.batch_size { - arrow_reader_builder = arrow_reader_builder.with_batch_size(batch_size); - } - - arrow_reader_builder + /// Returns an [`ArrowRecordBatchStream`] using the supplied reader settings. + /// + /// Reader settings are independent of file planning. Unset data-file concurrency + /// uses execution-host parallelism; zero is treated as one. + /// [`TableScanBuilder::with_concurrency_limit`] only controls manifest planning. + /// + /// ```no_run + /// # use iceberg::{Result, arrow::{ArrowReaderOptions, ParquetReadOptions}, table::Table}; + /// # async fn read(table: &Table) -> Result<()> { + /// let options = ArrowReaderOptions::builder() + /// .with_batch_size(1024) + /// .with_row_selection_enabled(true) + /// .with_parquet_options( + /// ParquetReadOptions::builder() + /// .with_metadata_size_hint(512 * 1024) + /// .build(), + /// ) + /// .build(); + /// let stream = table.scan().build()?.to_arrow(options).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// When upgrading, move reader-setting calls from `TableScanBuilder` to + /// `ArrowReaderOptions::builder()` and pass the result to `to_arrow(options)`. + /// Use `to_arrow(Default::default())` for default settings. To preserve a + /// data-file concurrency bound previously set through the scan, set + /// `with_data_file_concurrency_limit(limit)` on the reader options explicitly. + /// + /// Locally constructed readers accept the same options through + /// [`ArrowReaderBuilder::with_options`], including Parquet I/O tuning. + /// This replaces the complete configuration. Configure individual settings + /// through the options builders; per-field setters on `ArrowReaderBuilder` + /// have been removed. [`ArrowReaderBuilder::with_parquet_options`] replaces + /// only the nested Parquet configuration. + pub async fn to_arrow(&self, options: ArrowReaderOptions) -> Result { + ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone()) + .with_options(options) .build() .read(self.plan_files().await?) .map(|result| result.stream()) @@ -750,7 +684,7 @@ mod tests { use uuid::Uuid; use super::resolve_field_id; - use crate::arrow::ArrowReaderBuilder; + use crate::arrow::{ArrowReaderBuilder, ArrowReaderOptions, ParquetReadOptions}; use crate::expr::{BoundPredicate, Reference}; use crate::io::FileIO; use crate::metadata_columns::{ @@ -771,7 +705,7 @@ mod tests { }; use crate::test_utils::test_runtime; use crate::util::snapshot::snapshot_id_as_of_time; - use crate::{ErrorKind, TableIdent}; + use crate::{ErrorKind, Runtime, TableIdent}; #[tokio::test] async fn test_table_scan_columns() { @@ -892,7 +826,7 @@ mod tests { .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -948,7 +882,6 @@ mod tests { let table_scan = table .scan() .snapshot_id(3051729675574597004) - .with_row_selection_enabled(true) .build() .unwrap(); assert_eq!( @@ -1289,7 +1222,13 @@ mod tests { .unwrap(); tasks.sort_by_key(|task| task.data_file_path().to_string()); planned_files.push(tasks); - let batches: Vec<_> = scan.to_arrow().await.unwrap().try_collect().await.unwrap(); + let batches: Vec<_> = scan + .to_arrow(ArrowReaderOptions::default()) + .await + .unwrap() + .try_collect() + .await + .unwrap(); assert!(!batches.is_empty()); schemas.push(batches[0].schema()); let mut values: Vec = batches @@ -1563,7 +1502,13 @@ mod tests { #[tokio::test] async fn test_plan_files_on_table_without_any_snapshots() { let table = TableTestFixture::new_empty().table; - let batch_stream = table.scan().build().unwrap().to_arrow().await.unwrap(); + let batch_stream = table + .scan() + .build() + .unwrap() + .to_arrow(ArrowReaderOptions::default()) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert!(batches.is_empty()); } @@ -1574,12 +1519,7 @@ mod tests { fixture.setup_manifest_files().await; // Create table scan for current snapshot and plan files - let table_scan = fixture - .table - .scan() - .with_row_selection_enabled(true) - .build() - .unwrap(); + let table_scan = fixture.table.scan().build().unwrap(); let mut tasks = table_scan .plan_files() @@ -1936,7 +1876,7 @@ mod tests { .build() .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -1961,14 +1901,16 @@ mod tests { fixture.setup_manifest_files().await; // Create table scan for current snapshot and plan files - let table_scan = fixture - .table - .scan() - .with_row_selection_enabled(true) - .build() - .unwrap(); + let table_scan = fixture.table.scan().build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -1979,17 +1921,82 @@ mod tests { } #[tokio::test] - async fn test_open_parquet_no_deletions_by_separate_reader() { + async fn reader_options_preserve_reads_after_dropping_scan() { let mut fixture = TableTestFixture::new(); - fixture.setup_manifest_files().await; - - // Create table scan for current snapshot and plan files - let table_scan = fixture + fixture.setup_multi_row_group_manifest(&[150, 299]).await; + let scan = fixture .table .scan() - .with_row_selection_enabled(true) + .select(["y"]) + .with_filter(Reference::new("y").greater_than(Datum::long(1100))) .build() .unwrap(); + let options = ArrowReaderOptions::builder() + .with_batch_size(17) + .with_data_file_concurrency_limit(1) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(4096) + .with_range_fetch_concurrency(2) + .build(), + ) + .with_row_selection_enabled(true) + .with_bloom_filter_enabled(true) + .build(); + let expected: Vec<_> = scan + .to_arrow(options.clone()) + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let tasks: Vec<_> = scan + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let serialized = serde_json::to_value(&options).unwrap(); + assert_eq!( + serialized, + serde_json::json!({ + "batch_size": 17, "data_file_concurrency_limit": 1, + "row_group_filtering_enabled": true, "row_selection_enabled": true, + "bloom_filter_enabled": true, + "parquet_options": { + "metadata_size_hint": 4096, "range_coalesce_bytes": 1048576, "range_fetch_concurrency": 2, + "preload_column_index": true, "preload_offset_index": true, "preload_page_index": false, + }, + }) + ); + let deserialized: ArrowReaderOptions = serde_json::from_value(serialized).unwrap(); + assert_eq!(deserialized, options); + let options = deserialized; + drop(scan); + let reader = + ArrowReaderBuilder::new(FileIO::new_with_fs(), Runtime::try_current().unwrap()) + .with_options(options) + .build(); + let actual: Vec<_> = reader + .read(Box::pin(stream::iter(tasks.into_iter().map(Ok)))) + .unwrap() + .stream() + .try_collect() + .await + .unwrap(); + assert!(!actual.is_empty()); + assert!(actual.iter().all(|batch| batch.num_rows() <= 17)); + assert_eq!(actual, expected); + } + + #[tokio::test] + async fn test_open_parquet_no_deletions_by_separate_reader() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + // Create table scan for current snapshot and plan files + let table_scan = fixture.table.scan().build().unwrap(); let mut plan_task: Vec<_> = table_scan .plan_files() @@ -2032,15 +2039,16 @@ mod tests { fixture.setup_manifest_files().await; // Create table scan for current snapshot and plan files - let table_scan = fixture - .table - .scan() - .select(["x", "z"]) - .with_row_selection_enabled(true) - .build() - .unwrap(); + let table_scan = fixture.table.scan().select(["x", "z"]).build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2056,7 +2064,14 @@ mod tests { // test empty scan let table_scan = fixture.table.scan().select_empty().build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_columns(), 0); @@ -2071,12 +2086,17 @@ mod tests { // Filter: y < 3 let mut builder = fixture.table.scan(); let predicate = Reference::new("y").less_than(Datum::long(3)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2099,12 +2119,17 @@ mod tests { // Filter: y >= 5 let mut builder = fixture.table.scan(); let predicate = Reference::new("y").greater_than_or_equal_to(Datum::long(5)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2127,12 +2152,17 @@ mod tests { // Filter: dbl == 150.0 let mut builder = fixture.table.scan(); let predicate = Reference::new("dbl").equal_to(Datum::double(150.0f64)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2152,12 +2182,17 @@ mod tests { // Filter: i32 == 150 let mut builder = fixture.table.scan(); let predicate = Reference::new("i32").equal_to(Datum::int(150i32)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2177,12 +2212,17 @@ mod tests { // Filter: i64 == 150 let mut builder = fixture.table.scan(); let predicate = Reference::new("i64").equal_to(Datum::long(150i64)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2202,12 +2242,17 @@ mod tests { // Filter: bool == true let mut builder = fixture.table.scan(); let predicate = Reference::new("bool").equal_to(Datum::bool(true)); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2227,12 +2272,17 @@ mod tests { // Filter: y is null let mut builder = fixture.table.scan(); let predicate = Reference::new("y").is_null(); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches.len(), 0); @@ -2246,12 +2296,17 @@ mod tests { // Filter: y is not null let mut builder = fixture.table.scan(); let predicate = Reference::new("y").is_not_null(); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_rows(), 1024); @@ -2267,12 +2322,17 @@ mod tests { let predicate = Reference::new("y") .less_than(Datum::long(5)) .and(Reference::new("z").greater_than_or_equal_to(Datum::long(4))); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_rows(), 500); @@ -2303,12 +2363,17 @@ mod tests { let predicate = Reference::new("y") .less_than(Datum::long(5)) .or(Reference::new("z").greater_than_or_equal_to(Datum::long(4))); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_rows(), 1024); @@ -2340,12 +2405,17 @@ mod tests { // Filter: a STARTSWITH "Ice" let mut builder = fixture.table.scan(); let predicate = Reference::new("a").starts_with(Datum::string("Ice")); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2364,12 +2434,17 @@ mod tests { // Filter: a NOT STARTSWITH "Ice" let mut builder = fixture.table.scan(); let predicate = Reference::new("a").not_starts_with(Datum::string("Ice")); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2389,12 +2464,17 @@ mod tests { let mut builder = fixture.table.scan(); let predicate = Reference::new("a").is_in([Datum::string("Sioux"), Datum::string("Iceberg")]); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2414,12 +2494,17 @@ mod tests { let mut builder = fixture.table.scan(); let predicate = Reference::new("a").is_not_in([Datum::string("Sioux"), Datum::string("Iceberg")]); - builder = builder - .with_filter(predicate) - .with_row_selection_enabled(true); + builder = builder.with_filter(predicate); let table_scan = builder.build().unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -2684,11 +2769,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_FILE]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Verify we have 2 columns: x and _file @@ -2743,11 +2834,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_FILE, "z"]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_columns(), 3); @@ -2774,11 +2871,17 @@ mod tests { .table .scan() .select([RESERVED_COL_NAME_FILE]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Should have exactly 1 column @@ -2807,11 +2910,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_FILE]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Collect all unique file paths from the batches @@ -2852,11 +2961,17 @@ mod tests { .table .scan() .select([RESERVED_COL_NAME_FILE, "x", "y"]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_columns(), 3); @@ -2878,11 +2993,17 @@ mod tests { .table .scan() .select(["x", "y", RESERVED_COL_NAME_FILE]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches[0].num_columns(), 3); @@ -2912,11 +3033,17 @@ mod tests { RESERVED_COL_NAME_FILE, // _file repeated "y", // y repeated ]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Verify we have exactly 6 columns (duplicates are allowed and preserved) @@ -3146,11 +3273,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Verify we have 3 columns: x, _spec_id, and z @@ -3208,12 +3341,15 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER]) - .with_row_selection_enabled(true) .build() .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .await .unwrap() .try_collect() @@ -3247,12 +3383,15 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER]) - .with_row_selection_enabled(true) .build() .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) .await .unwrap() .try_collect() @@ -3272,11 +3411,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_SPEC_ID]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Verify we have 2 columns: x and _spec_id @@ -3325,11 +3470,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_SPEC_ID, "z"]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); // Verify the x column exists and has correct data @@ -3380,11 +3531,17 @@ mod tests { .table .scan() .select(["x", RESERVED_COL_NAME_POS, RESERVED_COL_NAME_FILE]) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches.len(), 2); @@ -3430,11 +3587,17 @@ mod tests { .scan() .select([RESERVED_COL_NAME_POS, "x", "y"]) .with_filter(predicate) - .with_row_selection_enabled(true) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches.len(), 2); @@ -3470,12 +3633,18 @@ mod tests { .table .scan() .select([RESERVED_COL_NAME_POS, "a", RESERVED_COL_NAME_POS, "x"]) - .with_row_selection_enabled(true) .with_filter(predicate) .build() .unwrap(); - let batch_stream = table_scan.to_arrow().await.unwrap(); + let batch_stream = table_scan + .to_arrow( + ArrowReaderOptions::builder() + .with_row_selection_enabled(true) + .build(), + ) + .await + .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); assert_eq!(batches.len(), 2); @@ -3543,7 +3712,7 @@ mod tests { .select(["x", RESERVED_COL_NAME_POS]) .build() .unwrap() - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -3617,7 +3786,7 @@ mod tests { .select(["x", RESERVED_COL_NAME_POS]) .build() .unwrap() - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -3669,10 +3838,9 @@ mod tests { .scan() .select(["y", RESERVED_COL_NAME_POS]) .with_filter(predicate) - .with_row_group_filtering_enabled(true) .build() .unwrap() - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() diff --git a/crates/integration_tests/tests/conflict_commit_test.rs b/crates/integration_tests/tests/conflict_commit_test.rs index 57153d44f0..73da9736ef 100644 --- a/crates/integration_tests/tests/conflict_commit_test.rs +++ b/crates/integration_tests/tests/conflict_commit_test.rs @@ -119,7 +119,7 @@ async fn test_append_data_file_conflict() { .select_all() .build() .unwrap() - .to_arrow() + .to_arrow(Default::default()) .await .unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); diff --git a/crates/integration_tests/tests/read_evolved_schema.rs b/crates/integration_tests/tests/read_evolved_schema.rs index f7416be2d4..844ad81136 100644 --- a/crates/integration_tests/tests/read_evolved_schema.rs +++ b/crates/integration_tests/tests/read_evolved_schema.rs @@ -48,7 +48,7 @@ async fn test_evolved_schema() { let predicate = Reference::new("language").not_equal_to(Datum::string("Rust")); let scan = table.scan().with_filter(predicate).build(); - let batch_stream = scan.unwrap().to_arrow().await.unwrap(); + let batch_stream = scan.unwrap().to_arrow(Default::default()).await.unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); @@ -84,7 +84,7 @@ async fn test_evolved_schema() { // Reference::new("foo").not_equal_to(Datum::int(22)); let scan = table.scan().build(); - let batch_stream = scan.unwrap().to_arrow().await.unwrap(); + let batch_stream = scan.unwrap().to_arrow(Default::default()).await.unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); let mut actual = vec![ @@ -114,7 +114,7 @@ async fn test_evolved_schema() { .await .unwrap(); let scan = table.scan().build(); - let batch_stream = scan.unwrap().to_arrow().await.unwrap(); + let batch_stream = scan.unwrap().to_arrow(Default::default()).await.unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); let mut actual_foo = vec![ diff --git a/crates/integration_tests/tests/read_positional_deletes.rs b/crates/integration_tests/tests/read_positional_deletes.rs index fe8c9830d4..64d07f86fc 100644 --- a/crates/integration_tests/tests/read_positional_deletes.rs +++ b/crates/integration_tests/tests/read_positional_deletes.rs @@ -61,7 +61,7 @@ async fn test_read_table_with_positional_deletes() { assert_eq!(plan[0].deletes().len(), 1); // we should see two rows deleted, returning 10 rows instead of 12 - let batch_stream = scan.to_arrow().await.unwrap(); + let batch_stream = scan.to_arrow(Default::default()).await.unwrap(); let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); let num_rows: usize = batches.iter().map(|v| v.num_rows()).sum(); assert_eq!(num_rows, 10); diff --git a/crates/storage/opendal/README.md b/crates/storage/opendal/README.md index ac3a8b62ef..87ebdbe171 100644 --- a/crates/storage/opendal/README.md +++ b/crates/storage/opendal/README.md @@ -80,7 +80,7 @@ async fn main() -> iceberg::Result<()> { .await?; let scan = table.scan().select_all().build()?; - let stream = scan.to_arrow().await?; + let stream = scan.to_arrow(Default::default()).await?; Ok(()) }