Splitting a mishappen csv file using pyspark RDD. EMR. Yarn memory exception errors - apache-spark

Question

I have been working on this code for a while. Following I have listed the code and most of the cluster attributes I am using on EMR. The purpose of the code is to split some csv files in two at a certain line number based on some basic iteration (I have included a simple split in the code below).
I frequently get this error "Container killed by YARN for exceeding memory limits" and have followed these design principles (link below) to resolve it, but I just don't know why this would get memory problems. I have over 22GB for yarn overhead, and the files are in the MB to single digit GB ranges.
I am using sometimes r5a.12xlarges to no avail. I just really don't see any kind of memory leak in this code. It also seems very slow, was only able to process something like 20GB in 16 hours output to S3. Is this a good way to parallelize this split operation? Is there a memory leak? What gives?
https://aws.amazon.com/premiumsupport/knowledge-center/emr-spark-yarn-memory-limit/
[
{
"Classification": "spark",
"Properties": {
"spark.maximizeResourceAllocation": "true"
}
},
{
"Classification": "spark-defaults",
"Properties": {
"spark.yarn.executor.memoryOverheadFactor":".2"
}
},
{
"Classification": "spark-env",
"Configurations": [
{
"Configurations": [],
"Properties": {
"PYSPARK_PYTHON": "python36"
},
"Classification": "export"
}
],
"Properties": {
}
}
]
def writetxt(txt: Union[List[str], pandas.DataFrame], path: str) -> None:
s3 = boto3.resource('s3')
s3path = S3Url(path)
object = s3.Object(s3path.bucket, s3path.key)
if isinstance(txt, pandas.DataFrame):
csv_buffer = StringIO()
txt.to_csv(csv_buffer)
object.put(Body=csv_buffer.getvalue())
else:
object.put(Body='\n'.join(txt).encode())
def main(
x: Iterator[Tuple[str, str]],
output_files: str
) -> None:
filename, content = x
filename = os.path.basename(S3Url(filename).key)
content = content.splitlines()
# Split the csv file
columnAttributes, csvData = data[:100], data[100:]
writetxt(csvData, os.path.join(output_files, 'data.csv', filename))
writetxt(columnAttributes, os.path.join(output_files, 'attr.csv', filename))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Split some mishapen csv files.')
parser.add_argument('input_files', type=str,
help='The location of the input files.')
parser.add_argument('output_files', type=str,
help='The location to put the output files.')
parser.add_argument('--nb_partitions', type=int, default=4)
args = parser.parse_args()
# creating the context
sc = SparkContext(appName="Broadcom Preprocessing")
# We use minPartitions because otherwise small files get put in the same partition together
# by default, which we have a lot of
# We use foreachPartition to reduce the number of function calls, which slow down spark
distFiles = sc.wholeTextFiles(args.input_files, minPartitions=args.nb_partitions) \
.foreach(partial(main, output_files=args.output_files))

Answer 1

I think your memory issues are because you're doing the actual data-splitting with Python code. Spark processes run in the JVM, but when you call custom Python code, the related data must be serialized over to a Python process (on each worker node) in order to execute. This adds a lot of overhead. I believe you can accomplish what you're trying to do entirely with Spark operations - meaning the final program will run entirely in the JVM-based Spark processes.
Try something like this:
from pyspark.sql.types import IntegerType
from pyspark.sql.window import Window
from pyspark.sql.functions import *
input_path = "..."
split_num = 100
# load filenames & contents
filesDF = spark.createDataFrame( sc.wholeTextFiles(input_path), ['filename','contents'] )
# break into individual lines & number them
linesDF = filesDF.select( "filename", \
row_number().over(Window.partitionBy("filename").orderBy("filename")).alias("line_number"), \
explode(split(col("contents"), "\n")).alias("contents") )
# split into headers & body
headersDF = linesDF.where(col("line_number") == lit(1))
bodyDF = linesDF.where(col("line_number") > lit(1))
# split the body in 2 based
splitLinesDF = bodyDF.withColumn("split", when(col("line_number") < lit(split_num), 0).otherwise(1))
split_0_DF = splitLinesDF.where(col("split") == lit(0)).select("filename", "line_number", "contents").union(headersDF).orderBy("filename", "line_number")
split_1_DF = splitLinesDF.where(col("split") == lit(1)).select("filename", "line_number", "contents").union(headersDF).orderBy("filename", "line_number")
# collapse all lines back down into a file
firstDF = split_0_DF.groupBy("filename").agg(concat_ws("\n",collect_list(col("contents"))).alias("contents"))
secondDF = split_1_DF.groupBy("filename").agg(concat_ws("\n",collect_list(col("contents"))).alias("contents"))
# pandas-UDF for more memory-efficient transfer of data from Spark to Python
#pandas_udf(returnType=IntegerType())
def writeFile( filename, contents ):
<save to S3 here>
# write each row to a file
firstDF.select( writeFile( col("filename"), col("contents") ) )
secondDF.select( writeFile( col("filename"), col("contents") ) )
Finally you will need to use some custom python code to save each split file to S3 (or, you could just code everything in Scala/Java). It is far, far more efficient to do this via pandas UDFs than to pass a standard python function to .foreach(...). Internally, spark will serialize the data to Arrow format in chunks (one per partition), which will be very efficient.
Additionally, it looks like you're trying to put the entire object into S3 in a single request. If the data is too large for this, it will fail. You should check out the S3 streaming upload functionality.