From e63e6896f1f270397cf6497095ef9133eb8ac082 Mon Sep 17 00:00:00 2001 From: t3rb3d Date: Sun, 11 Oct 2026 20:10:06 +0200 Subject: [PATCH 1/2] Enable page-index pruning for Iceberg scans --- crates/datafusion/src/read/scan.rs | 163 ++++++++++++++++++- crates/datafusion/src/test_utils/harness.rs | 171 +++++++++++++++++++- 2 files changed, 328 insertions(+), 6 deletions(-) diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index afcf3ee..c94b410 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -182,13 +182,19 @@ impl DataSource for IcebergDataSource { fn open( &self, _partition: usize, - _context: Arc, + context: Arc, ) -> Result { let fut = get_batch_stream( self.table.clone(), self.snapshot_id, Some(self.projection.clone()), self.predicate.clone(), + context + .session_config() + .options() + .execution + .parquet + .enable_page_index, ); let stream = futures::stream::once(fut).try_flatten(); @@ -274,6 +280,7 @@ async fn get_batch_stream( snapshot_id: Option, column_names: Option>, predicate: Option, + enable_page_index: bool, ) -> Result> + Send>>> { let scan_builder = match snapshot_id { Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), @@ -287,7 +294,10 @@ async fn get_batch_stream( if let Some(predicate) = predicate { scan_builder = scan_builder.with_filter(predicate); } - let table_scan = scan_builder.build().map_err(to_datafusion_error)?; + let table_scan = scan_builder + .with_row_selection_enabled(enable_page_index) + .build() + .map_err(to_datafusion_error)?; let stream = table_scan .to_arrow() @@ -296,3 +306,152 @@ async fn get_batch_stream( .map_err(to_datafusion_error); Ok(Box::pin(stream)) } + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::{ArrayRef, Decimal128Array, Int64Array}; + use datafusion::arrow::compute::concat_batches; + use datafusion::prelude::SessionContext; + use iceberg::arrow::schema_to_arrow_schema; + use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; + use parquet::basic::Compression; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + use super::*; + use crate::test_utils::IcebergTestHarness; + + const ROWS: usize = 262_144; + const QUERY: &str = "SELECT * FROM events WHERE id BETWEEN 8190 AND 8194 ORDER BY id"; + + fn page_index_fixture() -> (Schema, RecordBatch) { + let mut fields = vec![Arc::new(NestedField::required( + 1, + "id", + Type::Primitive(PrimitiveType::Long), + ))]; + let mut columns: Vec = + vec![Arc::new(Int64Array::from_iter_values(0..ROWS as i64))]; + for (index, precision) in [9, 18, 38].into_iter().enumerate() { + fields.push(Arc::new(NestedField::optional( + index as i32 + 2, + format!("amount_{precision}"), + Type::Primitive(PrimitiveType::Decimal { + precision, + scale: 2, + }), + ))); + let values = Decimal128Array::from_iter((0..ROWS).map(|row| { + (!(row / 8192).is_multiple_of(2)) + .then_some(row as i128 - ROWS as i128 / 2) + })) + .with_precision_and_scale(precision as u8, 2) + .unwrap(); + columns.push(Arc::new(values)); + } + let schema = Schema::builder().with_fields(fields).build().unwrap(); + let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let batch = RecordBatch::try_new(arrow_schema, columns).unwrap(); + (schema, batch) + } + + fn page_index_properties(enabled: bool) -> WriterProperties { + WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_dictionary_enabled(false) + .set_max_row_group_row_count(Some(ROWS)) + .set_data_page_row_count_limit(8192) + .set_statistics_enabled(if enabled { + EnabledStatistics::Page + } else { + EnabledStatistics::Chunk + }) + .set_offset_index_disabled(!enabled) + .build() + } + + async fn query(ctx: &SessionContext, sql: &str) -> Result { + let batches = ctx.sql(sql).await?.collect().await?; + Ok(concat_batches(&batches[0].schema(), &batches)?) + } + + #[tokio::test] + async fn test_page_index_setting_prunes_reads() -> Result<()> { + let (schema, batch) = page_index_fixture(); + let expected = batch.slice(8190, 5); + for catalog_backed in [false, true] { + let (harness, _directory) = IcebergTestHarness::with_parquet( + schema.clone(), + &batch, + page_index_properties(true), + catalog_backed, + ) + .await?; + let ctx = harness.context(); + let mut read_bytes = [0; 3]; + for (bytes, enabled) in + read_bytes.iter_mut().zip([None, Some(false), Some(true)]) + { + if let Some(enabled) = enabled { + ctx.sql(&format!( + "SET datafusion.execution.parquet.enable_page_index = '{enabled}'", + )).await?.collect().await?; + } + assert_eq!(query(ctx, QUERY).await?, expected); + *bytes = harness.take_parquet_read_bytes(); + } + let [default, disabled, enabled] = read_bytes; + assert_eq!(default, enabled); + assert!( + enabled < disabled, + "page index on: {enabled}, off: {disabled}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_page_index_setting_without_indexes() -> Result<()> { + let (schema, batch) = page_index_fixture(); + let expected = batch.slice(8190, 5); + for catalog_backed in [false, true] { + let (harness, _directory) = IcebergTestHarness::with_parquet( + schema.clone(), + &batch, + page_index_properties(false), + catalog_backed, + ) + .await?; + for enabled in [false, true] { + harness.context().sql(&format!( + "SET datafusion.execution.parquet.enable_page_index = '{enabled}'", + )).await?.collect().await?; + assert_eq!(query(harness.context(), QUERY).await?, expected); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_page_index_pruning_with_nullable_decimals_and_projection() -> Result<()> + { + let (schema, batch) = page_index_fixture(); + let expected = batch.slice(8190, 2).project(&[0])?; + for catalog_backed in [false, true] { + let (harness, _directory) = IcebergTestHarness::with_parquet( + schema.clone(), + &batch, + page_index_properties(true), + catalog_backed, + ) + .await?; + for precision in [9, 18, 38] { + let sql = format!( + "SELECT id FROM events WHERE id BETWEEN 8190 AND 8194 \ + AND amount_{precision} IS NULL ORDER BY id", + ); + assert_eq!(query(harness.context(), &sql).await?, expected); + } + } + Ok(()) + } +} diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index d106467..19ba965 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -16,8 +16,12 @@ // under the License. use std::collections::HashMap; +#[cfg(test)] +use std::ops::Range; use std::path::PathBuf; use std::sync::Arc; +#[cfg(test)] +use std::sync::atomic::{AtomicU64, Ordering}; use async_trait::async_trait; use bytes::Bytes; @@ -49,10 +53,15 @@ const NAMESPACE_NAME: &str = "public"; pub struct IcebergTestHarness { ctx: SessionContext, + #[cfg(test)] + parquet_read_bytes: Arc, } impl IcebergTestHarness { pub async fn new() -> Result { + let storage_factory = FixtureStorageFactory::default(); + #[cfg(test)] + let parquet_read_bytes = storage_factory.parquet_read_bytes.clone(); let state = SessionStateBuilder::new() .with_default_features() // Plan snapshots must not vary with the runner's CPU count. @@ -62,13 +71,13 @@ impl IcebergTestHarness { .with_table_factory( "ICEBERG".to_string(), Arc::new(IcebergTableProviderFactory::new_with_storage_factory( - Arc::new(FixtureStorageFactory::default()), + Arc::new(storage_factory.clone()), )), ) .build(); let catalog = MemoryCatalogBuilder::default() - .with_storage_factory(Arc::new(FixtureStorageFactory::default())) + .with_storage_factory(Arc::new(storage_factory)) .load( CATALOG_NAME, HashMap::from([( @@ -96,7 +105,11 @@ impl IcebergTestHarness { let ctx = SessionContext::new_with_state(state); ctx.register_catalog(CATALOG_NAME, Arc::new(provider)); - Ok(Self { ctx }) + Ok(Self { + ctx, + #[cfg(test)] + parquet_read_bytes, + }) } pub async fn query(&self, sql: &str) -> Result<(String, String)> { @@ -114,12 +127,17 @@ impl IcebergTestHarness { #[derive(Debug, Clone, Serialize, Deserialize)] struct FixtureStorageFactory { root: PathBuf, + #[cfg(test)] + #[serde(skip)] + parquet_read_bytes: Arc, } impl Default for FixtureStorageFactory { fn default() -> Self { Self { root: PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), + #[cfg(test)] + parquet_read_bytes: Arc::default(), } } } @@ -129,6 +147,8 @@ impl StorageFactory for FixtureStorageFactory { fn build(&self, _config: &StorageConfig) -> IcebergResult> { Ok(Arc::new(FixtureStorage { root: self.root.clone(), + #[cfg(test)] + parquet_read_bytes: self.parquet_read_bytes.clone(), })) } } @@ -136,6 +156,9 @@ impl StorageFactory for FixtureStorageFactory { #[derive(Debug, Clone, Serialize, Deserialize)] struct FixtureStorage { root: PathBuf, + #[cfg(test)] + #[serde(skip)] + parquet_read_bytes: Arc, } impl FixtureStorage { @@ -170,7 +193,15 @@ impl Storage for FixtureStorage { } async fn reader(&self, path: &str) -> IcebergResult> { - self.local().reader(&self.local_path(path)?).await + let reader = self.local().reader(&self.local_path(path)?).await?; + #[cfg(test)] + if path.ends_with(".parquet") { + return Ok(Box::new(CountingFileRead { + reader, + bytes: self.parquet_read_bytes.clone(), + })); + } + Ok(reader) } async fn write(&self, path: &str, bytes: Bytes) -> IcebergResult<()> { @@ -208,3 +239,135 @@ impl Storage for FixtureStorage { Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) } } + +#[cfg(test)] +struct CountingFileRead { + reader: Box, + bytes: Arc, +} + +#[cfg(test)] +#[async_trait] +impl FileRead for CountingFileRead { + async fn read(&self, range: Range) -> IcebergResult { + let bytes = self.reader.read(range).await?; + self.bytes.fetch_add(bytes.len() as u64, Ordering::Relaxed); + Ok(bytes) + } +} + +#[cfg(test)] +impl IcebergTestHarness { + pub(crate) fn context(&self) -> &SessionContext { + &self.ctx + } + + pub(crate) fn take_parquet_read_bytes(&self) -> u64 { + self.parquet_read_bytes.swap(0, Ordering::Relaxed) + } + + pub(crate) async fn with_parquet( + schema: iceberg::spec::Schema, + batch: &datafusion::arrow::array::RecordBatch, + properties: parquet::file::properties::WriterProperties, + catalog_backed: bool, + ) -> Result<(Self, tempfile::TempDir)> { + use std::collections::HashMap; + + use datafusion::datasource::TableProvider; + use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; + use iceberg::spec::{DataContentType, DataFileBuilder, DataFileFormat, Struct}; + use iceberg::transaction::{ApplyTransactionAction, Transaction}; + use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation}; + use parquet::arrow::ArrowWriter; + + use crate::{ + IcebergCatalogTableProvider, IcebergStaticTableProvider, to_datafusion_error, + }; + + let directory = tempfile::TempDir::new()?; + let factory = FixtureStorageFactory { + root: directory.path().to_path_buf(), + ..Default::default() + }; + let parquet_read_bytes = factory.parquet_read_bytes.clone(); + let catalog = Arc::new( + MemoryCatalogBuilder::default() + .with_storage_factory(Arc::new(factory)) + .load( + "fixture", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_string(), + WAREHOUSE_URI.trim_end_matches('/').to_string(), + )]), + ) + .await + .map_err(to_datafusion_error)?, + ); + let namespace = NamespaceIdent::new("default".to_string()); + catalog + .create_namespace(&namespace, HashMap::new()) + .await + .map_err(to_datafusion_error)?; + let table = catalog + .create_table( + &namespace, + TableCreation::builder() + .name("events".to_string()) + .schema(schema) + .build(), + ) + .await + .map_err(to_datafusion_error)?; + let path = format!("{}/data.parquet", table.metadata().location()); + let mut data = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut data, batch.schema(), Some(properties))?; + writer.write(batch)?; + let metadata = writer.close()?; + assert_eq!(metadata.num_row_groups(), 1); + let file = DataFileBuilder::default() + .content(DataContentType::Data) + .partition_spec_id(0) + .file_path(path.clone()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(data.len() as u64) + .record_count(batch.num_rows() as u64) + .partition(Struct::empty()) + .build() + .unwrap(); + table + .file_io() + .new_output(&path) + .map_err(to_datafusion_error)? + .write(data.into()) + .await + .map_err(to_datafusion_error)?; + let transaction = Transaction::new(&table); + let table = transaction + .fast_append() + .add_data_files([file]) + .apply(transaction) + .map_err(to_datafusion_error)? + .commit(catalog.as_ref()) + .await + .map_err(to_datafusion_error)?; + let provider: Arc = if catalog_backed { + Arc::new( + IcebergCatalogTableProvider::try_new(catalog, namespace, "events") + .await?, + ) + } else { + Arc::new(IcebergStaticTableProvider::try_new_from_table(table).await?) + }; + let ctx = SessionContext::new(); + ctx.register_table("events", provider)?; + Ok(( + Self { + ctx, + parquet_read_bytes, + }, + directory, + )) + } +} From f95e1a64202315bac15d0a1235181c46caf3151a Mon Sep 17 00:00:00 2001 From: t3rb3d Date: Tue, 20 Oct 2026 22:49:02 +0200 Subject: [PATCH 2/2] Simplify page-index coverage --- crates/datafusion/src/read/scan.rs | 149 ----------------- crates/datafusion/src/test_utils/harness.rs | 171 +------------------- crates/datafusion/tests/page_index.rs | 50 ++++++ 3 files changed, 54 insertions(+), 316 deletions(-) create mode 100644 crates/datafusion/tests/page_index.rs diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index c94b410..3d9f752 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -306,152 +306,3 @@ async fn get_batch_stream( .map_err(to_datafusion_error); Ok(Box::pin(stream)) } - -#[cfg(test)] -mod tests { - use datafusion::arrow::array::{ArrayRef, Decimal128Array, Int64Array}; - use datafusion::arrow::compute::concat_batches; - use datafusion::prelude::SessionContext; - use iceberg::arrow::schema_to_arrow_schema; - use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; - use parquet::basic::Compression; - use parquet::file::properties::{EnabledStatistics, WriterProperties}; - - use super::*; - use crate::test_utils::IcebergTestHarness; - - const ROWS: usize = 262_144; - const QUERY: &str = "SELECT * FROM events WHERE id BETWEEN 8190 AND 8194 ORDER BY id"; - - fn page_index_fixture() -> (Schema, RecordBatch) { - let mut fields = vec![Arc::new(NestedField::required( - 1, - "id", - Type::Primitive(PrimitiveType::Long), - ))]; - let mut columns: Vec = - vec![Arc::new(Int64Array::from_iter_values(0..ROWS as i64))]; - for (index, precision) in [9, 18, 38].into_iter().enumerate() { - fields.push(Arc::new(NestedField::optional( - index as i32 + 2, - format!("amount_{precision}"), - Type::Primitive(PrimitiveType::Decimal { - precision, - scale: 2, - }), - ))); - let values = Decimal128Array::from_iter((0..ROWS).map(|row| { - (!(row / 8192).is_multiple_of(2)) - .then_some(row as i128 - ROWS as i128 / 2) - })) - .with_precision_and_scale(precision as u8, 2) - .unwrap(); - columns.push(Arc::new(values)); - } - let schema = Schema::builder().with_fields(fields).build().unwrap(); - let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap()); - let batch = RecordBatch::try_new(arrow_schema, columns).unwrap(); - (schema, batch) - } - - fn page_index_properties(enabled: bool) -> WriterProperties { - WriterProperties::builder() - .set_compression(Compression::UNCOMPRESSED) - .set_dictionary_enabled(false) - .set_max_row_group_row_count(Some(ROWS)) - .set_data_page_row_count_limit(8192) - .set_statistics_enabled(if enabled { - EnabledStatistics::Page - } else { - EnabledStatistics::Chunk - }) - .set_offset_index_disabled(!enabled) - .build() - } - - async fn query(ctx: &SessionContext, sql: &str) -> Result { - let batches = ctx.sql(sql).await?.collect().await?; - Ok(concat_batches(&batches[0].schema(), &batches)?) - } - - #[tokio::test] - async fn test_page_index_setting_prunes_reads() -> Result<()> { - let (schema, batch) = page_index_fixture(); - let expected = batch.slice(8190, 5); - for catalog_backed in [false, true] { - let (harness, _directory) = IcebergTestHarness::with_parquet( - schema.clone(), - &batch, - page_index_properties(true), - catalog_backed, - ) - .await?; - let ctx = harness.context(); - let mut read_bytes = [0; 3]; - for (bytes, enabled) in - read_bytes.iter_mut().zip([None, Some(false), Some(true)]) - { - if let Some(enabled) = enabled { - ctx.sql(&format!( - "SET datafusion.execution.parquet.enable_page_index = '{enabled}'", - )).await?.collect().await?; - } - assert_eq!(query(ctx, QUERY).await?, expected); - *bytes = harness.take_parquet_read_bytes(); - } - let [default, disabled, enabled] = read_bytes; - assert_eq!(default, enabled); - assert!( - enabled < disabled, - "page index on: {enabled}, off: {disabled}" - ); - } - Ok(()) - } - - #[tokio::test] - async fn test_page_index_setting_without_indexes() -> Result<()> { - let (schema, batch) = page_index_fixture(); - let expected = batch.slice(8190, 5); - for catalog_backed in [false, true] { - let (harness, _directory) = IcebergTestHarness::with_parquet( - schema.clone(), - &batch, - page_index_properties(false), - catalog_backed, - ) - .await?; - for enabled in [false, true] { - harness.context().sql(&format!( - "SET datafusion.execution.parquet.enable_page_index = '{enabled}'", - )).await?.collect().await?; - assert_eq!(query(harness.context(), QUERY).await?, expected); - } - } - Ok(()) - } - - #[tokio::test] - async fn test_page_index_pruning_with_nullable_decimals_and_projection() -> Result<()> - { - let (schema, batch) = page_index_fixture(); - let expected = batch.slice(8190, 2).project(&[0])?; - for catalog_backed in [false, true] { - let (harness, _directory) = IcebergTestHarness::with_parquet( - schema.clone(), - &batch, - page_index_properties(true), - catalog_backed, - ) - .await?; - for precision in [9, 18, 38] { - let sql = format!( - "SELECT id FROM events WHERE id BETWEEN 8190 AND 8194 \ - AND amount_{precision} IS NULL ORDER BY id", - ); - assert_eq!(query(harness.context(), &sql).await?, expected); - } - } - Ok(()) - } -} diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index 19ba965..d106467 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -16,12 +16,8 @@ // under the License. use std::collections::HashMap; -#[cfg(test)] -use std::ops::Range; use std::path::PathBuf; use std::sync::Arc; -#[cfg(test)] -use std::sync::atomic::{AtomicU64, Ordering}; use async_trait::async_trait; use bytes::Bytes; @@ -53,15 +49,10 @@ const NAMESPACE_NAME: &str = "public"; pub struct IcebergTestHarness { ctx: SessionContext, - #[cfg(test)] - parquet_read_bytes: Arc, } impl IcebergTestHarness { pub async fn new() -> Result { - let storage_factory = FixtureStorageFactory::default(); - #[cfg(test)] - let parquet_read_bytes = storage_factory.parquet_read_bytes.clone(); let state = SessionStateBuilder::new() .with_default_features() // Plan snapshots must not vary with the runner's CPU count. @@ -71,13 +62,13 @@ impl IcebergTestHarness { .with_table_factory( "ICEBERG".to_string(), Arc::new(IcebergTableProviderFactory::new_with_storage_factory( - Arc::new(storage_factory.clone()), + Arc::new(FixtureStorageFactory::default()), )), ) .build(); let catalog = MemoryCatalogBuilder::default() - .with_storage_factory(Arc::new(storage_factory)) + .with_storage_factory(Arc::new(FixtureStorageFactory::default())) .load( CATALOG_NAME, HashMap::from([( @@ -105,11 +96,7 @@ impl IcebergTestHarness { let ctx = SessionContext::new_with_state(state); ctx.register_catalog(CATALOG_NAME, Arc::new(provider)); - Ok(Self { - ctx, - #[cfg(test)] - parquet_read_bytes, - }) + Ok(Self { ctx }) } pub async fn query(&self, sql: &str) -> Result<(String, String)> { @@ -127,17 +114,12 @@ impl IcebergTestHarness { #[derive(Debug, Clone, Serialize, Deserialize)] struct FixtureStorageFactory { root: PathBuf, - #[cfg(test)] - #[serde(skip)] - parquet_read_bytes: Arc, } impl Default for FixtureStorageFactory { fn default() -> Self { Self { root: PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), - #[cfg(test)] - parquet_read_bytes: Arc::default(), } } } @@ -147,8 +129,6 @@ impl StorageFactory for FixtureStorageFactory { fn build(&self, _config: &StorageConfig) -> IcebergResult> { Ok(Arc::new(FixtureStorage { root: self.root.clone(), - #[cfg(test)] - parquet_read_bytes: self.parquet_read_bytes.clone(), })) } } @@ -156,9 +136,6 @@ impl StorageFactory for FixtureStorageFactory { #[derive(Debug, Clone, Serialize, Deserialize)] struct FixtureStorage { root: PathBuf, - #[cfg(test)] - #[serde(skip)] - parquet_read_bytes: Arc, } impl FixtureStorage { @@ -193,15 +170,7 @@ impl Storage for FixtureStorage { } async fn reader(&self, path: &str) -> IcebergResult> { - let reader = self.local().reader(&self.local_path(path)?).await?; - #[cfg(test)] - if path.ends_with(".parquet") { - return Ok(Box::new(CountingFileRead { - reader, - bytes: self.parquet_read_bytes.clone(), - })); - } - Ok(reader) + self.local().reader(&self.local_path(path)?).await } async fn write(&self, path: &str, bytes: Bytes) -> IcebergResult<()> { @@ -239,135 +208,3 @@ impl Storage for FixtureStorage { Ok(OutputFile::new(Arc::new(self.clone()), path.to_string())) } } - -#[cfg(test)] -struct CountingFileRead { - reader: Box, - bytes: Arc, -} - -#[cfg(test)] -#[async_trait] -impl FileRead for CountingFileRead { - async fn read(&self, range: Range) -> IcebergResult { - let bytes = self.reader.read(range).await?; - self.bytes.fetch_add(bytes.len() as u64, Ordering::Relaxed); - Ok(bytes) - } -} - -#[cfg(test)] -impl IcebergTestHarness { - pub(crate) fn context(&self) -> &SessionContext { - &self.ctx - } - - pub(crate) fn take_parquet_read_bytes(&self) -> u64 { - self.parquet_read_bytes.swap(0, Ordering::Relaxed) - } - - pub(crate) async fn with_parquet( - schema: iceberg::spec::Schema, - batch: &datafusion::arrow::array::RecordBatch, - properties: parquet::file::properties::WriterProperties, - catalog_backed: bool, - ) -> Result<(Self, tempfile::TempDir)> { - use std::collections::HashMap; - - use datafusion::datasource::TableProvider; - use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; - use iceberg::spec::{DataContentType, DataFileBuilder, DataFileFormat, Struct}; - use iceberg::transaction::{ApplyTransactionAction, Transaction}; - use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation}; - use parquet::arrow::ArrowWriter; - - use crate::{ - IcebergCatalogTableProvider, IcebergStaticTableProvider, to_datafusion_error, - }; - - let directory = tempfile::TempDir::new()?; - let factory = FixtureStorageFactory { - root: directory.path().to_path_buf(), - ..Default::default() - }; - let parquet_read_bytes = factory.parquet_read_bytes.clone(); - let catalog = Arc::new( - MemoryCatalogBuilder::default() - .with_storage_factory(Arc::new(factory)) - .load( - "fixture", - HashMap::from([( - MEMORY_CATALOG_WAREHOUSE.to_string(), - WAREHOUSE_URI.trim_end_matches('/').to_string(), - )]), - ) - .await - .map_err(to_datafusion_error)?, - ); - let namespace = NamespaceIdent::new("default".to_string()); - catalog - .create_namespace(&namespace, HashMap::new()) - .await - .map_err(to_datafusion_error)?; - let table = catalog - .create_table( - &namespace, - TableCreation::builder() - .name("events".to_string()) - .schema(schema) - .build(), - ) - .await - .map_err(to_datafusion_error)?; - let path = format!("{}/data.parquet", table.metadata().location()); - let mut data = Vec::new(); - let mut writer = - ArrowWriter::try_new(&mut data, batch.schema(), Some(properties))?; - writer.write(batch)?; - let metadata = writer.close()?; - assert_eq!(metadata.num_row_groups(), 1); - let file = DataFileBuilder::default() - .content(DataContentType::Data) - .partition_spec_id(0) - .file_path(path.clone()) - .file_format(DataFileFormat::Parquet) - .file_size_in_bytes(data.len() as u64) - .record_count(batch.num_rows() as u64) - .partition(Struct::empty()) - .build() - .unwrap(); - table - .file_io() - .new_output(&path) - .map_err(to_datafusion_error)? - .write(data.into()) - .await - .map_err(to_datafusion_error)?; - let transaction = Transaction::new(&table); - let table = transaction - .fast_append() - .add_data_files([file]) - .apply(transaction) - .map_err(to_datafusion_error)? - .commit(catalog.as_ref()) - .await - .map_err(to_datafusion_error)?; - let provider: Arc = if catalog_backed { - Arc::new( - IcebergCatalogTableProvider::try_new(catalog, namespace, "events") - .await?, - ) - } else { - Arc::new(IcebergStaticTableProvider::try_new_from_table(table).await?) - }; - let ctx = SessionContext::new(); - ctx.register_table("events", provider)?; - Ok(( - Self { - ctx, - parquet_read_bytes, - }, - directory, - )) - } -} diff --git a/crates/datafusion/tests/page_index.rs b/crates/datafusion/tests/page_index.rs new file mode 100644 index 0000000..4f07527 --- /dev/null +++ b/crates/datafusion/tests/page_index.rs @@ -0,0 +1,50 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::error::Result; +use iceberg_datafusion::test_utils::IcebergTestHarness; + +#[tokio::test] +async fn page_index_setting_without_indexes() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + + // The checked-in taxi files have no column or offset indexes. + for enabled in [false, true] { + harness + .query(&format!( + "SET datafusion.execution.parquet.enable_page_index = '{enabled}'" + )) + .await?; + let (_, batches) = harness + .query( + "SELECT COUNT(*) AS trips FROM taxi \ + WHERE pickup_date = DATE '2024-01-10' \ + AND payment_type IN (1, 2) AND trip_distance >= 2.0", + ) + .await?; + insta::allow_duplicates! { + insta::assert_snapshot!(batches, @r" + +-------+ + | trips | + +-------+ + | 9891 | + +-------+ + "); + } + } + Ok(()) +}