[Python] No space error while running Spark from Airflow

I am running the following code in Airflow from MWAA. The input data is approximately 400 GB in size I am running the code in environment class mw1.large.

Up to 1000 DAG capacity

20 concurrent tasks (by default)

Components:

Web servers: 2 vCPU 4GB RAM each

Workers: 4 vCPU 8GB RAM each

Schedulers: 4 vCPU 8GB RAM each

Database: 2 vCPU 8GB RAM

I am getting no space issue after running the code below. Can you suggest optimizations as I have tried multiple combinations but nothing seems to be working.

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count
from datetime import datetime

# Define S3 bucket details
S3_BUCKET = "ms-bioinfo-biorelate-transfer"
S3_PREFIX = "uploads/2025/01"

# Input files
S3_INPUT_FILES = {
"Clinical_Trial": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_Clinical_Trial.parquet",
"Drug_Label": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_Drug_Label.parquet",
"NIH_Grant": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_NIH_Grant.parquet",
"PMC": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_PMC.parquet",
"Patent": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_Patent.parquet",
"Preprint": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_Preprint.parquet",
"PubMed": f"s3a://{S3_BUCKET}/{S3_PREFIX}/Annotations_PubMed.parquet"
}

# Output path
S3_OUTPUT_PATH = f"s3a://{S3_BUCKET}/{S3_PREFIX}/annotation_count_data"

# DAG default arguments
default_args = {
"owner": "airflow",
"depends_on_past": False,
"start_date": datetime(2025, 1, 1),
"email_on_failure": False,
"email_on_retry": False,
}

# Define DAG
dag = DAG(
'biorelate_transformation_pipeline_optimized',
default_args=default_args,
description='Optimized Biorelate Data Transformation Pipeline',
schedule_interval=None,
catchup=False,
tags=['biorelate'],
)

def process_file(file_key, file_path):
"""
Processes a single Parquet file, aggregates data, and writes the output to S3.
"""
spark = SparkSession.builder \
.appName(f"biorelate_annotation_count_{file_key}") \
.config("spark.hadoop.fs.s3a.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem") \
.config("spark.jars.packages", "org.apache.hadoop:hadoop-aws:3.3.4,org.apache.hadoop:hadoop-common:3.3.4") \
.config("spark.executor.memory", "3G") \
.config("spark.driver.memory", "3G") \
.config("spark.executor.memoryOverhead", "1G") \
.config("spark.memory.offHeap.enabled", "true") \
.config("spark.memory.offHeap.size", "512m") \
.config("spark.network.timeout", "600s") \
.config("spark.executor.heartbeatInterval", "60s") \
.config("spark.sql.shuffle.partitions", "500") \
.config("spark.sql.files.maxPartitionBytes", "64MB") \
.config("spark.sql.parquet.filterPushdown", "true") \
.config("spark.sql.parquet.mergeSchema", "false") \
.config("spark.sql.adaptive.enabled", "true") \
.config("spark.sql.parquet.compression.codec", "snappy") \
.config("spark.hadoop.fs.s3a.fast.upload", "true") \
.config("spark.hadoop.fs.s3a.multipart.size", "104857600") \
.getOrCreate()

print(f"Processing file: {file_path}")

try:
df = spark.read.parquet(file_path).select("source_type", "entity_id", "entity_name", "entity_type", "document_id")
df = df.dropDuplicates()

count_df = df.groupBy("source_type", "entity_id", "entity_name", "entity_type") \
.agg(count(col('document_id')).alias('total_num_of_pub_per_attribute'))

output_path = f"{S3_OUTPUT_PATH}/Annotations_{file_key}_processed.parquet"

count_df.write.mode("overwrite").option("maxRecordsPerFile", "5000").parquet(output_path)

print(f"Completed processing {file_key}")

# Explicit garbage collection
del df, count_df
spark.catalog.clearCache()

except Exception as e:
print(f"Error processing {file_key}: {e}")
spark.stop()
raise
finally:
spark.stop()


# Define static tasks for each file
process_clinical_trial = PythonOperator(
task_id='process_clinical_trial',
python_callable=process_file,
op_args=["Clinical_Trial", S3_INPUT_FILES["Clinical_Trial"]],
dag=dag,
)

process_drug_label = PythonOperator(
task_id='process_drug_label',
python_callable=process_file,
op_args=["Drug_Label", S3_INPUT_FILES["Drug_Label"]],
dag=dag,
)

process_nih_grant = PythonOperator(
task_id='process_nih_grant',
python_callable=process_file,
op_args=["NIH_Grant", S3_INPUT_FILES["NIH_Grant"]],
dag=dag,
)

process_pmc = PythonOperator(
task_id='process_pmc',
python_callable=process_file,
op_args=["PMC", S3_INPUT_FILES["PMC"]],
dag=dag,
)

process_patent = PythonOperator(
task_id='process_patent',
python_callable=process_file,
op_args=["Patent", S3_INPUT_FILES["Patent"]],
dag=dag,
)

process_preprint = PythonOperator(
task_id='process_preprint',
python_callable=process_file,
op_args=["Preprint", S3_INPUT_FILES["Preprint"]],
dag=dag,
)

process_pubmed = PythonOperator(
task_id='process_pubmed',
python_callable=process_file,
op_args=["PubMed", S3_INPUT_FILES["PubMed"]],
dag=dag,
)

process_pmc >> process_patent >> process_preprint >> process_pubmed >> process_clinical_trial >> process_drug_label >> process_nih_grant

Continue reading...