[Python] How do I efficiently write results from python's ProcessPoolExecutor to a single...

I use python to read a large csv file, transform the data within (mostly string operations), and then write the results to a single parquet file. In the transformation process, rows are independent. 1 row in the csv file yields 1 row in the parquet file.

In my current implementation, I divide the data into chunks (via pandas.read_csv()) and then transform the chunked data, allocating the work across most of the cores in my CPU (via multiprocessing.ProcessPoolExecutor). Once all tasks are complete, I concatenate the results into a single, large pandas dataframe and then write it to a parquet file.

The code below provides a minimal example.

Concatenating the results maxes out memory usage. Is there a safe, more efficient process that writes each chunk of transformed data once it becomes available? I've read about ThreadPoolExecutor, Queues, and futures but I haven't found a clean example of how to combine all of these multiprocessing techniques in a case like mine.

import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd
import regex as re
from concurrent.futures import ProcessPoolExecutor
from multiprocessing import cpu_count, Queue

filename = 'input.csv'

# Create example of input file
n_rows = int(25 * 1e6) # a mid-sized dataset/ I am anticipating much larger files
input_df = pd.DataFrame([f'User_id_{i}_says_"hello!".' for i in range(n_rows)], columns=['A'])
input_df.to_csv(filename, header=True, index=False)

# Example of data transformation
def clean_text(s):
regex = re.compile('[^a-zA-Z]')
return regex.sub('', s)

def process(df):
# Split string into 5 columns.
out = pd.DataFrame(df['A'].str.split('_').to_list())

# Convert numeric id to string
out[2] = out[2].astype(str)

# Remove non-alphanumerics from column 4 entries.
out[4] = out[4].apply(lambda x: clean_text(x))

# Rename columns
out.rename(columns={i: f'A{i}' for i in out.columns }, inplace=True)

return out

if __name__ == "__main__":

# Set parquet data schema
schema = pa.schema((f'A{i}', pa.string()) for i in range(5))

# Multiprocessing implementation
core_count = cpu_count() - 3 # OP's machine has 24 cores

# Read input data in chunks
chunk_size = int(1e5)
chunks = pd.read_csv('input.csv', header=0, na_filter=False, chunksize=chunk_size)

# Allocate transformation task
with pq.ParquetWriter('processed.parquet', schema=schema, compression='GZIP') as writer:
with ProcessPoolExecutor(core_count) as executor:
results = executor.map(process, chunks)
df = pd.concat(results, axis=0)

# Convert df to record batch
transformed_batch = pa.RecordBatch.from_pandas(df, schema=schema)
writer.write_batch(transformed_batch)

Continue reading...