anuragmantri commented on code in PR #55:
URL: https://github.com/apache/datafusion-iceberg/pull/55#discussion_r4321925909
##########
crates/datafusion/src/read/scan.rs:
##########
@@ -53,20 +67,47 @@ pub struct IcebergDataSource {
predicate: Option<Predicate>,
/// Optional limit on the number of rows to return
fetch: Option<usize>,
+ /// The data files to merge in the order the source reports, listed while
+ /// the scan was planned. `None` lists the files when the source runs and
+ /// reports no order.
+ sorted: Option<SortedTasks>,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+/// Data files that all record the same sort order, and the part of that order
+/// the source reports, over the columns of its output.
+#[derive(Debug, Clone)]
+struct SortedTasks {
+ tasks: Vec<FileScanTask>,
+ ordering: LexOrdering,
}
impl IcebergDataSource {
/// Starts building a source that scans `table` and returns the columns of
/// `schema`, all of them unless
[`IcebergDataSourceBuilder::with_projection`]
/// selects some. Columns are read from Iceberg by name.
- pub fn builder(table: Table, schema: ArrowSchemaRef) ->
IcebergDataSourceBuilder {
+ ///
+ /// `sorted_tasks` are the data files to merge in the sort order they all
+ /// record, which the source then reports to DataFusion. With `None`, the
+ /// source lists its files when it runs and reports no order, as a table
+ /// provider's source does. A source rebuilt from another, such as in
+ /// another process, must be given what that source's
+ /// [`sorted_tasks`](Self::sorted_tasks) returns, because a plan built
over a
+ /// sorted source may rely on its order.
[`IcebergDataSourceBuilder::build`]
+ /// fails if the tasks don't share a sort order the source can report.
+ pub fn builder(
+ table: Table,
+ schema: ArrowSchemaRef,
+ sorted_tasks: Option<Vec<FileScanTask>>,
Review Comment:
Went with option 1 in 8014ccd. `IcebergDataSource::builder(table, schema)`
is unchanged from #63, and the sorted tasks go through a new
`with_sorted_tasks` setter. `build()` still rejects tasks that don't share a
sort order the source can report. `plan_sort_order` is gone. In its place, a
crate-private `plan_sorted_tasks(&builder, state)` lists the files and only
returns the tasks, and the providers pass them to `with_sorted_tasks`. So a
provider builds a sorted source through the same builder calls a codec uses to
rebuild one. The case @andygrove raised is caught at test time:
`test_rebuilt_sorted_scan_keeps_order` checks that a rebuild with the tasks
keeps the order and one without them reports none. Making `with_sorted_tasks`
ignore its tasks fails 11 of the 14 sorted-scan tests.
##########
crates/datafusion/tests/integration_datafusion_test.rs:
##########
@@ -1280,3 +1309,876 @@ async fn
test_scan_after_schema_evolution_reads_provider_columns()
Ok(())
}
+
+/// The id of the sort order of the tables `create_sorted_table` creates.
+const SORTED: Option<i32> = Some(1);
+
+/// Creates table `name`, sorted by `id` ascending, with columns `id`, `data`,
+/// a struct `info` and a list `tags`, whose values are derived from `id`. Each
+/// entry of `files` becomes one data file of the given ids, in the given
order,
+/// which records the given sort order id.
+async fn create_sorted_table(
+ catalog: &Arc<dyn Catalog>,
+ namespace: &NamespaceIdent,
+ name: &str,
+ files: &[(Option<i32>, &[i32])],
+) -> Result<Table, Box<dyn Error>> {
+ let sort_order = SortOrder::builder()
+ .with_sort_field(
+ SortField::builder()
+ .source_id(1)
+ .transform(Transform::Identity)
+ .direction(SortDirection::Ascending)
+ .null_order(NullOrder::First)
+ .build(),
+ )
+ .build_unbound()?;
+ let creation = TableCreation::builder()
+ .location(temp_path())
+ .name(name.to_string())
+ .schema(
+ Schema::builder()
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int))
+ .into(),
+ NestedField::required(
+ 2,
+ "data",
+ Type::Primitive(PrimitiveType::String),
+ )
+ .into(),
+ NestedField::optional(
+ 3,
+ "info",
+ Type::Struct(StructType::new(vec![
+ NestedField::optional(
+ 4,
+ "tag",
+ Type::Primitive(PrimitiveType::String),
+ )
+ .into(),
+ ])),
+ )
+ .into(),
+ NestedField::optional(
+ 5,
+ "tags",
+ Type::List(ListType::new(
+ NestedField::list_element(
+ 6,
+ Type::Primitive(PrimitiveType::Int),
+ false,
+ )
+ .into(),
+ )),
+ )
+ .into(),
+ ])
+ .build()?,
+ )
+ .sort_order(sort_order)
+ .build();
+ let table = catalog.create_table(namespace, creation).await?;
+ assert_eq!(
+ Some(table.metadata().default_sort_order_id() as i32),
+ SORTED
+ );
+ let files: Vec<_> = files
+ .iter()
+ .map(|(sort_order_id, ids)| {
+ (*sort_order_id, ids.iter().copied().map(Some).collect())
+ })
+ .collect();
+ append_sorted_files(catalog, &table, &files).await
+}
+
+/// Commits one data file of `table` per entry of `files`, with the given ids,
+/// in the given order, recording the given sort order id. The other columns'
+/// values are derived from `id`.
+async fn append_sorted_files(
+ catalog: &Arc<dyn Catalog>,
+ table: &Table,
+ files: &[(Option<i32>, Vec<Option<i32>>)],
+) -> Result<Table, Box<dyn Error>> {
+ if files.is_empty() {
+ return Ok(table.clone());
+ }
+ let arrow_schema =
+ Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?);
+ let (DataType::Struct(info_fields), DataType::List(tags_field)) = (
+ arrow_schema.field(2).data_type().clone(),
+ arrow_schema.field(3).data_type().clone(),
+ ) else {
+ unreachable!("info is a struct and tags a list");
+ };
+ let data_dir = format!("{}/data", table.metadata().location());
+ std::fs::create_dir_all(&data_dir)?;
+ let snapshots = table.metadata().snapshots().count();
+ let mut data_files = Vec::new();
+ for (i, (sort_order_id, ids)) in files.iter().enumerate() {
+ let path = format!("{data_dir}/{snapshots}-{i}.parquet");
+ let label = |prefix: &str, id: &Option<i32>| match id {
+ Some(id) => format!("{prefix} {id}"),
+ None => format!("{prefix} null"),
+ };
+ let batch = RecordBatch::try_new(
+ arrow_schema.clone(),
+ vec![
+ Arc::new(Int32Array::from(ids.clone())),
+ Arc::new(StringArray::from_iter_values(
+ ids.iter().map(|id| label("row", id)),
+ )),
+ Arc::new(StructArray::new(
+ info_fields.clone(),
+ vec![Arc::new(StringArray::from_iter_values(
+ ids.iter().map(|id| label("tag", id)),
+ ))],
+ None,
+ )),
+ Arc::new(ListArray::new(
+ tags_field.clone(),
+ OffsetBuffer::from_lengths(ids.iter().map(|_| 2)),
+
Arc::new(Int32Array::from_iter_values(ids.iter().flat_map(|id| {
+ let id = id.unwrap_or_default();
+ [id, -id]
+ }))),
+ None,
+ )),
+ ],
+ )?;
+ let mut writer = ArrowWriter::try_new(
+ std::fs::File::create(&path)?,
+ arrow_schema.clone(),
+ None,
+ )?;
+ writer.write(&batch)?;
+ writer.close()?;
+
+ let mut builder = DataFileBuilder::default();
+ builder
+ .content(DataContentType::Data)
+ .file_path(path.clone())
+ .file_format(DataFileFormat::Parquet)
+ .file_size_in_bytes(std::fs::metadata(&path)?.len())
+ .record_count(ids.len() as u64)
+ .partition_spec_id(0)
+ .partition(Struct::empty());
+ if let Some(sort_order_id) = sort_order_id {
+ builder.sort_order_id(*sort_order_id);
+ }
+ data_files.push(builder.build()?);
+ }
+ let tx = Transaction::new(table);
+ let tx = tx.fast_append().add_data_files(data_files).apply(tx)?;
+ Ok(tx.commit(catalog.as_ref()).await?)
+}
+
+/// Stands in for another engine committing, without writing data, the current
+/// schema of `table` with the field of the same id as `field` replaced by it.
+/// iceberg-rust has no action that makes a column optional or promotes its
+/// type, so this rewrites the table's current metadata file in place.
+async fn replace_field_externally(
+ table: &Table,
+ field: NestedField,
+) -> Result<(), Box<dyn Error>> {
+ let current = table.metadata().current_schema();
+ let schema = Schema::builder()
+ .with_fields(current.as_struct().fields().iter().map(|current| {
+ if current.id == field.id {
+ Arc::new(field.clone())
+ } else {
+ current.clone()
+ }
+ }))
+ .build()?;
+ let metadata = table
+ .metadata()
+ .clone()
+ .into_builder(None)
+ .add_current_schema(schema)?
+ .build()?
+ .metadata;
+ let location =
MetadataLocation::from_str(table.metadata_location_result()?)?;
+ metadata.write_to(table.file_io(), &location).await?;
+ Ok(())
+}
+
+/// The values, or nulls, of the `id` column, which must be the first, of
+/// `batches`.
+fn nullable_ids(batches: &[RecordBatch]) -> Vec<Option<i32>> {
+ batches
+ .iter()
+ .flat_map(|batch| batch.column(0).as_primitive::<Int32Type>().iter())
+ .collect()
+}
+
+/// A session that registers the Iceberg options, with
+/// `iceberg.execution.preserve_data_ordering` set to `preserve`.
+async fn ordering_session(
+ preserve: bool,
+ target_partitions: usize,
+) -> Result<SessionContext, Box<dyn Error>> {
+ let config = SessionConfig::new()
+ .with_target_partitions(target_partitions)
+ .with_option_extension(IcebergDataFusionConfig::default());
+ let ctx = SessionContext::new_with_config(config);
+ ctx.sql(&format!(
+ "SET iceberg.execution.preserve_data_ordering = {preserve}"
+ ))
+ .await?
+ .collect()
+ .await?;
+ Ok(ctx)
+}
+
+/// Plans `table` as a full scan in `ctx`.
+async fn plan_sorted_scan(
+ table: &Table,
+ ctx: &SessionContext,
+) -> Result<Arc<dyn ExecutionPlan>, Box<dyn Error>> {
+ let provider =
IcebergStaticTableProvider::try_new_from_table(table.clone()).await?;
+ Ok(provider.scan(&ctx.state(), None, &[], None).await?)
+}
+
+/// The values of the `id` column, which must be the first, of `batches`.
+fn ids(batches: &[RecordBatch]) -> Vec<i32> {
+ batches
+ .iter()
+ .flat_map(|batch| {
+ batch
+ .column(0)
+ .as_primitive::<Int32Type>()
+ .values()
+ .to_vec()
+ })
+ .collect()
+}
+
+async fn sorted_namespace() -> Result<(Arc<dyn Catalog>, NamespaceIdent),
Box<dyn Error>>
+{
+ let catalog = get_iceberg_catalog().await;
+ let namespace = NamespaceIdent::new("sorted".to_string());
+ set_test_namespace(&catalog, &namespace).await?;
+ Ok((Arc::new(catalog), namespace))
+}
+
+#[tokio::test]
+async fn test_sorted_scan_merges_files_with_overlapping_ranges()
Review Comment:
Agreed. In 8014ccd the sorted-scan tests and their helpers are in
`crates/datafusion/tests/sorted_iceberg_scan.rs`, and
`integration_datafusion_test.rs` is back to main apart from one line, where its
codec-style `rebuild_scan` now passes `sorted_tasks()`. I looked at reusing the
Taxi dataset, but its only sort order is the unsorted one (order id 0), so its
files can't exercise a sorted scan without being rewritten. The tests need
specific shapes anyway, such as overlapping key ranges, nulls in the sort key,
files that disagree on their sort order, and a schema changed after the
provider was built, so the file writes its own small tables.
--
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]
