Upsertion memory usage grows exponentially as table size grows

[mwa28]

Apache Iceberg version

0.9.0

Please describe the bug 🐞

Hello,

I am trying out the new upsert method on a table created as follows on the AWS Glue Catalog :

catalog = load_catalog("glue", **{"type": "glue"})
schema = Schema(
    NestedField(1, "dt_insert", StringType(), required=True),
    NestedField(2, "controller_id", StringType(), required=True),
    NestedField(3, "timestamp", StringType(), required=True),
    NestedField(4, "data_type", StringType(), required=True),
    NestedField(5, "parameter", StringType(), required=True),
    NestedField(6, "value", StringType(), required=True),
    NestedField(7, "data_source", StringType(), required=True),
    NestedField(8, "unique_id", StringType(), required=True),
    NestedField(9, "dt_import_utc", StringType(), required=True),
    identifier_field_ids=[8],  # 'unique_id' is the primary key
)
catalog.create_table(
        identifier=<database>.<table_name>,
        schema=schema,
        partition_spec=PartitionSpec(
            PartitionField(
                source_id=9,
                field_id=1000,
                transform=IdentityTransform(),
                name="dt_import_utc",
            ),
        ),
        sort_order=SortOrder(
            SortField(source_id=2),  # controller_id
            SortField(source_id=4),  # data_type
            SortField(source_id=3),  # dt_insert
            SortField(source_id=1),  # timestamp
        ),
        location="s3://bucket/prefix/catalog",
        properties={
      "write.format.default" : "parquet",
      "write.target-file-size-bytes" : 134217728, # 128 MB
      "write.metadata.delete-after-commit.enabled" : True,
      "write.metadata.previous-versions-max" : 5
    })

unique_id is a the merged string of columns 1, 2, 3, 4 and 5 with an underscore.

My upsertion is a simple pandas dataframe dumped into a dict and being transformed to pyarrow for upsertion (I know pyarrow accepts directly a pandas df but doing this for an internal reason that requires writing the df in a stage and read in another, so i am dumping as json for readability).

iceberg_table = catalog.load_table(table_name)
table =  pa.Table.from_pydict(to_upsert, schema=iceberg_table.schema().as_arrow())
iceberg_table.upsert(df=table, join_cols=["unique_id"])

With AWS Athena, a MERGE INTO statement takes about 3 seconds to run on the table and scans 1.11MB of data before completion.

MERGE INTO <database>.<table> target
            USING <database>.<temp_table> source
            ON (target."unique_id" = source."unique_id")
            WHEN MATCHED THEN
                UPDATE SET "dt_insert" = source."dt_insert", "controller_id" = source."controller_id", "timestamp" = source."timestamp", "data_type" = source."data_type", "parameter" = source."parameter", "value" = source."value", "data_source" = source."data_source", "dt_import_utc" = source."dt_import_utc", "unique_id" = source."unique_id"
            WHEN NOT MATCHED THEN
                INSERT ("dt_insert", "controller_id", "timestamp", "data_type", "parameter", "value", "data_source", "dt_import_utc", "unique_id")
                VALUES (source."dt_insert", source."controller_id", source."timestamp", source."data_type", source."parameter", source."value", source."data_source", source."dt_import_utc", source."unique_id")

(statement generated by using the to_iceberg from awswrangler)

Meanwhile when i try with pyiceberg upsert, it is using more than 10240MB. I am running on AWS lambda and it is causing an out of memory error.

I have no issue with the append() function, it completes fairly quickly but it seems that the upsert needs further optimization to be able to efficiently retrieve only relevant data.

Current table size is at 18.5GB for both Athena based statement and pyiceberg upsert function call.

Willingness to contribute

• I can contribute a fix for this bug independently
• I would be willing to contribute a fix for this bug with guidance from the Iceberg community
• I cannot contribute a fix for this bug at this time

[mwa28]

From what i was able to gather, the uspert() can benefit from partition awareness to enable a partial table scan instead of a full table scan.

An AI-generated code fix would look like as follows :

def upsert(
    self,
    df: pa.Table,
    join_cols: Optional[List[str]] = None,
    when_matched_update_all: bool = True,
    when_not_matched_insert_all: bool = True,
    case_sensitive: bool = True,
) -> UpsertResult:
# Get partition columns from the table spec
    partition_cols = [field.name for field in self.spec().fields]

    # Partition pruning: filter matched_iceberg_table to only relevant partitions
    if partition_cols:
        partition_values = df.select(partition_cols).to_pandas().drop_duplicates()
        mask = None
        for _, row in partition_values.iterrows():
            cond = functools.reduce(
                operator.and_,
                [pc.field(col) == row[col] for col in partition_cols]
            )
            mask = cond if mask is None else (mask | cond)
        if mask is not None and len(matched_iceberg_table) > 0:
            matched_iceberg_table = matched_iceberg_table.filter(mask)

# ...rest of upsert logic...

I have not reviewed what the AI generated code, it could be complete gibberish but i do agree on that there is a lack of partition awareness in the function which i thought would have been pre-baked into an iceberg table but it doesn't seem like it at the moment. Maybe a better idea would be to leverage the filesystem beneath similar to how spark does or pushdown predicates, that can access only the required folder/partition ?

[Fokko]

@mwa28 Thanks for raising this. 10240 is a lot of memory, and doesn't sound right.

Luckily @koenvo ran into the same issue, and pushed a fix, which got merged recently: #1995. Any chance to try out that implementation? It hasn't been released, but it would be good to try it out in production cases.

[koenvo]

@mwa28 Thanks for raising this. 10240 is a lot of memory, and doesn't sound right.

Luckily @koenvo ran into the same issue, and pushed a fix, which got merged recently: #1995. Any chance to try out that implementation? It hasn't been released, but it would be good to try it out in production cases.

Curious how these changes impact the memory usage in your case.

Would it be possible to add the partition key to your join_keys? Maybe we could always add partition keys when not passed as join_keys yet..?

[mwa28]

Curious how these changes impact the memory usage in your case.

Would it be possible to add the partition key to your join_keys? Maybe we could always add partition keys when not passed as join_keys yet..?

Sorry i haven't had the time to look at this yet, kinda time consuming for me to build an non-released version for my lambda at the moment, i need some time before i can get back to it.

I like the idea of adding the partition_key to the join_keys by default but i think this should be an intrinsic operation rather than giving the user the ability to override the value if set by default and then having to ask them to keep the partition column in the join keys.

I think your modification should be onto something, i will keep you posted once i have the time to test it myself.