Skip to content

feat(datafusion): Support insert_into in IcebergTableProvider #1511

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.

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
wants to merge 35 commits into
base: main
Choose a base branch
from

Conversation

CTTY
Copy link
Contributor

@CTTY CTTY commented Jul 15, 2025

Which issue does this PR close?

What changes are included in this PR?

Are these changes tested?

@@ -440,10 +440,12 @@ impl PartnerAccessor<ArrayRef> for ArrowArrayAccessor {
Ok(schema_partner)
}

// todo generate field_pos in datafusion instead of passing to here
Copy link
Contributor Author

Choose a reason for hiding this comment

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

I found it tricky to handle this case: the input from datafusion won't have field id, and we will need to assign them manually. maybe there is a way to do name mapping here?

Copy link
Contributor

Choose a reason for hiding this comment

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

Could you help me to understand why we need to change this?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is a temporary hack to an issue that I don't know how exactly to fix for now: the RecordBatch from Datafusion won't have PARQUET_FIELD_ID_META_KEY in its schema's metadata, causing the schema visiting to fail here.

I'm thinking maybe we can bound the schema in datafusion via name mapping, but have not got the chance to explore more

Copy link
Contributor

Choose a reason for hiding this comment

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

Why we need to convert RecordBatch's schema to iceberg schema?

Copy link
Contributor

Choose a reason for hiding this comment

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

The method you mentioned is typically used to convert parquet file's schema to iceberg schema.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This method is used when using ParquetWriter to write RecordBatch. When it's counting nan values, it will need to walk through both RecordBatch's schema and Iceberg schema in a partner fashion:

.compute(self.schema.clone(), batch_c)?;

Basically the call stack is NanValueCountVisitor::compute -> visit_struct_with_partner -> ArrowArrayAccessor::field_partner -> get_field_id

@CTTY CTTY force-pushed the ctty/df-insert branch from 7843b0d to 2f9efa8 Compare July 16, 2025 03:37
Copy link
Contributor

@liurenjie1024 liurenjie1024 left a comment

Choose a reason for hiding this comment

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

Thanks @CTTY for this pr, just finished round of review. My suggestion is to start with unpartitioned table first.

@@ -432,3 +433,69 @@ async fn test_metadata_table() -> Result<()> {

Ok(())
}

#[tokio::test]
async fn test_insert_into() -> Result<()> {
Copy link
Contributor

Choose a reason for hiding this comment

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

I'm not a big fan of adding this kind of integration tests. How about adding sqllogictests?


// Create data file writer builder
let data_file_writer_builder = DataFileWriterBuilder::new(
ParquetWriterBuilder::new(
Copy link
Contributor

Choose a reason for hiding this comment

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

This should be RollingFileWriter

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Good point! Added rolling_writer.rs

Comment on lines +252 to +258
// // Apply the action and commit the transaction
// let updated_table = action
// .apply(tx)
// .map_err(to_datafusion_error)?
// .commit(catalog.as_ref())
// .await
// .map_err(to_datafusion_error)?;
Copy link
Contributor

Choose a reason for hiding this comment

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

Why comment out this?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is for testing only and will be uncommented in the future. I commented it out for now because the memory catalog doesn't support update table for now and this will fail my local test

Comment on lines +162 to +163
println!("----StructArray from record stream: {:?}", struct_arr);
println!("----Schema.as_struct from table: {:?}", schema.as_struct());
Copy link
Contributor

Choose a reason for hiding this comment

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

We should use log here.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is for testing only, and I'm planning to remove these log lines

@@ -440,10 +440,12 @@ impl PartnerAccessor<ArrayRef> for ArrowArrayAccessor {
Ok(schema_partner)
}

// todo generate field_pos in datafusion instead of passing to here
Copy link
Contributor

Choose a reason for hiding this comment

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

Could you help me to understand why we need to change this?

Copy link
Contributor

@liurenjie1024 liurenjie1024 left a comment

Choose a reason for hiding this comment

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

Thanks @CTTY , the direction looks good to me!

@@ -440,10 +440,12 @@ impl PartnerAccessor<ArrayRef> for ArrowArrayAccessor {
Ok(schema_partner)
}

// todo generate field_pos in datafusion instead of passing to here
Copy link
Contributor

Choose a reason for hiding this comment

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

Why we need to convert RecordBatch's schema to iceberg schema?

@@ -440,10 +440,12 @@ impl PartnerAccessor<ArrayRef> for ArrowArrayAccessor {
Ok(schema_partner)
}

// todo generate field_pos in datafusion instead of passing to here
Copy link
Contributor

Choose a reason for hiding this comment

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

The method you mentioned is typically used to convert parquet file's schema to iceberg schema.


// Verify each serialized file contains expected data
for json in &serialized_files {
assert!(json.contains("path/to/file"));
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: Why not assert the json output? We could use snapshot test to make it easier, see https://docs.rs/expect-test/latest/expect_test/

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I think a snapshot test makes more sense

/// This trait extends `IcebergWriter` with the ability to determine when to start
/// writing to a new file based on the size of incoming data.
#[async_trait]
pub trait RollingFileWriter: IcebergWriter {
Copy link
Contributor

Choose a reason for hiding this comment

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

Sorry I forgot that currently we don't have RollingFileWriter. It's fine to leave it in this pr, but it would be better to use a separate pr to add this. RollingFileWriter should be a FileWriter rather an IcebergWriter.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants