Support partitioned writes

[Fokko]

Feature Request / Improvement

Support partitioned writes

So I think we want to tackle the static overwrite first, and then we can compute the predicate for the dynamic overwrite to support that. We can come up with a separate API. I haven't really thought this trough, and we can still change this.
I think the most important steps are the breakdown of the work. There is a lot involved, but luckily we already get the test suite from the full overwrite.

Steps I can see:

• Extend the summary generation to support partitioned writes (here in Java)
• Add support for the append files.
  • How are we going to fan out the writing of the data. We have an Arrow table, what is an efficient way to compute the partitions and scale out the work. For example, are we going to sort the table on the partition column and do a full pass through it? Or are we going to compute all the affected partitions, and then scale out?
• Add support for static overwrites
• Add support for dynamic overwrites

Other things on my mind:

• In Iceberg it can be that some files are still on an older partitioning, we should make sure that we handle those correctly based on the that we provide.
• How to handle delete files; it might be that the delete files become unrelated because the affected datafiles are replaced. We could first ignore this.

The good part:

• In PyIceberg we're first going to ignore the fast-appends (this is when you create a new manifest, and add it to the manifest list). Instead we'll just take the existing manifest(s) and rewrite it into a single new manifest which makes it a bit easier to reason about the snapshot (and therefore the snapshot summaries). The reason is that this caused quite a few bugs in Java, and it can be added always on a later moment.

[jqin61]

Hi @Fokko and Iceberg community, I and @syun64 are continuing working on testing the write capability in Write support pr. We are excited about it as it will help us with our use case a lot.

Our use case also includes overwriting partitions of tables. So I am highly interested in opportunities to contribute to this issue. Would it be alright for me to start working on this issue, based on the Write support pr if no one else has already begun?

[Fokko]

Hey @jqin61 Thanks for replying here. I'm not aware of the fact that anyone already started on this. It would be great if you can take a stab at it 🚀

[jqin61]

In Iceberg it can be that some files are still on an older partitioning, we should make sure that we handle those correctly based on the that we provide.

It seems Spark's iceberg support has such overwrite behaviors under partition scheme evolution:

• dynamic overwrite: data files generated from old partition spec will not be replaced even if some of the records match the overwriting data
• static overwrite with PARTITION values specified: same as above
• static overwrite without PARTITION values: all data is deleted regardless of what partition specs they conform to.

As Fokko mentioned, we need to make sure in the implementation we use the latest partition spec_id when overwriting partitions so that the data in the old partition spec is not touched.

[jqin61]

How are we going to fan out the writing of the data. We have an Arrow table, what is an efficient way to compute the partitions and scale out the work. For example, are we going to sort the table on the partition column and do a full pass through it? Or are we going to compute all the affected partitions, and then scale out?

It just comes to me that when spark writes to iceberg, it requires the input dataframe to be sorted by the partition value otherwise an error will be raised during writing. Do we want to take the same assumption for pyiceberg?

If not, if we have to use arrow.compute.filter() to extract each partition before serialization, it seems a global sort before the filter() on the entire table is unnecessary since the filter makes no assumption of the array organization?

To extract the partitions by filter(), would it be helpful if we firstly build an API in pyarrow which does a full scan of the array and bucket-sorts it into partitions and returns buckets (partitions) as a list of arrow arrays? These arrays could be further passed as input to writing jobs which are executed in a multi-threading way.

[Fokko]

I currently see two approaches:

• First get the unique partitions, and then filter for each of the partitions the relevant data. It is nice that we know the partition upfront, but the downside is that we need to do a pass through the dataframe to get the relevant rows for that partition.
• Sort the entire dataframe on the partition columns, which is expensive, but we need to do this only once. Once it is sorted, we can just do a single pass through the dataframe. If we want to scale out, we can split out the rows: [0, n/2], [n/2, 2].

I'm not sure what is the best. I think the first one works better if you have few partitions, and the latter one is more efficient when you have many partitions.

To extract the partitions by filter(), would it be helpful if we firstly build an API in pyarrow which does a full scan of the array and bucket-sorts it into partitions and returns buckets (partitions) as a list of arrow arrays? These arrays could be further passed as input to writing jobs which are executed in a multi-threading way.

Starting with the API is always a great idea. My only concern is that we make sure that we don't take copies of the data, since that might blow up the memory quite quickly.

Hope this helps!

[Fokko]

@jqin61 I did some more thinking over the weekend, and I think that the approach that you suggested is the most flexible. I forgot about the sort-order that we also want to add at some point. Then we would need to sort twice 👎

[jqin61]

Based on the existing discussion, there are 3 major possible directions for detecting partitions and writing each partition in a multi-threaded way to maximize I/O. It seems there isn’t any approach simple enough that we could purely leverage the existing Pyarrow APIs in Pyiceberg. I marshalled Fokko's suggestions and list these approaches for discussion purpose:

