[Python] Python memory error and dask dying while joing on dataframes

I am currently working on processing a huge data and my code is continuously falling into memory error. I have tried the code in both pandas and dask. I am not sure if is it because of the logic i am using and hence I am providing with the full codes. Please help me with some inputs. Thanks.

Note: The codes and data are not the real ones.

Below is the use case:

I need to enrich a very large main dataset with fields from a
similarly large lookup table, based on a shared key, but with
custom fallback rules and without exhausting memory.
Business Rules.

1. Unique Key - If a key occurs exactly once in the lookup table,
simply pull in that row’s data.

2. Duplicate Keys
- When multiple lookup rows share the same key, check each row’s date range (y ≤
join_date ≤ x).
- If one or more rows satisfy the range, use the first matching row
(in original file order).
- If none match, take the last duplicate (again preserving file
order).
3. Missing Key
- If a key from the main table does not exist in the lookup table, leave the
lookup fields empty (NULL).

**Technical Constraints**

1. A straightforward left‐merge in pandas or
Dask blows out memory (kernel crashes / OOM).
2. Main table: ~70 million rows, ~14 million distinct keys (with duplicates across
different join dates).
3. Lookup table: similarly large, with
duplicate keys having different x/y ranges.
4. I am running in local machine with 32Gib RAM

In Dask:
I have tried setting the key columns for merge as index before merging in order to
prevent shuffling. Partitioned data with partition size of 250MB, 150Mb and 200MB.
I have also tried partitioning based on keys but it was creating skewed
partitions.


Below is the pandas code i have tried

import pandas as pd
import numpy as np

value_col = ['taste'] # the value column that need to be taken from the lookup table
key_col = 'Fruits' # key column for the lookup
start_date = 'y' # column of the lookup date
end_date = 'x' # column of the lookup date
join_date = 'date' # column of the lookup date
nb_table = 1

df = pd.read_parquet()
dataset['df1'] = pd.read_csv()



for i in range(1,nb_table+1):

df = df.reset_index()

#key_col is the lookup key and the datsets are merged based on those keys
# This merge generates Cartesian products for duplicate keys, allowing conditions to be applied later to retain only the necessary rows.
df = pd.merge(df,dataset[f'df{i}'][key_col+value_col], on=key_col,how='left',indicator=True)

# Using the index column for grouping so that duplicate keys are not grouped together.
df['group_num'] = df.groupby(by = var_dict[f'__varclein{i}']+['index']).ngroup()

# convert date columns to required format
# creates the flag column with 1 for the key which falls in the date range and zero which doesn't fall in the date range
df[x] = pd.to_datetime(df[x], format="%Y-%m-%d", errors='coerce')
df[y] = pd.to_datetime(df[y], format="%Y-%m-%d", errors='coerce')
df[join_date] = pd.to_datetime(df[join_date], format="%Y-%m-%d", errors='coerce')
df['flag'] = ((df[join_date] <= df[end_date]) & (df[join_date] >=
df[start_date])).astype('int8')


df['index'] = df.index

# Compute the first row per group where flag == 1
# This filters the DataFrame to only rows with flag==1, then uses
groupby('group').first() to pick the first occurrence.
# first_zero now contains, for groups that contain at least one 1, the first row's original index.
first_one = (
df[df.flag == 1]
.groupby(key_col+['group_num'], as_index=False)
.first()[key_col+['group_num']+['index']]
)

# Compute the last row overall for each group
# last_row contains the last row's original index for each group.
last_row = (
df.groupby(key_col+['group_num'], as_index=False)
.last()[key_col+['group_num']+['index']]
)


# Merge the two results and decide which row to keep per group
# Merge the "last_row" (which covers all groups) with "first_one" (which may be missing for groups with no ones).
merged = pd.merge(last_row, first_one, on=key_col+['group_num'], how='left', suffixes=('_last', '_first'))


# if the first_one index is missing (NaN), then use the last_row index; otherwise, use the first_zero index.
merged['final_idx'] = np.where(merged['index_first'].isnull(),
merged['index_last'],
merged['index_first'])

# Select those rows from the original DataFrame
result = df.loc[merged['final_idx']].copy()

