gabotechs commented on code in PR #67:
URL: https://github.com/apache/datafusion-iceberg/pull/67#discussion_r4348130695
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -70,6 +98,45 @@ impl IcebergDataSource {
}
}
+ /// Planned, ungrouped tasks, if file discovery was enabled for this scan.
+ /// A distributed consumer may assign these tasks to its own workers
instead
+ /// of using the local output partitions of this source.
+ pub fn planned_files(&self) -> Option<&PlannedFileScan> {
+ self.planned.as_ref().map(|planned| planned.files.as_ref())
+ }
+
Review Comment:
How about doing something like this:
```rust
pub fn planned_files(&self) -> Option<&[FileScanTask]> {
self.planned.as_ref().map(|planned| planned.files.tasks())
}
```
So that we can avoid widening the surface of the public API by exposing
`PlannedFileScan`?
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -31,11 +31,37 @@ use
datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayFormatType, Partitioning};
use datafusion::prelude::Expr;
use futures::{Stream, TryStreamExt};
+use iceberg::arrow::ArrowReader;
use iceberg::expr::Predicate;
+use iceberg::scan::{FileScanTask, FileScanTaskStream, TableScan};
use iceberg::table::Table;
+use super::file_planning::PlannedFileScan;
use super::predicate::convert_filters_to_predicate;
-use crate::to_datafusion_error;
+use crate::{IcebergPlanningConfig, to_datafusion_error};
+
+// A reader is created once per physical scan. Its clones share the delete
loader's
+// cache while each read() call still owns its own scan metrics.
+#[derive(Clone)]
+struct ScanReader(ArrowReader);
Review Comment:
A bit weird to have this extra wrapping just so that we can `Debug` it.
Even if more verbose, A manual `Debug` impl for the whole
`IcebergDataSource` might be less weird.
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -145,32 +216,128 @@ impl IcebergDataSourceBuilder {
self
}
+ /// Finish configuring the scan and optionally select its files before
execution.
+ ///
+ /// Honors [`IcebergPlanningConfig::eager_file_planning`]. The returned
builder
+ /// only exposes `build`, so filters, projection, and snapshot cannot
change
+ /// after files have been selected.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the projection is invalid, file planning fails, or
+ /// the retained tasks exceed the query memory pool.
+ pub async fn with_planned_files(
+ self,
+ context: Arc<TaskContext>,
+ ) -> Result<PlannedIcebergDataSourceBuilder> {
+ if !context
Review Comment:
This is a bit weird from a public API standpoint:
It says `with_planned_files()`, but it might not return a builder with
planned files. Specially as this is part of the public API, I'd be very
intentional about how this method behaves, and people might find surprising
that they are calling `with_planned_files` and it might not actually plan any
file.
##########
crates/datafusion/src/read/file_planning.rs:
##########
@@ -0,0 +1,81 @@
+// 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.
+
+//! Query-local file work selected before execution, independently of task
assignment.
+
+use std::mem::{size_of, size_of_val};
+
+use datafusion::common::Result;
+use datafusion::execution::TaskContext;
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use futures::TryStreamExt;
+use iceberg::scan::{FileScanTask, TableScan};
+
+use crate::to_datafusion_error;
+
+/// File work selected by one Iceberg scan before assignment to output
partitions.
+///
+/// The tasks retain their bound predicates, projection, and delete
references. Keeping
+/// this value alive also keeps the query memory reservation for the retained
tasks.
+/// This is not a reader and contains no coordinator-local reader state.
+#[derive(Debug)]
+pub struct PlannedFileScan {
Review Comment:
This struct seems like an intermediate wrapper for holding the memory
reservation, do you think it would be possible to make it `pub(crate)` instead
of `pub`?
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -145,32 +216,128 @@ impl IcebergDataSourceBuilder {
self
}
+ /// Finish configuring the scan and optionally select its files before
execution.
+ ///
+ /// Honors [`IcebergPlanningConfig::eager_file_planning`]. The returned
builder
+ /// only exposes `build`, so filters, projection, and snapshot cannot
change
+ /// after files have been selected.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the projection is invalid, file planning fails, or
+ /// the retained tasks exceed the query memory pool.
+ pub async fn with_planned_files(
+ self,
+ context: Arc<TaskContext>,
+ ) -> Result<PlannedIcebergDataSourceBuilder> {
+ if !context
+ .session_config()
+ .options()
+ .extensions
+ .get::<IcebergPlanningConfig>()
+ .is_some_and(|config| config.eager_file_planning)
+ {
+ return Ok(PlannedIcebergDataSourceBuilder {
+ builder: self,
+ planned: None,
+ });
+ }
+ let target_partitions = context.session_config().target_partitions();
+ let schema = self.projected_schema()?;
+ let projection: Vec<String> = schema
+ .fields()
+ .iter()
+ .map(|field| field.name().clone())
+ .collect();
+ let scan = build_table_scan(
+ &self.table,
+ self.snapshot_id,
+ Some(&projection),
+ self.predicate.as_ref(),
+ )?;
+ // Dropping physical planning drops collection and its memory
reservation.
+ let planned = PlannedFileScan::collect(scan, &context).await?;
+ let partitions =
planned.tasks().len().max(1).min(target_partitions.max(1));
+
+ Ok(PlannedIcebergDataSourceBuilder {
+ builder: self,
+ planned: Some(PlannedScan {
+ files: Arc::new(planned),
+ partitions,
+ }),
+ })
+ }
+
/// Builds the source, which reads the columns of its projected schema by
name.
///
/// # Errors
///
/// Returns an error if the projection is out of bounds for the schema.
pub fn build(self) -> Result<IcebergDataSource> {
- let schema = match &self.projection {
- None => self.schema,
- Some(projection) => Arc::new(self.schema.project(projection)?),
- };
+ self.build_with_planned_files(None)
+ }
+
+ fn projected_schema(&self) -> Result<ArrowSchemaRef> {
+ match &self.projection {
+ None => Ok(self.schema.clone()),
+ Some(projection) => Ok(Arc::new(self.schema.project(projection)?)),
+ }
+ }
+
+ fn build_with_planned_files(
+ self,
+ planned: Option<PlannedScan>,
+ ) -> Result<IcebergDataSource> {
+ let schema = self.projected_schema()?;
let projection = schema
.fields()
.iter()
.map(|field| field.name().clone())
.collect();
+ let reader = self.table.reader_builder();
+ let reader = if planned
+ .as_ref()
+ .is_some_and(|planned| planned.partitions > 1)
+ {
+ reader.with_data_file_concurrency_limit(1)
+ } else {
+ reader
+ };
+ let reader = ScanReader(reader.build());
Ok(IcebergDataSource {
table: self.table,
snapshot_id: self.snapshot_id,
schema,
projection,
predicate: self.predicate,
fetch: self.fetch,
+ planned,
+ reader,
})
}
}
+/// A configured scan ready to build after optional eager file planning.
+///
+/// Created by [`IcebergDataSourceBuilder::with_planned_files`]. Configuration
+/// is fixed so the retained tasks stay consistent with the resulting source.
+#[derive(Debug, Clone)]
+pub struct PlannedIcebergDataSourceBuilder {
Review Comment:
Do we really need a separate builder for this? I think it should be possible
to just have `IcebergDataSourceBuilder`.
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -231,11 +424,18 @@ impl DataSource for IcebergDataSource {
if let Some(fetch) = self.fetch {
write!(f, ", fetch={fetch}")?;
}
+ if let Some(planned) = self.planned_files() {
+ write!(f, ", planned_files={}", planned.tasks().len())?;
+ }
Ok(())
}
fn output_partitioning(&self) -> Partitioning {
- Partitioning::UnknownPartitioning(1)
+ Partitioning::UnknownPartitioning(
+ self.planned
+ .as_ref()
+ .map_or(1, |planned| planned.partitions),
+ )
}
Review Comment:
It's a shame we cannot run the lazy FileScanTask discovery mode partitioned.
I don't think there should be a fundamental limitation that prevents us from
doing it here.
##########
crates/datafusion/src/config.rs:
##########
@@ -0,0 +1,32 @@
+// 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::common::extensions_options;
+use datafusion::config::ConfigExtension;
+
+extensions_options! {
+ /// Opt-in file planning during physical planning for local Iceberg scans.
+ pub struct IcebergPlanningConfig {
+ /// Discover and retain file tasks in TableProvider::scan instead of
at execution.
+ pub eager_file_planning: bool, default = false
+ }
+}
+
+impl ConfigExtension for IcebergPlanningConfig {
+ // Distinct from DFD's `iceberg` options so both integrations can share a
session.
+ const PREFIX: &'static str = "iceberg_local";
Review Comment:
🙈
##########
crates/datafusion/src/table_provider/catalog_table_provider.rs:
##########
@@ -126,23 +126,25 @@ impl TableProvider for IcebergCatalogTableProvider {
async fn scan(
&self,
- _state: &dyn Session,
+ state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
- // Load fresh table metadata from catalog
+ // Load fresh table metadata from catalog.
let table = self
.catalog
.load_table(&self.table_ident)
.await
.map_err(to_datafusion_error)?;
- // Create scan with fresh metadata (always use current snapshot)
+ // Create scan with fresh metadata (always use current snapshot).
let source = IcebergDataSource::builder(table, self.schema.clone())
.with_projection(projection)
.with_filters(filters)
.with_fetch(limit)
+ .with_planned_files(state.task_ctx())
+ .await?
.build()?;
Review Comment:
If eager file scan is disabled, we probably should not be calling this here.
I know the disabled param check leaves inside the function, but it might be
better to surface it here in order to be explicit about it.
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -231,11 +424,18 @@ impl DataSource for IcebergDataSource {
if let Some(fetch) = self.fetch {
write!(f, ", fetch={fetch}")?;
}
+ if let Some(planned) = self.planned_files() {
+ write!(f, ", planned_files={}", planned.tasks().len())?;
+ }
Ok(())
}
fn output_partitioning(&self) -> Partitioning {
- Partitioning::UnknownPartitioning(1)
+ Partitioning::UnknownPartitioning(
+ self.planned
+ .as_ref()
+ .map_or(1, |planned| planned.partitions),
+ )
}
Review Comment:
There does not seem to be any issue open about this, it might be worth
opening one now and linking it here.
##########
crates/datafusion/tests/common/position_deletes.rs:
##########
@@ -0,0 +1,197 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::HashMap;
+use std::error::Error;
+use std::fs;
+use std::io;
+use std::path::Path;
+use std::sync::Arc;
+
+use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::prelude::SessionConfig;
+use iceberg::Runtime;
+use iceberg::metadata_columns::{
+ RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS,
+ RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS,
+};
+use iceberg::spec::{
+ DataContentType, DataFileBuilder, DataFileFormat, ManifestListWriter,
+ ManifestWriterBuilder, Operation, Snapshot, Summary,
+};
+use iceberg::table::Table;
+use iceberg_datafusion::IcebergCatalogTableProvider;
+use iceberg_datafusion::test_utils::IcebergTestHarness;
+use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
+use tempfile::TempDir;
+
+/// A taxi table with its first data file's first row deleted.
+/// Owns the copied files for the lifetime of scans against this table.
+pub struct PositionDeleteFixture {
+ pub table: Table,
+ _directory: TempDir,
+}
+
+pub async fn taxi_with_position_delete() -> Result<PositionDeleteFixture,
Box<dyn Error>>
Review Comment:
:thinking: I'm realizing that the Taxi dataset that was committed is not
really working well...
I've seen several PRs with people facing fundamental limitations with that
dataset that require a bunch of customization and overwriting of the dataset
for making it suitable for certain tests.
We probably need to rethink the integration testing approach based on a
statically committed dataset, it does not seem to be flexible enough to cover
all the testing needs in the project.
##########
crates/datafusion/src/table_provider/static_table_provider.rs:
##########
@@ -120,17 +120,19 @@ impl TableProvider for IcebergStaticTableProvider {
async fn scan(
&self,
- _state: &dyn Session,
+ state: &dyn Session,
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
- // Use cached table (no refresh)
+ // Use cached table (no refresh).
let source = IcebergDataSource::builder(self.table.clone(),
self.schema.clone())
.with_snapshot_id(self.snapshot_id)
.with_projection(projection)
.with_filters(filters)
.with_fetch(limit)
+ .with_planned_files(state.task_ctx())
+ .await?
.build()?;
Review Comment:
Same as above
##########
crates/datafusion/src/config.rs:
##########
@@ -0,0 +1,32 @@
+// 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::common::extensions_options;
+use datafusion::config::ConfigExtension;
+
+extensions_options! {
+ /// Opt-in file planning during physical planning for local Iceberg scans.
+ pub struct IcebergPlanningConfig {
+ /// Discover and retain file tasks in TableProvider::scan instead of
at execution.
+ pub eager_file_planning: bool, default = false
+ }
+}
+
+impl ConfigExtension for IcebergPlanningConfig {
+ // Distinct from DFD's `iceberg` options so both integrations can share a
session.
+ const PREFIX: &'static str = "iceberg_local";
Review Comment:
Maybe just "iceberg" and remove the comment. Alternatively, I think the
configuration structure shipped in
https://github.com/apache/datafusion-iceberg/pull/55 was quite good, we could
back port that here if that makes sense.
##########
crates/datafusion/tests/eager_file_planning.rs:
##########
@@ -0,0 +1,362 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::error::Error;
+use std::result::Result as StdResult;
+use std::sync::Arc;
+
+use datafusion::common::{Result, stats::Precision};
+use datafusion::datasource::TableProvider;
+use datafusion::datasource::source::{DataSource, DataSourceExec};
+use datafusion::execution::memory_pool::GreedyMemoryPool;
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::physical_plan::common::collect;
+use datafusion::physical_plan::{ExecutionPlanProperties, collect as
collect_plan};
+use datafusion::prelude::{SessionConfig, col, lit};
+use datafusion::scalar::ScalarValue;
+use iceberg_datafusion::read::IcebergDataSource;
+use iceberg_datafusion::test_utils::IcebergTestHarness;
+use iceberg_datafusion::{
+ IcebergCatalogTableProvider, IcebergPlanningConfig,
IcebergStaticTableProvider,
+};
+
+mod common;
+
+use common::position_deletes::taxi_with_position_delete;
+
+fn eager_config() -> SessionConfig {
+ SessionConfig::new()
+ .with_target_partitions(4)
+ .with_option_extension(IcebergPlanningConfig::default())
+ .set_bool("iceberg_local.eager_file_planning", true)
+}
+
+fn iceberg_source(
+ plan: &dyn datafusion::physical_plan::ExecutionPlan,
+) -> &IcebergDataSource {
+ plan.downcast_ref::<DataSourceExec>()
+ .expect("the Iceberg provider creates a DataSourceExec")
+ .data_source()
+ .downcast_ref::<IcebergDataSource>()
+ .expect("the plan contains the Iceberg data source")
+}
+
+#[tokio::test]
+async fn plans_file_tasks_before_execution_without_changing_statistics() ->
Result<()> {
+ let eager = IcebergTestHarness::with_config(eager_config()).await?;
+ let lazy = IcebergTestHarness::new().await?;
+ let eager_provider = eager.context().table_provider("taxi").await?;
+ let lazy_provider = lazy.context().table_provider("taxi").await?;
+ let eager_plan = eager_provider
+ .scan(&eager.context().state(), None, &[], None)
+ .await?;
+ let lazy_plan = lazy_provider
+ .scan(&lazy.context().state(), None, &[], None)
+ .await?;
+ let eager_source = iceberg_source(eager_plan.as_ref());
+ let lazy_source = iceberg_source(lazy_plan.as_ref());
+
+ assert!(lazy_source.planned_files().is_none());
+ assert!(lazy_source.planned_partition_tasks(0)?.is_none());
+ assert!(lazy_source.planned_partition_tasks(1).is_err());
+ assert!(eager_source.planned_partition_tasks(4).is_err());
+ let tasks = eager_source.planned_files().unwrap().tasks();
+ for (partition, indices) in [vec![0, 4], vec![1, 5], vec![2, 6], vec![3]]
+ .into_iter()
+ .enumerate()
+ {
+ let assigned: Vec<_> = eager_source
+ .planned_partition_tasks(partition)?
+ .unwrap()
+ .map(std::ptr::from_ref)
+ .collect();
+ let expected: Vec<_> = indices
+ .into_iter()
+ .map(|index| std::ptr::from_ref(&tasks[index]))
+ .collect();
+ assert_eq!(assigned, expected);
+ }
+
+ assert_eq!(lazy_plan.output_partitioning().partition_count(), 1);
+ assert_eq!(eager_source.planned_files().unwrap().tasks().len(), 7);
+ assert_eq!(eager_plan.output_partitioning().partition_count(), 4);
+ assert!(eager_source.open(4, eager.context().task_ctx()).is_err());
+ assert!(lazy_source.open(1, lazy.context().task_ctx()).is_err());
+ assert_eq!(
+ eager_source.partition_statistics(None)?.num_rows,
+ Precision::Absent
+ );
+ assert_eq!(
+ eager_source.partition_statistics(None)?.total_byte_size,
+ lazy_source.partition_statistics(None)?.total_byte_size
+ );
+
+ let empty = eager_provider
+ .scan(
+ &eager.context().state(),
+ None,
+ &[col("pickup_date").eq(lit(ScalarValue::Date32(Some(0))))],
+ None,
+ )
+ .await?;
+ assert!(
+ iceberg_source(empty.as_ref())
+ .planned_files()
+ .unwrap()
+ .tasks()
+ .is_empty()
+ );
+ assert_eq!(empty.output_partitioning().partition_count(), 1);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn eager_partition_count_is_bounded_by_selected_files() -> Result<()> {
+ let eager =
+
IcebergTestHarness::with_config(eager_config().with_target_partitions(16))
+ .await?;
+ let provider = eager.context().table_provider("taxi").await?;
+ let plan = provider
+ .scan(&eager.context().state(), None, &[], None)
+ .await?;
+ assert_eq!(plan.output_partitioning().partition_count(), 7);
+ Ok(())
+}
+
+#[tokio::test]
+async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> {
+ let eager = IcebergTestHarness::with_config(eager_config()).await?;
+ let lazy = IcebergTestHarness::new().await?;
+ for (query, planned_files) in [
+ (
+ "SELECT pickup_date, COUNT(*) AS trips FROM taxi \
+ WHERE pickup_date = DATE '2024-01-10' \
+ GROUP BY pickup_date ORDER BY pickup_date",
+ 1,
+ ),
+ (
+ "SELECT COUNT(*) AS trips FROM taxi \
+ WHERE pickup_date = DATE '2024-01-08' AND trip_distance > 50",
+ 1,
+ ),
+ (
+ "SELECT COUNT(*) AS trips FROM taxi \
+ WHERE pickup_date = DATE '2025-01-01'",
+ 0,
+ ),
+ ] {
+ let (plan, eager_batches) = eager.query(query).await?;
+ let (_, lazy_batches) = lazy.query(query).await?;
+ assert!(
+ plan.contains(&format!("planned_files={planned_files}")),
+ "eager scan not visible in plan: {plan}"
+ );
+ assert_eq!(eager_batches, lazy_batches);
+ }
+ Ok(())
+}
+
+#[tokio::test]
+async fn eager_multi_partition_limit_is_global() -> Result<()> {
+ let eager = IcebergTestHarness::with_config(eager_config()).await?;
+ let (plan, batches) = eager.query("SELECT pickup_date FROM taxi LIMIT
3").await?;
+ assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}");
+ assert_eq!(batches.matches("2024-01-").count(), 3, "{batches}");
+ Ok(())
+}
+
+#[tokio::test]
+async fn eager_ordered_limit_matches_lazy_results() -> Result<()> {
+ let eager = IcebergTestHarness::with_config(eager_config()).await?;
+ let lazy = IcebergTestHarness::new().await?;
+ let sql = "SELECT vendor_id, pickup_at FROM taxi \
+ ORDER BY pickup_at, vendor_id, pickup_location_id LIMIT 3";
+ let (plan, eager_batches) = eager.query(sql).await?;
+ let (_, lazy_batches) = lazy.query(sql).await?;
+ assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}");
+ assert_eq!(eager_batches, lazy_batches);
+ Ok(())
+}
+
+#[tokio::test]
+async fn eager_and_lazy_scans_apply_position_deletes() -> StdResult<(),
Box<dyn Error>> {
+ let fixture = taxi_with_position_delete().await?;
+ let eager = IcebergTestHarness::with_config(eager_config()).await?;
+ let table = fixture.table.clone();
+ let provider =
IcebergStaticTableProvider::try_new_from_table(table).await?;
+ let lazy_state =
IcebergTestHarness::with_config(SessionConfig::new()).await?;
+
+ let state = eager.context().state();
+ let eager_plan = provider.scan(&state, None, &[], None).await?;
+ assert_eq!(
+ iceberg_source(eager_plan.as_ref())
+ .planned_files()
+ .expect("eager plan has files")
+ .tasks()
+ .iter()
+ .filter(|task| !task.deletes().is_empty())
+ .count(),
+ 1
+ );
+ assert_eq!(
+ iceberg_source(eager_plan.as_ref())
+ .planned_files()
+ .expect("eager plan has files")
+ .tasks()
+ .len(),
+ 7
+ );
+ let lazy_plan = provider
+ .scan(&lazy_state.context().state(), None, &[], None)
+ .await?;
+ let eager_batches = collect_plan(eager_plan, state.task_ctx()).await?;
+ let lazy_batches = collect_plan(lazy_plan,
lazy_state.context().task_ctx()).await?;
+ let eager_rows: usize = eager_batches.iter().map(|batch|
batch.num_rows()).sum();
+ let lazy_rows: usize = lazy_batches.iter().map(|batch|
batch.num_rows()).sum();
+ assert_eq!(lazy_rows, 174_999);
+ assert_eq!(eager_rows, lazy_rows);
+ Ok(())
+}
+
+#[test]
+fn eager_scan_uses_the_tables_custom_runtime() -> Result<()> {
+ let io_runtime = tokio::runtime::Builder::new_multi_thread()
+ .enable_all()
+ .build()?;
+ let (harness, provider) = io_runtime.block_on(async {
+ let harness = IcebergTestHarness::with_config(eager_config()).await?;
+ let provider = harness.context().table_provider("taxi").await?;
+ let catalog_provider = provider
+ .downcast_ref::<IcebergCatalogTableProvider>()
+ .expect("the fixture creates a catalog Iceberg provider");
+ let table = catalog_provider
+ .catalog()
+ .load_table(catalog_provider.table_ident())
+ .await
+ .map_err(iceberg_datafusion::to_datafusion_error)?;
+ let provider =
IcebergStaticTableProvider::try_new_from_table(table).await?;
+ Ok::<_, datafusion::error::DataFusionError>((harness, provider))
+ })?;
+ // The table retains its own runtime while planning is called from another
runtime.
+ let planning_runtime = tokio::runtime::Builder::new_current_thread()
+ .enable_all()
+ .build()?;
+ planning_runtime.block_on(async {
+ let state = harness.context().state();
+ let plan = provider.scan(&state, None, &[], None).await?;
+ assert_eq!(
+ iceberg_source(plan.as_ref())
+ .planned_files()
+ .unwrap()
+ .tasks()
+ .len(),
+ 7
+ );
+ let mut rows = 0;
+ for partition in 0..plan.output_partitioning().partition_count() {
+ let batches = collect(plan.execute(partition,
state.task_ctx())?).await?;
+ rows += batches.iter().map(|batch|
batch.num_rows()).sum::<usize>();
+ }
+ assert_eq!(rows, 175_000);
+ Ok(())
+ })
+}
+
+#[tokio::test]
+async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> {
+ let runtime = Arc::new(
+ RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::new(GreedyMemoryPool::new(1)))
+ .build()?,
+ );
+ let eager =
+ IcebergTestHarness::with_config_and_runtime(eager_config(),
runtime.clone())
+ .await?;
+ let provider = eager.context().table_provider("taxi").await?;
+ let error = provider
+ .scan(&eager.context().state(), None, &[], None)
+ .await
+ .expect_err("file metadata cannot fit in one byte");
+ assert!(
+ error.to_string().contains("Iceberg file planning"),
+ "{error}"
+ );
+ assert_eq!(runtime.memory_pool.reserved(), 0);
Review Comment:
I see a lot of text-based assertions in these tests being done manually
rather than using `insta`.
Is there any fundamental limitation for not using it here? or is it mostly
about saving some lines.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