Filter Out Partitions
As Fokko suggested, we could filter the table to get partitions before writing but we will need an API on Arrow to get unique partition values (e.g. extend compute.unique() from array/scalar to table).

partitions: list[dict] = pyarrow.compute.unique(arrow_table)

With it, we could filter the table to get partitions and provide them as inputs to concurrent jobs.

arrow_table = pa.table({
    'column1': ['A', 'B', 'A', 'C', 'B'],
    'column2': [1, 2, 1, 3, 2],
    'column3': ['X', 'Y', 'X', 'Z', 'Y']
})
partition_keys = table.select(['column1', 'column2'])
# The existing unique does not have support on table, we need to create API on Arrow side.
partitions: list[dict] = pyarrow.compute.unique(partition_keys)
# = [
    {'column1': 'A', 'column2': 1},
    {'column1': 'B', 'column2': 2},
    {'column1': 'C', 'column3': 3}
]


def filter_and_write_table(partition, index):
    # Create a boolean mask for rows that match the criteria
    mask = pc.and_(*(pc.equal(table[col], val) for col, val in partition.items()))
    
    # Filter the table
    filtered_table = table.filter(mask)

    # Write the filtered table to a Parquet file
    parquet_file = f'output_partition_{index}.parquet'
    pq.write_table(filtered_table, parquet_file)

with ThreadPoolExecutor() as executor:
    for i, partition in enumerate(partitions):
        executor.submit(filter_and_write_table, partition, i)

Sort and Single-direction Writing
As Fokko suggested, we could sort the table first. We then slice the table and do a one-direction scan for each slice to write out partitioned files.
Suppose we have such an arrow API that takes a sorted table, scans through it, creates a new file whenever encountering a row with a new partition value, and raises an error if it encounters a row with a partition value it already passes, just like how spark writes to an iceberg table.

def write_table_partitions(sorted_table, partition_columns, dir)

Then we could do

partition_columns = ['column1', 'column2']
sorted_table = table.sort_by([('column1', 'ascending'), ('column2', 'ascending')])
directory_path = '/path/to/output/directory'

# Break down the sorted table to slices with zero-copy
slices = slice_table(sorted_table, slice_options)

# Call the API
with ThreadPoolExecutor() as executor:
    # Submit tasks to the executor
    for i, partition in enumerate(partitions):
        executor.submit(write_table_partitions, sorted_table, partition_columns, dir)

Bucketing
We could create an arrow API to return the partitioned tables/record batches based on the inputs of a table and alist of partition columns in a way that the algorithm does a full scan of the arrow table in O(table_length) time and bucket-sorts it and creates a table/record batch for each bucket:

table_partitions = pyarrow.compute.partition(arrow_table, partition_columns)

We could write each batch:

def write_table_to_parquet(table, directory_path, file_suffix):
    # Construct the full path for the Parquet file
    file_path = os.path.join(directory_path, f'record_batch_{file_suffix}.parquet')

    # Write the table to a Parquet file
    pq.write_table(table, file_path)

with ThreadPoolExecutor() as executor:
    for i, partition_as_table in enumerate(table_partitions):
        executor.submit(write_table_to_parquet, partition_as_table, directory_path, i)

As Fokko pointed out, the filter method will not be efficient if there are many partitions - the filter takes O(table_length) time and although each thread can filter on its own, on a single node, the execution will be O(table_length * number_of_partitions) for all the jobs. Technically we only need one same scan to get all the buckets.

It seems the sort method is not as efficient compared to the bucketing method because the relative order of partitions does not matter, so a general sort algorithm on the partition column might be overkill (compared with bucketing).

I feel like all 3 directions require some implementation on Arrow itself (I did not find any approach simple enough that we could purely leverage the existing Pyarrow APIs to implement any of the directions). And I want to get opinions on whether pursuing arrow API level utilities smells good. Thank you!

Specifically, for the third direction of bucketing and returning materialized tables/batches, since Arrow has dataset.write_dataset() which supports partition-respected writing, I did some reading to see how it partitions and whether we could leverage anything from it.

https://github.com/apache/arrow/blob/main/cpp/src/arrow/dataset/partition.cc#L118 is where the partition happens. The partition algorithm is a full scan with bucket sort leveraging Grouper class utilities in arrow's compute component. Specifically:
1.Grouper.consume() initiates the groups based on the presenting partition columns
https://github.com/apache/arrow/blob/55afcf0450aa2b611e78335bdbfd77e55ae3bc9f/cpp/src/arrow/compute/row/grouper.cc#L422
2.Grouper.MakeGroupings() builds a ListArray where each list represents a partition and each element in the list represents the row_id of the original table.
https://github.com/apache/arrow/blob/55afcf0450aa2b611e78335bdbfd77e55ae3bc9f/cpp/src/arrow/compute/row/grouper.cc#L886C45-L886C58
3.Grouper.ApplyGroupings() efficiently converts the grouped representation of row_ids into actual rows.
https://github.com/apache/arrow/blob/55afcf0450aa2b611e78335bdbfd77e55ae3bc9f/cpp/src/arrow/compute/row/grouper.cc#L875