if flag_col !=None:
condition1 = (result['flag']==1)
condition2 = (result['_merge']!='left_only')
condition3 = ((result['flag']!=1)

# if the date falls in the date range and there is key in the lookup table then value of flag_col is 1
result[flag_col] = np.where(condition1 & condition2,1,result[flag_col])

# if the date doesn't fall in the date range but there is key in the lookup table then value of flag_col is 2
result[flag_col] = np.where(condition3 & condition2,2,result[flag_col])

# if there is no key in the lookup table then value of flag_col is 3
result[flag_col] = np.where(~condition2,3,result[flag_col])

df = result.drop(columns=['index','_merge'])


Below is the dask code I have tried

value_col = ['taste'] # the value column that need to be taken from the lookup table based on the key and conditions
key_col = 'Fruits' # key column for the lookup
start_date = 'y' # column of the lookup date
end_date = 'x' # column of the lookup date
join_date = 'date' # column of the lookup date

df = dd.read_parquet()
dataset['df1'] = dd.read_csv()
dataset['df1'] = dataset['df1'].repartition(partition_size='250MB').persist()
df = df.repartition(partition_size='250MB').persist()

for i in range(1,nb_table+1):

df = df.reset_index()

df['merge_key'] = df[key_col][0]].astype(str)
dataset[f'df{i}']['merge_key'] = dataset[f'df{i}'][key_col][0]].astype(str)
for c in key_col[1:]:
df['merge_key'] = df['merge_key'] + '_' + df[c].astype(str)
dataset[f'df{i}']['merge_key'] = dataset[f'df{i}']['merge_key'] + '_' + dataset[f'df{i}'][c].astype(str)

df = df.set_index('merge_key').persist()
dataset[f'df{i}'] = dataset[f'df{i}'].set_index('merge_key').persist()


df = dd.merge(df, dataset[f'df{i}'][key_col+value_col],left_index=True, right_index=True,how='left',indicator=True)
df = df.reset_index(drop='merge_key')
df = df.repartition(partition_size='250MB').persist()

group_cols = key_col + ['index']
for col in group_cols:
df[col] = df[col].fillna('').astype(str)
df['group_key'] = functools.reduce(lambda a, b: a + '_' + b, [df[col] for col in group_cols])

# Hash the group key column to simulate ngroup()
df['group_num'] = df['group_key'].map_partitions(
lambda s: s.apply(hash), meta=('group_num', 'int64')
)

# creating conditional flags
df[x] = dd.to_datetime(df[x], format="%Y-%m-%d", errors='coerce')
df[y] = dd.to_datetime(df[y], format="%Y-%m-%d", errors='coerce')
df[join_date] = dd.to_datetime(df[join_date], format="%Y-%m-%d", errors='coerce')
df['flag'] = ((df[join_date] <= df[end_date]) & (df[join_date] >= df[start_date])).astype('int8')


df['index'] = df.index

group_cols = key_col + ['group_num']

# Get first row with flag == 1
df_zero = df[df['flag'] == 1]
first_one = (
df_zero.groupby(group_cols)
.first()
.reset_index()[group_cols + ['index']]
)

# Get last row per group (regardless of flag)
last_row = (
df.groupby(group_cols)
.last()
.reset_index()[group_cols + ['index']]
)
merged = dd.merge(last_row, first_one,on=key_col+['group_num'],how='left',suffixes=('_last', '_first'),indicator=True)


merged['final_idx'] = merged['index_first'].where(
~merged['index_first'].isnull(), # use index_first when it's not null
merged['index_last'] # else use index_last
)

result = merged[key_col+['final_idx']]
result = result.merge(df, left_on=key_col+['final_idx'], right_on=key_col+['index'], how='left')

var_name = flag_col
if var_name is not None:

# Compute flag conditions efficiently
not_left_only = result['_merge'] != 'left_only'

# First condition: flag == 1 and merged is set to 1
mask_1 = (result['flag'] == 1) & not_left_only
result[var_name] = result[var_name].where(~mask_1, 1)

# Second condition: flag != 1 and merged is set to 2
mask_2 = (result['flag'] != 1) & not_left_only
result[var_name] = result[var_name].where(~mask_2, 2)

# Final: where not merged is set to 2
result[var_name] = result[var_name].where(not_left_only, 3)

# Drop unnecessary columns
df = result.drop(columns=['index', '_merge'])
df = df.repartition(partition_size='250MB').persist()


Below is the snapshot of the inputs and output

first table - main data where key is the fruits column

second table - Lookup table where x and y column are the date ranges and Fruits is the key col

Third table - it is the final output where flag is the conditional column created






Continue reading...