dlt-hub / dlt-hub/dlt

OOM (Memory Leak/High Consumption) when extracting millions of Parquet files from S3 via Arrow batches

Open
#3,544 2 comments 0 reactions 1 assignee Claimed by @rudolfix View on GitHub
question
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.

Image

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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.