feat(python): Add DataFrame.write_iceberg

[kevinjqliu]

Resolves #14610

This PR introduces the write_iceberg function to the Dataframe API. write_iceberg requires a pyiceberg.table.Table object and a mode (either append or overwrite).

Note, partitioned write is currently not supported (blocked on iceberg-python #208)

Constructing the Iceberg catalog and Table object is left outside of this function. The test (test_write_iceberg) provides an example of creating an in-memory catalog and a Iceberg Table object.

Example

import polars as pl
df = pl.DataFrame(
    {
        "foo": [1, 2, 3, 4, 5],
        "bar": [6, 7, 8, 9, 10],
        "ham": ["a", "b", "c", "d", "e"],
    }
)

table_path = "/tmp/path/to/iceberg-table/"

# create catalog and table
from pyiceberg.catalog.sql import SqlCatalog
catalog = SqlCatalog(
    "default", uri="sqlite:///:memory:", warehouse=f"file://{table_path}"
)
catalog.create_namespace("foo")
table = catalog.create_table(
    "foo.bar",
    schema=df.to_arrow().schema,
)

df.write_iceberg(table, mode="overwrite")
print(table.location())

# read back the iceberg table 
pl.scan_iceberg(table).collect()

Files created

(.venv) ➜  py-polars git:(kevinjqliu/iceberg-write) tree /tmp/path/to/iceberg-table/                  
/tmp/path/to/iceberg-table/
└── default.db
    └── table
        ├── data
        │   └── 00000-0-3a2c5080-0e03-4dc5-a21a-73b8423ae181.parquet
        └── metadata
            ├── 00000-4ed15e3e-cafb-436c-b0ec-3bd08e6079ab.metadata.json
            ├── 00001-47503a3a-c2d9-4024-8605-28e92b6448c0.metadata.json
            ├── 3cd50fc2-83c6-4e5d-a68a-c34b56418631-m0.avro
            └── snap-3013071076850462756-0-3cd50fc2-83c6-4e5d-a68a-c34b56418631.avro

4 directories, 5 files

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.89%. Comparing base (dddf0b7) to head (d2f64d2).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #15018   +/-   ##
=======================================
  Coverage   79.88%   79.89%           
=======================================
  Files        1513     1513           
  Lines      203466   203471    +5     
  Branches     2892     2893    +1     
=======================================
+ Hits       162546   162555    +9     
+ Misses      40372    40369    -3     
+ Partials      548      547    -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kevinliu-stripe]

waiting for pyiceberg 0.6.1 release, should be released soon, which will fix the schema issue described above

[mkarbo]

@kevinliu-stripe fyi 0.6.1 has been released now :)

https://github.com/apache/iceberg-python/releases/tag/pyiceberg-0.6.1

[stinodego]

Thanks for the PR. A few comments.

[stinodego]

Why does this method take a Table and not a Pathlike, like other write functions?

[kevinjqliu]

Writing to an iceberg table requires creating the pyiceberg.table.Table object, which also requires an iceberg catalog.
I think it's best to delegate the creation of the catalog and table to be outside of this function and just pass in the table object.

The original implementation takes a str/Pathlike representing the "warehouse location" for the iceberg table. In Iceberg world, writing to a table usually requires an update to the catalog.

[stinodego]

This logic is so simple that this does not warrant its own method, in my opinion.

[kevinjqliu]

The main goal is to implement the top-level write_iceberg function for the dataframe API. The actual implementation code can be moved somewhere else

[stinodego]

The point is that, if users already have an Iceberg table object, they can just write table.append(df.to_arrow). It's even shorter than df.write_iceberg(table, mode="append"). So there is not much added value to a write_iceberg method.

If there is some complex logic required to set up the iceberg table, or if to_arrow is not sufficient to handle all data types correctly, we can consider adding a write_iceberg to save all users some implementation hassle. But right now it doesn't seem warranted.

[kevinjqliu]

That is true. Pyiceberg is well integrated with arrow. With an arrow dataframe and a PyIceberg table object, one can just invoke the pyiceberg write functions .append/.overwrite.

Given the above, is there still value in implementing a simple write_iceberg method?

Looking at the write_delta method, aside from the merge functionality, the function just pass data as arrow into the write_deltalake function.

[kevinjqliu]

If there is some complex logic required to set up the iceberg table

There's a scenario where a user might want to write an Iceberg table to a location in blob store. In such case, the write_iceberg function can take care of creating an in-memory catalog and iceberg table object before writing.

That was the initial version of this PR
0801012

[glesperance]

Seems like a write_iceberg would be worth it just by virtue of having full support from polars / parity with delta.

[stinodego]

This needs much more extensive testing. Try using the df fixture (it has a mix of many data types), or define your own dataframe with lots of different data types.

[kevinjqliu]

Thanks! I'll take a look at this fixture

[stinodego]

See comments. Needs testing on more different data types before this can be merged.

[glesperance]

This is great. Excited to try this out

[mkarbo]

Excited to try this one out!

[kevinjqliu]

hey @stinodego thanks for the previous review, could you take another look at this PR?

I've changed the test to use the df test fixture.
Currently, pyiceberg does not support time64[ns] and returning the dtype as Enum or Categorical. Hopefully, we can add those incrementally.