Other than being used in the dataset writing, Grouper from Arrow's compute component is used to support other exposed compute APIs such as aggregation functions. At the end of the day, what we want (in order to support Pyiceberg's partitioned write) is an API that returns record batches/tables based on an input table and an input partition scheme, so maybe we could expose such a new API under compute leveraging Grouper.

[asheeshgarg]

@jqin61 just wondering if we can use this directly
https://arrow.apache.org/docs/python/generated/pyarrow.dataset.partitioning.html

[jqin61]

@jqin61 just wondering if we can use this directly https://arrow.apache.org/docs/python/generated/pyarrow.dataset.partitioning.html

Thank you Ashish! I overlooked it, as you mention, we could just use write_dataset() with specified args of partitioning base_nametemplate to write out the partitioned datafiles as iceberg needs.

[asheeshgarg]

@Fokko @jqin61 I am also interested in this to move forward as we also deal with lot of write involves partitions. Happy to collaborate on to this.
For write_dataset() we might need to look if we need to add field meta data at parquet in terms of field-id etc while committing data files.

[sungwy]

@Fokko @jqin61 I am also interested in this to move forward as we also deal with lot of write involves partitions. Happy to collaborate on to this. For write_dataset() we might need to look if we need to add field meta data at parquet in terms of field-id etc while committing data files.

Yes - I think we learned this from our earlier attempts: https://github.com/apache/iceberg-python/pull/41/files/1398a2fb01341087a1334482db84a193843a2362#r1427302782

As @jqin61 pointed out in a previous PR, adding these to the schema should output parquet files with the correct field_id.

[asheeshgarg]

@Fokko @jqin61
Today I tried basic example on partition write:

from pyiceberg.io.pyarrow import schema_to_pyarrow
import pyarrow as pa
from pyarrow import parquet as pq
data = {'key': ['001', '001', '002', '002'],
        'value_1': [10, 20, 100, 200],
        'value_2': ['a', 'b', 'a', 'b']}
my_partitioning = pa.dataset.partitioning(pa.schema([pa.field("key", pa.string())]), flavor='hive')
TABLE_SCHEMA = Schema(
    NestedField(field_id=1, name="key", field_type=StringType(), required=False),
    NestedField(field_id=2, name="value_1", field_type=StringType(), required=False),
    NestedField(field_id=3, name="value_2", field_type=StringType(), required=False),
)
schema = schema_to_pyarrow(TABLE_SCHEMA)
patbl = pa.Table.from_pydict(data)
pq.write_to_dataset(patbl,'partitioned_data',partitioning=my_partitioning,schema=schema)

If I don't use schema in write it works fine. But if I pass the schema create

schema = schema_to_pyarrow(TABLE_SCHEMA)

It fails with

ArrowTypeError: Item has schema
key: string
value_1: int64
value_2: string
which does not match expected schema
key: string
  -- field metadata --
  PARQUET:field_id: '1'
value_1: string
  -- field metadata --
  PARQUET:field_id: '2'
value_2: string
  -- field metadata --
  PARQUET:field_id: '3'

I also tried the parquet write the way we are doing currently:

writer = pq.ParquetWriter("test", schema=schema, version="1.0") 
writer.write_table(patbl)
ValueError: Table schema does not match schema used to create file: 
table:
key: string
value_1: int64
value_2: string vs. 
file:
key: string
  -- field metadata --
  PARQUET:field_id: '1'
value_1: string
  -- field metadata --
  PARQUET:field_id: '2'
value_2: string
  -- field metadata --
  PARQUET:field_id: '3

Do we do any other transformation for the schema before we write in current write support.

[Fokko]

Hey @jqin61

Thanks for the elaborate post, and sorry for my slow reply. I did want to take the time to write a good answer.

Probably the following statement needs another map step:

partitions: list[dict] = pyarrow.compute.unique(arrow_table)

The above is true for an identity partition, but often we take truncate the month, day or hour from a field, and use that as a partition. Another example is the bucketing partition where we hash the field, and determine in which bucket it will fall.

With regard of utilizing the Arrow primitives that are already there. I think that's a great idea, we just have to make sure that they are flexible enough for Iceberg. There are a couple of questions that pop into my mind:

• Can we support all Icebergs partition strategies, such as bucketing, truncating etc.
• Are we able to extract the metrics similar that we do for non-partitioned writes.

@asheeshgarg Thanks for giving it a try. Looking at the schema, there is a discrapency. The test-data that you generate has value_1 as an int64, and the table expects a string. I think the error is correct here.

[asheeshgarg]

@Fokko thanks for pointing out the mismatch it worked. After modifying the datatype it worked.