Skip to content

DataFusion 52 migration #3052

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
comphead wants to merge 7 commits into
base: main
Choose a base branch
from
Draft

DataFusion 52 migration #3052

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
62c847f
DataFusion 52 migration
comphead Jan 7, 2026
275a870
DataFusion 52 migration
comphead Jan 7, 2026
38b4ce7
DataFusion 52 migration
comphead Jan 7, 2026
e5d1b0a
DataFusion 52 migration
comphead Jan 9, 2026
20d5b74
DataFusion 52 migration
comphead Jan 14, 2026
3a695ed
DataFusion 52 migration
comphead Jan 14, 2026
1db2215
DataFusion 52 migration
comphead Jan 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
816 changes: 402 additions & 414 deletions native/Cargo.lock
View file
Open in desktop

Large diffs are not rendered by default.

7 changes: 4 additions & 3 deletions native/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,10 @@ arrow = { version = "57.0.0", features = ["prettyprint", "ffi", "chrono-tz"] }
async-trait = { version = "0.1" }
bytes = { version = "1.10.0" }
parquet = { version = "57.0.0", default-features = false, features = ["experimental"] }
datafusion = { version = "51.0.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
datafusion-datasource = { version = "51.0.0" }
datafusion-spark = { version = "51.0.0" }
datafusion = { version = "52.0.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] }
datafusion-datasource = { version = "52.0.0" }
datafusion-spark = { version = "52.0.0" }
datafusion-physical-expr-adapter = { version = "52.0.0" }
datafusion-comet-spark-expr = { path = "spark-expr" }
datafusion-comet-proto = { path = "proto" }
chrono = { version = "0.4", default-features = false, features = ["clock"] }
Expand Down
3 changes: 2 additions & 1 deletion native/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ tempfile = "3.24.0"
itertools = "0.14.0"
paste = "1.0.14"
datafusion = { workspace = true, features = ["parquet_encryption", "sql"] }
datafusion-physical-expr-adapter = { workspace = true }
datafusion-datasource = { workspace = true }
datafusion-spark = { workspace = true }
once_cell = "1.18.0"
Expand Down Expand Up @@ -95,7 +96,7 @@ jni = { version = "0.21", features = ["invocation"] }
lazy_static = "1.4"
assertables = "9"
hex = "0.4.3"
datafusion-functions-nested = { version = "51.0.0" }
datafusion-functions-nested = { version = "52.0.0" }

[features]
backtrace = ["datafusion/backtrace"]
Expand Down
27 changes: 9 additions & 18 deletions native/core/src/execution/operators/iceberg_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,7 @@ use iceberg::io::FileIO;

use crate::execution::operators::ExecutionError;
use crate::parquet::parquet_support::SparkParquetOptions;
use crate::parquet::schema_adapter::SparkSchemaAdapterFactory;
use datafusion::datasource::schema_adapter::SchemaAdapterFactory;
use crate::parquet::schema_adapter::adapt_batch_with_expressions;
use datafusion_comet_spark_expr::EvalMode;
use datafusion_datasource::file_stream::FileStreamMetrics;

Expand Down Expand Up @@ -299,22 +298,14 @@ impl IcebergFileStream {
.map_err(|e| DataFusionError::Execution(format!("Iceberg scan error: {}", e)))
.and_then(move |batch| {
let spark_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);
let adapter_factory = SparkSchemaAdapterFactory::new(spark_options, None);
let file_schema = batch.schema();
let adapter = adapter_factory
.create(Arc::clone(&target_schema), Arc::clone(&file_schema));

let result = match adapter.map_schema(file_schema.as_ref()) {
Ok((schema_mapper, _projection)) => {
schema_mapper.map_batch(batch).map_err(|e| {
DataFusionError::Execution(format!("Batch mapping failed: {}", e))
})
}
Err(e) => Err(DataFusionError::Execution(format!(
"Schema mapping failed: {}",
e
))),
};
let result =
adapt_batch_with_expressions(batch, &target_schema, &spark_options)
.map_err(|e| {
DataFusionError::Execution(format!(
"Batch adaptation failed: {}",
e
))
});
futures::future::ready(result)
});

Expand Down
37 changes: 18 additions & 19 deletions native/core/src/execution/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1093,18 +1093,11 @@ impl PhysicalPlanner {
let files =
self.get_partitioned_files(&scan.file_partitions[self.partition as usize])?;
let file_groups: Vec<Vec<PartitionedFile>> = vec![files];
let partition_fields: Vec<Field> = partition_schema
.fields()
.iter()
.map(|field| {
Field::new(field.name(), field.data_type().clone(), field.is_nullable())
})
.collect_vec();

let scan = init_datasource_exec(
required_schema,
Some(data_schema),
Some(partition_schema),
Some(partition_fields),
object_store_url,
file_groups,
Some(projection_vector),
Expand Down Expand Up @@ -3385,6 +3378,7 @@ mod tests {
use arrow::array::{Array, DictionaryArray, Int32Array, ListArray, RecordBatch, StringArray};
use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema};
use datafusion::catalog::memory::DataSourceExec;
use datafusion::config::TableParquetOptions;
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::object_store::ObjectStoreUrl;
use datafusion::datasource::physical_plan::{
Expand All @@ -3394,6 +3388,7 @@ mod tests {
use datafusion::logical_expr::ScalarUDF;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::{assert_batches_eq, physical_plan::common::collect, prelude::SessionContext};
use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
use tempfile::TempDir;
use tokio::sync::mpsc;

Expand All @@ -3402,7 +3397,7 @@ mod tests {
use crate::execution::operators::ExecutionError;
use crate::execution::planner::literal_to_array_ref;
use crate::parquet::parquet_support::SparkParquetOptions;
use crate::parquet::schema_adapter::SparkSchemaAdapterFactory;
use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory;
use datafusion_comet_proto::spark_expression::expr::ExprStruct;
use datafusion_comet_proto::spark_expression::ListLiteral;
use datafusion_comet_proto::{
Expand Down Expand Up @@ -4004,18 +3999,22 @@ mod tests {
}
}

let source = ParquetSource::default().with_schema_adapter_factory(Arc::new(
SparkSchemaAdapterFactory::new(
SparkParquetOptions::new(EvalMode::Ansi, "", false),
None,
),
))?;
let source = Arc::new(
ParquetSource::new(Arc::new(read_schema.clone()))
.with_table_parquet_options(TableParquetOptions::new()),
) as Arc<dyn FileSource>;

let spark_parquet_options = SparkParquetOptions::new(EvalMode::Legacy, "UTC", false);

let expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory> = Arc::new(
SparkPhysicalExprAdapterFactory::new(spark_parquet_options, None),
);

let object_store_url = ObjectStoreUrl::local_filesystem();
let file_scan_config =
FileScanConfigBuilder::new(object_store_url, read_schema.into(), source)
.with_file_groups(file_groups)
.build();
let file_scan_config = FileScanConfigBuilder::new(object_store_url, source)
.with_expr_adapter(Some(expr_adapter_factory))
.with_file_groups(file_groups)
.build();

// Run native read
let scan = Arc::new(DataSourceExec::new(Arc::new(file_scan_config.clone())));
Expand Down
1 change: 0 additions & 1 deletion native/core/src/parquet/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -765,7 +765,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
required_schema,
Some(data_schema),
None,
None,
object_store_url,
file_groups,
None,
Expand Down
92 changes: 38 additions & 54 deletions native/core/src/parquet/parquet_exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
use crate::execution::operators::ExecutionError;
use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID};
use crate::parquet::parquet_support::SparkParquetOptions;
use crate::parquet::schema_adapter::SparkSchemaAdapterFactory;
use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory;
use arrow::datatypes::{Field, SchemaRef};
use datafusion::config::TableParquetOptions;
use datafusion::datasource::listing::PartitionedFile;
Expand All @@ -29,10 +29,11 @@ use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::object_store::ObjectStoreUrl;
use datafusion::physical_expr::expressions::BinaryExpr;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_expr_adapter::PhysicalExprAdapterFactory;
use datafusion::prelude::SessionContext;
use datafusion::scalar::ScalarValue;
use datafusion_comet_spark_expr::EvalMode;
use itertools::Itertools;
use datafusion_datasource::TableSchema;
use std::collections::HashMap;
use std::sync::Arc;

Expand Down Expand Up @@ -60,7 +61,6 @@ pub(crate) fn init_datasource_exec(
required_schema: SchemaRef,
data_schema: Option<SchemaRef>,
partition_schema: Option<SchemaRef>,
partition_fields: Option<Vec<Field>>,
object_store_url: ObjectStoreUrl,
file_groups: Vec<Vec<PartitionedFile>>,
projection_vector: Option<Vec<usize>>,
Expand All @@ -78,7 +78,26 @@ pub(crate) fn init_datasource_exec(
encryption_enabled,
);

let mut parquet_source = ParquetSource::new(table_parquet_options);
// Determine the schema to use for ParquetSource
let table_schema = if let Some(ref data_schema) = data_schema {
Comment on lines +81 to +82
Copy link

@adriangb adriangb Jan 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is also TableSchema::with_table_partition_cols which might make this easier

if let Some(ref partition_schema) = partition_schema {
let partition_fields: Vec<_> = partition_schema
.fields()
.iter()
.map(|f| {
Arc::new(Field::new(f.name(), f.data_type().clone(), f.is_nullable())) as _
})
.collect();
TableSchema::new(Arc::clone(data_schema), partition_fields)
} else {
TableSchema::from_file_schema(Arc::clone(data_schema))
}
} else {
TableSchema::from_file_schema(Arc::clone(&required_schema))
};

let mut parquet_source =
ParquetSource::new(table_schema).with_table_parquet_options(table_parquet_options);

// Create a conjunctive form of the vector because ParquetExecBuilder takes
// a single expression
Expand All @@ -104,37 +123,27 @@ pub(crate) fn init_datasource_exec(
);
}

let file_source = parquet_source.with_schema_adapter_factory(Arc::new(
SparkSchemaAdapterFactory::new(spark_parquet_options, default_values),
))?;
let expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory> = Arc::new(
SparkPhysicalExprAdapterFactory::new(spark_parquet_options, default_values),
);

let file_source: Arc<dyn FileSource> = Arc::new(parquet_source);

let file_groups = file_groups
.iter()
.map(|files| FileGroup::new(files.clone()))
.collect();

let file_scan_config = match (data_schema, projection_vector, partition_fields) {
(Some(data_schema), Some(projection_vector), Some(partition_fields)) => {
get_file_config_builder(
data_schema,
partition_schema,
file_groups,
object_store_url,
file_source,
)
.with_projection_indices(Some(projection_vector))
.with_table_partition_cols(partition_fields)
Copy link
Contributor Author

@comphead comphead Jan 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

with_partition_cols now derived automatically
https://datafusion.apache.org/library-user-guide/upgrading.html#refactoring-of-filesource-constructors-and-filescanconfigbuilder-to-accept-schemas-upfront

.build()
}
_ => get_file_config_builder(
required_schema,
partition_schema,
file_groups,
object_store_url,
file_source,
)
.build(),
};
let mut file_scan_config_builder =
FileScanConfigBuilder::new(object_store_url, file_source).with_file_groups(file_groups);

if let Some(projection_vector) = projection_vector {
file_scan_config_builder = file_scan_config_builder
.with_projection_indices(Some(projection_vector))?
.with_expr_adapter(Some(expr_adapter_factory));
}

let file_scan_config = file_scan_config_builder.build();

Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
}
Expand Down Expand Up @@ -165,28 +174,3 @@ fn get_options(

(table_parquet_options, spark_parquet_options)
}

fn get_file_config_builder(
schema: SchemaRef,
partition_schema: Option<SchemaRef>,
file_groups: Vec<FileGroup>,
object_store_url: ObjectStoreUrl,
file_source: Arc<dyn FileSource>,
) -> FileScanConfigBuilder {
match partition_schema {
Some(partition_schema) => {
let partition_fields: Vec<Field> = partition_schema
.fields()
.iter()
.map(|field| {
Field::new(field.name(), field.data_type().clone(), field.is_nullable())
})
.collect_vec();
FileScanConfigBuilder::new(object_store_url, Arc::clone(&schema), file_source)
.with_file_groups(file_groups)
.with_table_partition_cols(partition_fields)
}
_ => FileScanConfigBuilder::new(object_store_url, Arc::clone(&schema), file_source)
.with_file_groups(file_groups),
}
}
Loading