gabotechs commented on code in PR #55:
URL: https://github.com/apache/datafusion-iceberg/pull/55#discussion_r4311048885
##########
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:
The whole point of having a builder pattern here is so that we can avoid
breaking the API with new arguments that might or might not be provided.
Can we adhere here to the current pattern and avoid breaking the current
signature? this implies delegating `IcebergDataSource` construction fully to
`IcebergDataSourceBuilder` as expected, and avoid creating extra side
constructors like `IcebergDataSourceBuilder::plan_sort_order`.
##########
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:
This file is reaching a dangerous size, and the amount of things being
tested here is too big to be scalable in the long run.
Could we refactor these tests into an isolated integration test in something
like `crates/datafusion/tests/sorted_iceberg_scan.rs`? With a bit of luck, we
can even reuse the existing Taxi dataset for reducing the amount of preparatory
churn.
--
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]
