OOM (Memory Leak/High Consumption) when extracting millions of Parquet files from S3 via Arrow batches
- Dominant language
- Python
- Stars
- 5.9k
- Forks
- 605
- Avg merge
- 1d 14h
- Merged PRs (30d)
- 38
Description
### dlt version
1.20.0
### Describe the problem
I am using dlt for load of Parquet files from S3 to BigQuery. The bucket is pretty big (over milion files).
I need to do an incremental load. File sizes differe between 20 kB up to 200 MB.
I also want to preserve file url and hive partitions encoded in that url as additional columns in my target table. However, I don't think this matters in this case.
Here are my utils for extracting hive partitions from file paths and adding them as columns in the destination (any better idea how to do it?):
```python
import pyarrow as pa
import pyarrow.dataset as ds
from pyarrow import parquet as pq
def extract_partitions(item: FileItemDict, schema: pa.Schema) -> FileItemDict:
partitioning = ds.HivePartitioning(schema)
expression = partitioning.parse(item["file_url"])
partitions = ds.get_partition_keys(expression)
item["partitions"] = partitions
logger.debug("Extracted partitions: %r", partitions)
return item
def _add_scalar_column(
batch: pa.RecordBatch, column_name: str, value, value_type: pa.DataType
) -> pa.RecordBatch:
"""Efficiently add a scalar column without creating intermediate list."""
scalar = pa.scalar(value, type=value_type)
column = pa.repeat(scalar, len(batch))
logger.debug("Adding column: %r", column_name)
return batch.append_column(column_name, column)
def _add_extra_columns(batch, columns: dict) -> pa.RecordBatch:
for column_name, column_info in columns.items():
batch = _add_scalar_column(
batch, column_name, column_info["value"], column_info["value_type"]
)
return batch
@dlt.transformer
def read_parquet_and_add_hive_columns(
files: Iterable[FileItemDict],
columns_mapper: Callable[[FileItemDict], dict],
batch_size: int = 1000,
) -> Iterator[TDataItems]:
for file in files:
file_url = unquote(file["file_url"])
columns = columns_mapper(file)
logger.info(
"Open file: %s [%s]", file_url, humanize.naturalsize(file["size_in_bytes"])
)
with (
file.open() as source_file,
pq.ParquetFile(source_file, pre_buffer=True) as parquet_file,
):
for batch in parquet_file.iter_batches(batch_size=batch_size):
yield _add_extra_columns(batch, columns)
logger.info(
"Finished extracting file: %s [%s]",
file_url,
humanize.naturalsize(file["size_in_bytes"]),
)
```
Here's my resource definition:
```python
PARTITION_SCHEMA = pa.schema(
[
("projectid", pa.string()),
("year", pa.int32()),
("month", pa.int32()),
("day", pa.int32()),
]
)
def columns_mapper(file: FileItemDict) -> dict:
partitions = file["partitions"]
return {
"project_id": {
"value": partitions["projectid"],
"value_type": pa.string(),
},
"event_date": {
"value": date(partitions["year"], partitions["month"], partitions["day"]),
"value_type": pa.date32(),
},
"file_url": {
"value": file["file_url"],
"value_type": pa.string(),
},
}
@dlt.resource
def event_store(bucket_url: str, event: str, last_ts: datetime = None) -> DltResource:
logger.info("Scan path: %s", path)
files = filesystem(
credentials=dlt.secrets["sources.event_store.credentials"],
bucket_url=path,
file_glob="**/*.parquet",
incremental=dlt.sources.incremental(
"modification_date", initial_value=INITIAL_TS
),
files_per_page=10,
)
if last_ts:
logger.info("Setting upper bound for modification_date to: %s", last_ts)
files = files.add_filter(lambda r: r["modification_date"] < last_ts)
files = files.add_map(
lambda item: extract_partitions(item, schema=PARTITION_SCHEMA)
)
return files | read_parquet_and_add_hive_columns(columns_mapper=columns_mapper)
```
I am using a transformer to read the files as Arrow batches and inject Hive partitioning data (extracted from the file path) as extra columns - I'm not sure if this is the only way.
Despite using generators and processing data in batches, the memory consumption in the extract phase is surprisingly high.
It seems the memory consumption grows when the size of the files that are being extracted increases. But I'd expect it to be steady, considering I'm iterating using batches.
I had to set the number of extract workers to 1 and max parallel items as well:
```
EXTRACT__WORKERS: 1
EXTRACT__MAX_PARALLEL_ITEMS: 1
```
This makes the whole process pretty slow, and still, consuming up to 2GB of memory seems a lot for what it is doing.
Increasing number of extract workers or parallel items causes my PODs being terminated with OOM error. I can't afford to declare memory request to 4GBs. All in all, I'm mostly just move data from S3 to Bigquery.
Any ideas where I can be doing things wrong?
### Expected behavior
_No response_
### Steps to reproduce
Run the pipeline that extracts significant number of parquet files from S3 to any destination.
Memory consumption is unreasonably high.
### Operating system
Linux
### Runtime environment
Kubernetes
### Python version
3.13
### dlt data source
S3
### dlt destination
Google BigQuery
### Other deployment details
_No response_
### Additional information
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.