I have a huge table (df1, ~1.5 billion rows) stored in a S3 bucket, divided in multiple parts or partitions with the Parquet extension. My goal is to filter it by keeping those rows where the value of a particular column exists in a column (with the same name) of another table (df2). I'm using Dask for this operation since the first table/DataFrame doesn't fit into memory.
This is what I've tried so far:
df1 = dd.read_parquet(f"s3://{MY_BUCKET}/{PATH_TO_HUGE_TABLE}/*", index=INDEX_COL)
df2 = dd.read_parquet(f"s3://{MY_BUCKET}/{PATH_TO_AUX_TABLE}/*", index=INDEX_COL, columns=A_SINGLE_COLUMN_LIST)
filtered_df1 = df1.merge(df2, left_index=True, right_index=True, how='inner')
filtered_df1 = filtered_df1.reset_index(drop=False)
filtered_df1.to_parquet(f"s3://{MY_BUCKET}/{PATH}", write_index=False)
The code above runs without problems till the workers start running out of memory and eventually the execution is aborted. I'm working with a quite powerful machine of 96 CPUs and more that 700GB of memory (and yet the huge table doesn't fit). I've tried to run it with 96 and 8 workers, but it happens the same in both cases.
I've also tried the following and then included the isin() method in a map_partitions() function. Both approaches had the same outcome that I explained above.
df2 = dd.read_parquet(f"s3://{MY_BUCKET}/{PATH_TO_AUX_TABLE}/*")
filter_values = df2[INDEX_COL].unique().compute() # len(filter_values) --> ~17e6
df1 = dd.read_parquet(f"s3://{MY_BUCKET}/{PATH_TO_HUGE_TABLE}/*")
filtered_df1 = df1[ df1[INDEX_COL].isin(filter_values) ]
filtered_df1.to_parquet(f"s3://{MY_BUCKET}/{PATH}", write_index=False)
Does anyone know another way of doing this with Dask that could work? Perhaps a different Python library? Any idea is appreciated, thank you.