When pyiceberg loads Iceberg tables containing large JSON data into memory, memory usage explodes in pyarrow

[thomas-pfeiffer]

Apache Iceberg version

0.11.0 (latest release)

Please describe the bug 🐞

We have a table with a column containing large JSON strings (basically think of large pydantic validation results). The resulting datafile in Iceberg itself is a roughly 6MB (GZip compressed) Parquet file, but when querying it, the memory consumption goes to ~4GB. (Just for a few seconds, but long enough to cause out-of-memory issues on some systems.)

The issue is that by default pyarrow loads the strings per row into memory, which blows up the memory. If we download the datafile and open it directly via pyarrow this behaviour can be reproduced.

There are 2 workarounds in pyarrow: 1. Don't load the problematic column (given that is possible in your use case) and 2. switch to dictionary-encoding for set column (example snippet below).

from pyarrow.parquet import read_table
table = read_table("datafile.parquet", read_dictionary=["problematic_column"])

Issue in pyiceberg:
Regardless if you use table.scan(...).to_arrow() or table.scan(...).to_arrow_batch_reader(), pyiceberg has afaik currently no option to specify the dictionary encoding for certain tables, hence pyarrow uses the default encoding and the memory usage explodes.

The to_arrow_batch_reader does not help here either, because -as per my understanding- in the batch reader of pyiceberg each batch represents an individual datafile. Hence, if there is one problematic 6MB datafile, it makes no difference if you use the batch reader or not. I also have the impression that when you iterate over the reader, pyarrow has already loaded the parquet file in a separate thread and this is where the memory explosion actually happens.

So the current only workaround in pyiceberg is option 1: Don't load the problematic column by specifying the selected_fields:

from pyiceberg.catalog import load_catalog

catalog = load_catalog("default")
table = catalog.load_table("your_table")
reader = table.scan(selected_fields=("all", "other", "columns")).to_arrow_batch_reader()
for batch in reader:
...

Expected behaviour:
There should be an option somewhere, e.g. in the data_scan to specify for which columns dictionary encoding should be used. This option should be forwarded to pyarrow internally somehow, so that pyarrow uses less memory.

Remark:
I would not change the default behaviour. It would be just good to have the option to configure the encoding in pyarrow when needed.

This issue is a follow up for #1205

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

[kevinjqliu]

Thanks for reporting this!

The to_arrow_batch_reader does not help here either, because -as per my understanding- in the batch reader of pyiceberg each batch represents an individual datafile. Hence, if there is one problematic 6MB datafile, it makes no difference if you use the batch reader or not. I also have the impression that when you iterate over the reader, pyarrow has already loaded the parquet file in a separate thread and this is where the memory explosion actually happens.

This is a bug. I've found out about this recently (see #3122). to_arrow_batch_reader should read from a single parquet file 1 batch at a time; thus reducing the memory consumption.

#2676 is the proper fix to correct this behavior. Would love to see if this helps with your issue.

There should be an option somewhere, e.g. in the data_scan to specify for which columns dictionary encoding should be used. This option should be forwarded to pyarrow internally somehow, so that pyarrow uses less memory.

I think thats a reasonable feature request. Ive opened #3170 to track this

[kevinjqliu]

The issue is that by default pyarrow loads the strings per row into memory, which blows up the memory. If we download the datafile and open it directly via pyarrow this behaviour can be reproduced.

👍 #2676 solves the memory issue for reading multiple batches. If reading a single batch blows up the memory, the only fix here is to use dictionary-encoding

[thomas-pfeiffer]

Adding here the documentation for the read_dictionary option of pyarrow: https://arrow.apache.org/docs/python/parquet.html#reading-types-as-dictionaryarray

I also think that the place where this additional parameter option would be needed is at:

iceberg-python/pyiceberg/io/pyarrow.py

Line 1100 in 536a7d0

return ds.ParquetFileFormat(**kwargs)

I guess we would need the possibility to add something like the statement below - "just" passed through all of pyiceberg:

ds.ParquetFileFormat(read_options=ds.ParquetReadOptions(dictionary_columns=["column1", "column2"]))