diff --git a/crates/catalog/s3tables/src/catalog.rs b/crates/catalog/s3tables/src/catalog.rs index 73127bd903..f0ace93a02 100644 --- a/crates/catalog/s3tables/src/catalog.rs +++ b/crates/catalog/s3tables/src/catalog.rs @@ -1329,7 +1329,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/README.md b/crates/iceberg/README.md index 08ce82bd2f..f6688c0a52 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. @@ -65,6 +65,36 @@ async fn main() -> Result<()> { } ``` +### Reader settings + +Pass reader settings separately from file planning: + +```rust +let options = iceberg::arrow::ArrowReaderOptions::builder() + .with_batch_size(Some(1024)) + .with_row_selection_enabled(true) + .with_parquet_options( + iceberg::arrow::ParquetReadOptions::builder() + .with_metadata_size_hint(Some(512 * 1024)) + .build(), + ) + .build(); +let stream = table.scan().build()?.to_arrow(options).await?; +``` + +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. Scan +`with_concurrency_limit()` controls manifest planning only; set +`with_data_file_concurrency_limit(Some(limit))` on the reader options to also +limit data-file reads. An unset data-file limit uses execution-host parallelism. + +Locally constructed readers accept the same options through +`ArrowReaderBuilder::with_options(options)`, including Parquet I/O tuning. +This replaces the complete configuration; subsequent convenience setters override +individual settings. `with_parquet_options()` replaces only the nested Parquet +configuration. + ## Storage Backends For storage backend support (S3, GCS, local filesystem, etc.), use the [`iceberg-storage-opendal`](https://crates.io/crates/iceberg-storage-opendal) crate. See its [README](../storage/opendal/README.md) for available backends and feature flags. diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 5f87d98810..2989c2a4ec 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -75,10 +75,45 @@ pub fn iceberg::arrow::ArrowReaderBuilder::with_batch_size(self, batch_size: usi 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_options(self, options: iceberg::arrow::ArrowReaderOptions) -> Self +pub fn iceberg::arrow::ArrowReaderBuilder::with_parquet_options(self, options: iceberg::arrow::ParquetReadOptions) -> 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 struct iceberg::arrow::ArrowReaderOptions +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 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::Copy for iceberg::arrow::ParquetReadOptions +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 +1472,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 +1483,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/mod.rs b/crates/iceberg/src/arrow/reader/mod.rs index 9d9c8be48e..7d56ae9760 100644 --- a/crates/iceberg/src/arrow/reader/mod.rs +++ b/crates/iceberg/src/arrow/reader/mod.rs @@ -19,8 +19,6 @@ use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; use crate::io::FileIO; -use crate::runtime::Runtime; -use crate::util::available_parallelism; /// Default gap between byte ranges below which they are coalesced into a /// single request. Matches object_store's `OBJECT_STORE_COALESCE_DEFAULT`. @@ -36,6 +34,7 @@ const DEFAULT_METADATA_SIZE_HINT: usize = 512 * 1024; mod file_reader; mod options; +mod parquet_options; mod pipeline; mod positional_deletes; mod predicate_visitor; @@ -43,127 +42,14 @@ mod projection; mod row_filter; mod row_lineage; pub use file_reader::ArrowFileReader; -pub(crate) use options::ParquetReadOptions; +pub use options::{ArrowReaderBuilder, ArrowReaderOptions, ArrowReaderOptionsBuilder}; +pub use parquet_options::{ParquetReadOptions, ParquetReadOptionsBuilder}; use predicate_visitor::{CollectFieldIdVisitor, PredicateConverter}; use projection::{ add_fallback_field_ids_to_arrow_schema, apply_name_mapping_to_arrow_schema, find_leaf_by_field_id, }; -/// Builder to create ArrowReader -pub struct ArrowReaderBuilder { - batch_size: Option, - 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, - 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; - 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; - self - } - - /// Build the ArrowReader. - pub fn build(self) -> ArrowReader { - ArrowReader { - batch_size: self.batch_size, - file_io: self.file_io.clone(), - delete_file_loader: CachingDeleteFileLoader::new( - self.file_io.clone(), - self.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, - } - } -} - /// Reads data from Parquet files #[derive(Clone)] pub struct ArrowReader { diff --git a/crates/iceberg/src/arrow/reader/options.rs b/crates/iceberg/src/arrow/reader/options.rs index ae6a3ed18e..46f6ced9f6 100644 --- a/crates/iceberg/src/arrow/reader/options.rs +++ b/crates/iceberg/src/arrow/reader/options.rs @@ -15,70 +15,302 @@ // 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, -}; +use super::{ArrowReader, ParquetReadOptions}; +use crate::arrow::caching_delete_file_loader::CachingDeleteFileLoader; +use crate::io::FileIO; +use crate::runtime::Runtime; +use crate::util::available_parallelism; + +/// Configures Arrow reads independently of file planning and local I/O resources. +/// +/// ```no_run +/// # use iceberg::{Result, arrow::{ArrowReaderOptions, ParquetReadOptions}, table::Table}; +/// # async fn read(table: &Table) -> Result<()> { +/// let options = ArrowReaderOptions::builder() +/// .with_batch_size(Some(1024)) +/// .with_parquet_options( +/// ParquetReadOptions::builder() +/// .with_metadata_size_hint(Some(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. + batch_size: Option, + /// Maximum in-flight data files, or execution-host parallelism when unset. + 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() + } +} + +/// Builder to create ArrowReader +pub struct ArrowReaderBuilder { + options: ArrowReaderOptions, + file_io: FileIO, + runtime: Runtime, +} + +impl ArrowReaderBuilder { + /// Create a new ArrowReaderBuilder + pub fn new(file_io: FileIO, runtime: Runtime) -> Self { + ArrowReaderBuilder { + options: ArrowReaderOptions::default(), + file_io, + 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.options.data_file_concurrency_limit = Some(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.options.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.options.row_group_filtering_enabled = row_group_filtering_enabled; + self + } + + /// Determines whether to enable row selection. + pub fn with_row_selection_enabled(mut self, row_selection_enabled: bool) -> Self { + self.options.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.options.bloom_filter_enabled = bloom_filter_enabled; + self + } -/// Options for tuning Parquet file I/O. -#[derive(Clone, Copy, Debug, TypedBuilder)] -#[builder(field_defaults(setter(prefix = "with_")))] -pub(crate) struct ParquetReadOptions { - /// Number of bytes to prefetch for parsing the Parquet metadata. + /// 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). - /// - /// 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, - /// Gap threshold for merging nearby byte ranges into a single request. + pub fn with_metadata_size_hint(mut self, metadata_size_hint: usize) -> Self { + self.options.parquet_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`. - #[builder(default = DEFAULT_RANGE_COALESCE_BYTES)] - pub(crate) range_coalesce_bytes: u64, - /// Maximum number of merged byte ranges to fetch concurrently. + /// 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.options.parquet_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`. - #[builder(default = DEFAULT_RANGE_FETCH_CONCURRENCY)] - pub(crate) range_fetch_concurrency: usize, - /// Whether to preload the column index when reading Parquet metadata. - #[builder(default = true)] - pub(crate) preload_column_index: bool, - /// Whether to preload the offset index when reading Parquet metadata. - #[builder(default = true)] - pub(crate) preload_offset_index: bool, - /// Whether to preload the page index when reading Parquet metadata. - #[builder(default = false)] - pub(crate) preload_page_index: bool, -} + /// 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.options.parquet_options.range_fetch_concurrency = range_fetch_concurrency; + self + } -impl ParquetReadOptions { - pub(crate) fn metadata_size_hint(&self) -> Option { - self.metadata_size_hint + /// 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 } - pub(crate) fn range_coalesce_bytes(&self) -> u64 { - self.range_coalesce_bytes + /// Replaces Parquet I/O settings while preserving other reader settings. + pub fn with_parquet_options(mut self, options: ParquetReadOptions) -> Self { + self.options.parquet_options = options; + self } - pub(crate) fn range_fetch_concurrency(&self) -> usize { - self.range_fetch_concurrency + /// Build the ArrowReader. + pub fn build(self) -> ArrowReader { + let concurrency_limit_data_files = self + .options + .data_file_concurrency_limit + .unwrap_or_else(|| available_parallelism().get()); + ArrowReader { + batch_size: self.options.batch_size, + file_io: self.file_io.clone(), + delete_file_loader: CachingDeleteFileLoader::new( + self.file_io.clone(), + concurrency_limit_data_files, + self.runtime.clone(), + ), + concurrency_limit_data_files, + 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, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::io::FileIO; + use crate::test_utils::test_runtime; + use crate::util::available_parallelism; + + #[test] + fn missing_options_use_defaults_and_unknown_options_are_ignored() { + let decoded: ArrowReaderOptions = + serde_json::from_str(r#"{"future_option":true}"#).unwrap(); + assert_eq!(decoded, ArrowReaderOptions::default()); + let decoded: ArrowReaderOptions = serde_json::from_str(r#"{"batch_size":17}"#).unwrap(); + assert_eq!( + decoded, + ArrowReaderOptions::builder() + .with_batch_size(Some(17)) + .build() + ); + assert_eq!( + serde_json::to_value(ArrowReaderOptions::default()).unwrap()["data_file_concurrency_limit"], + serde_json::Value::Null + ); } - pub(crate) fn preload_column_index(&self) -> bool { - self.preload_column_index + #[test] + fn partial_parquet_options_preserve_defaults() { + let options: ArrowReaderOptions = serde_json::from_str( + r#"{"parquet_options":{"range_fetch_concurrency":2,"future_option":true}}"#, + ) + .unwrap(); + assert_eq!( + options, + ArrowReaderOptions::builder() + .with_parquet_options( + ParquetReadOptions::builder() + .with_range_fetch_concurrency(2) + .build() + ) + .build() + ); } - pub(crate) fn preload_offset_index(&self) -> bool { - self.preload_offset_index + #[test] + fn parquet_override_preserves_reader_settings_and_allows_individual_overrides() { + let reader = ArrowReaderBuilder::new(FileIO::new_with_memory(), test_runtime()) + .with_options( + ArrowReaderOptions::builder() + .with_batch_size(Some(17)) + .build(), + ) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(Some(4096)) + .with_range_fetch_concurrency(2) + .build(), + ) + .with_range_fetch_concurrency(3) + .build(); + assert_eq!( + (reader.batch_size, reader.parquet_read_options), + ( + Some(17), + ParquetReadOptions::builder() + .with_metadata_size_hint(Some(4096)) + .with_range_fetch_concurrency(3) + .build() + ) + ); } - pub(crate) fn preload_page_index(&self) -> bool { - self.preload_page_index + #[test] + fn options_restore_local_defaults_or_apply_explicit_settings() { + for explicit in [false, true] { + let options = if explicit { + ArrowReaderOptions::builder() + .with_batch_size(Some(17)) + .with_data_file_concurrency_limit(Some(3)) + .with_row_group_filtering_enabled(false) + .with_row_selection_enabled(true) + .with_bloom_filter_enabled(true) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(Some(4096)) + .with_range_coalesce_bytes(8192) + .with_range_fetch_concurrency(2) + .build(), + ) + .build() + } else { + ArrowReaderOptions::default() + }; + let reader = ArrowReaderBuilder::new(FileIO::new_with_memory(), test_runtime()) + .with_batch_size(1024) + .with_data_file_concurrency_limit(99) + .with_metadata_size_hint(4096) + .with_range_coalesce_bytes(8192) + .with_range_fetch_concurrency(2) + .with_options(options) + .build(); + let expected = if explicit { + (Some(17), 3, false, true, true) + } else { + (None, available_parallelism().get(), true, false, false) + }; + assert_eq!( + ( + reader.batch_size, + reader.concurrency_limit_data_files, + reader.row_group_filtering_enabled, + reader.row_selection_enabled, + reader.bloom_filter_enabled + ), + expected + ); + assert_eq!( + ( + reader.parquet_read_options.metadata_size_hint(), + reader.parquet_read_options.range_coalesce_bytes(), + reader.parquet_read_options.range_fetch_concurrency() + ), + if explicit { + (Some(4096), 8192, 2) + } else { + (Some(512 * 1024), 1024 * 1024, 10) + } + ); + } } } diff --git a/crates/iceberg/src/arrow/reader/parquet_options.rs b/crates/iceberg/src/arrow/reader/parquet_options.rs new file mode 100644 index 0000000000..05081b6061 --- /dev/null +++ b/crates/iceberg/src/arrow/reader/parquet_options.rs @@ -0,0 +1,95 @@ +// 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. + +//! 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, +}; + +/// Options for tuning Parquet file I/O. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize, TypedBuilder)] +#[serde(default)] +#[builder(field_defaults(setter(prefix = "with_")))] +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, + /// 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, + /// 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. + #[builder(default = true, setter(skip))] + #[serde(skip)] + pub(crate) preload_column_index: bool, + /// Whether to preload the offset index when reading Parquet metadata. + #[builder(default = true, setter(skip))] + #[serde(skip)] + pub(crate) preload_offset_index: bool, + /// Whether to preload the page index when reading Parquet metadata. + #[builder(default = false, setter(skip))] + #[serde(skip)] + pub(crate) preload_page_index: bool, +} + +impl Default for ParquetReadOptions { + fn default() -> Self { + Self::builder().build() + } +} + +impl ParquetReadOptions { + pub(crate) fn metadata_size_hint(&self) -> Option { + self.metadata_size_hint + } + + pub(crate) fn range_coalesce_bytes(&self) -> u64 { + self.range_coalesce_bytes + } + + pub(crate) fn range_fetch_concurrency(&self) -> usize { + self.range_fetch_concurrency + } + + pub(crate) fn preload_column_index(&self) -> bool { + self.preload_column_index + } + + pub(crate) fn preload_offset_index(&self) -> bool { + self.preload_offset_index + } + + pub(crate) fn preload_page_index(&self) -> bool { + self.preload_page_index + } +} 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 cb77abeb74..737536d433 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; @@ -259,13 +242,6 @@ impl<'a> TableScanBuilder<'a> { 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 +251,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 +280,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 +339,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 +356,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 +366,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, } @@ -601,19 +512,9 @@ impl TableScan { } /// 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 + 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()) @@ -760,7 +661,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::{ @@ -781,7 +682,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() { @@ -902,7 +803,7 @@ mod tests { .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -958,7 +859,6 @@ mod tests { let table_scan = table .scan() .snapshot_id(3051729675574597004) - .with_row_selection_enabled(true) .build() .unwrap(); assert_eq!( @@ -1299,7 +1199,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 @@ -1573,7 +1479,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()); } @@ -1584,12 +1496,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() @@ -1973,7 +1880,7 @@ mod tests { .build() .unwrap(); let batches: Vec<_> = table_scan - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -1998,14 +1905,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(); @@ -2016,17 +1925,76 @@ 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(Some(17)) + .with_data_file_concurrency_limit(Some(1)) + .with_parquet_options( + ParquetReadOptions::builder() + .with_metadata_size_hint(Some(4096)) + .with_range_fetch_concurrency(2) + .build(), + ) + .with_row_selection_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": false, + "parquet_options": { "metadata_size_hint": 4096, "range_coalesce_bytes": 1048576, "range_fetch_concurrency": 2 }, + }) + ); + let options: ArrowReaderOptions = serde_json::from_value(serialized).unwrap(); + 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() @@ -2069,15 +2037,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(); @@ -2093,7 +2062,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); @@ -2108,12 +2084,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(); @@ -2136,12 +2117,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(); @@ -2164,12 +2150,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(); @@ -2189,12 +2180,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(); @@ -2214,12 +2210,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(); @@ -2239,12 +2240,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(); @@ -2264,12 +2270,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); @@ -2283,12 +2294,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); @@ -2304,12 +2320,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); @@ -2340,12 +2361,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); @@ -2377,12 +2403,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(); @@ -2401,12 +2432,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(); @@ -2426,12 +2462,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(); @@ -2451,12 +2492,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(); @@ -2721,11 +2767,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 @@ -2780,11 +2832,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); @@ -2811,11 +2869,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 @@ -2844,11 +2908,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 @@ -2889,11 +2959,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); @@ -2915,11 +2991,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); @@ -2949,11 +3031,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) @@ -3183,11 +3271,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 @@ -3245,12 +3339,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() @@ -3284,12 +3381,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() @@ -3309,11 +3409,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 @@ -3362,11 +3468,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 @@ -3417,11 +3529,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); @@ -3467,11 +3585,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); @@ -3507,12 +3631,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); @@ -3580,7 +3710,7 @@ mod tests { .select(["x", RESERVED_COL_NAME_POS]) .build() .unwrap() - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -3654,7 +3784,7 @@ mod tests { .select(["x", RESERVED_COL_NAME_POS]) .build() .unwrap() - .to_arrow() + .to_arrow(ArrowReaderOptions::default()) .await .unwrap() .try_collect() @@ -3706,10 +3836,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(()) }