[SPARK-41952][SQL] Fix Parquet zstd off-heap memory leak as a workaround for PARQUET-2160

[pan3793]

What changes were proposed in this pull request?

SPARK-41952 was raised for a while, but unfortunately, the Parquet community does not publish the patched version yet, as a workaround, we can fix the issue on the Spark side first.

We encountered this memory issue when migrating data from parquet/snappy to parquet/zstd, Spark executors always occupy unreasonable off-heap memory and have a high risk of being killed by NM.

See more discussions at apache/parquet-java#982 and apache/iceberg#5681

Why are the changes needed?

The issue is fixed in the parquet community PARQUET-2160, but the patched version is not available yet.

Does this PR introduce any user-facing change?

Yes, it's bug fix.

How was this patch tested?

The existing UT should cover the correctness check, I also verified this patch by scanning a large parquet/zstd table.

spark-shell --executor-cores 4 --executor-memory 6g --conf spark.executor.memoryOverhead=2g

spark.sql("select sum(hash(*)) from parquet_zstd_table ").show(false)

• before this patch

All executors get killed by NM quickly.

ERROR YarnScheduler: Lost executor 1 on hadoop-xxxx.****.org: Container killed by YARN for exceeding physical memory limits. 8.2 GB of 8 GB physical memory used. Consider boosting spark.executor.memoryOverhead.

• after this patch

Query runs well, no executor gets killed.

[dongjoon-hyun]

cc @sunchao and @parthchandra

[dongjoon-hyun]

BTW, do you happen to know why Apache Parquet community didn't fill the Fixed Version of that parquet JIRA or did release, @pan3793 ? According to the log, it seems to be merged on September, 2022.

• apache/parquet-java@53f65a8

[pan3793]

Based on apache/parquet-java#982 (comment), I guess that the Parquet community may think it's not a critical issue, but in my case, it's critical.

[LuciferYang]

For 3.2/.3.3/3.4 branch, I am ok to this fix way, and we can revert ParquetCodecFactory.java after upgrading parquet to new version in master.

[LuciferYang]

banch-3.3/3.2 use parquet 1.12.2, if this fix is accepted, would you mind submitting pr for these two branches? @pan3793

[pan3793]

banch-3.3/3.2 use parquet 1.12.2, if this fix is accepted, would you mind submitting pr for these two branches? @pan3793

sure.

[pan3793]

I verified this patch by scanning a large parquet/zstd table and updated the PR description.

[LuciferYang]

+1, LGTM

[LuciferYang]

I think we should add a TODO to remind us to revert this file when updating Parquet version

[pan3793]

updated

[LuciferYang]

should be TODO(SPARK-xxx):

[sunchao]

LGTM

[dongjoon-hyun]

Thank you all. And, feel free to merge to land this to the release branches, @sunchao .

cc @cloud-fan , @HyukjinKwon , @mridulm , @tgravescs

[sunchao]

Thanks! merged to master/branch-3.3/branch-3.2

[LuciferYang]

@sunchao also need merge to branch-3.4 ...

[sunchao]

Yes, it's in branch-3.4 as well

[yujhe-li]

@pan3793 We found that in some cases this workaround did not work when reading Parquet files in zstd codec and outputting it to a new path.

val df = spark.read.parquet("zstd_parquet_file")
df.repartition(3).write.mode("overwrite").parquet("/tmp/test")

In debug mode, you can see it won't set ParquetCodecFactory when initializing org.apache.parquet.hadoop.ParquetFileReader which is called from ParquetFileFormat.buildReaderWithPartitionValues.

[pan3793]

@yujhe possible to provide a text thread stack? I will find some time to take a look later

[yujhe-li]

Sure, here is the thread stack.

"Executor task launch worker for task 0.0 in stage 1.0 (TID 1)@7377" daemon prio=5 tid=0x36 nid=NA runnable
  java.lang.Thread.State: RUNNABLE
	  at org.apache.parquet.hadoop.ParquetFileReader.<init>(ParquetFileReader.java:770)
	  at org.apache.parquet.hadoop.ParquetFileReader.open(ParquetFileReader.java:657)
	  at org.apache.parquet.hadoop.ParquetRecordReader.initializeInternalReader(ParquetRecordReader.java:162)
	  at org.apache.parquet.hadoop.ParquetRecordReader.initialize(ParquetRecordReader.java:140)
	  at org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat.$anonfun$buildReaderWithPartitionValues$2(ParquetFileFormat.scala:405)
	  at org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat$$Lambda$2412.203214861.apply(Unknown Source:-1)
	  at org.apache.spark.sql.execution.datasources.FileScanRDD$$anon$1.org$apache$spark$sql$execution$datasources$FileScanRDD$$anon$$readCurrentFile(FileScanRDD.scala:209)
	  at org.apache.spark.sql.execution.datasources.FileScanRDD$$anon$1.nextIterator(FileScanRDD.scala:270)
	  at org.apache.spark.sql.execution.datasources.FileScanRDD$$anon$1.hasNext(FileScanRDD.scala:116)
	  at scala.collection.Iterator$$anon$9.hasNext(Iterator.scala:583)
	  at org.apache.spark.sql.execution.datasources.FileFormatDataWriter.writeWithIterator(FileFormatDataWriter.scala:91)
	  at org.apache.spark.sql.execution.datasources.FileFormatWriter$.$anonfun$executeTask$1(FileFormatWriter.scala:341)
	  at org.apache.spark.sql.execution.datasources.FileFormatWriter$$$Lambda$2521.582578288.apply(Unknown Source:-1)
	  at org.apache.spark.util.Utils$.tryWithSafeFinallyAndFailureCallbacks(Utils.scala:1538)
	  at org.apache.spark.sql.execution.datasources.FileFormatWriter$.executeTask(FileFormatWriter.scala:348)
	  at org.apache.spark.sql.execution.datasources.FileFormatWriter$.$anonfun$write$22(FileFormatWriter.scala:266)
	  at org.apache.spark.sql.execution.datasources.FileFormatWriter$$$Lambda$2415.242150313.apply(Unknown Source:-1)
	  at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:90)
	  at org.apache.spark.scheduler.Task.run(Task.scala:136)
	  at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$3(Executor.scala:548)
	  at org.apache.spark.executor.Executor$TaskRunner$$Lambda$1345.354344687.apply(Unknown Source:-1)
	  at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1504)
	  at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:551)
	  at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
	  at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
	  at java.lang.Thread.run(Thread.java:748)

[yujhe-li]

We found that this happens if we are reading Parquet file with nested columns in schema.

val path = "/tmp/parquet_zstd"
(1 to 100).map(i => (i, Seq(i)))
  .toDF("id", "value")
  .repartition(1)
  .write
  .mode("overwrite")
  .parquet(path)

val df = spark.read.parquet(path)
df.write.mode("overwrite").parquet("/tmp/dummy")

After tracing the code, ParquetCodecFactory only applies to VectorizedParquetRecordReader with vectorized reader enabled.

However, for schema with nested columns, the vectorized reader is disabled by default in Spark 3.3 (spark.sql.parquet.enableNestedColumnVectorizedReader=false). Therefore, this workaround does not work in this case.

[pan3793]

@yujhe Oops, I totally forgot this one...

Your analysis makes sense, I took a look at the parquet non-vectorized reading code path, injecting such a workaround does not clean as we did in this PR. As this is not a regression, and the new Spark version adopts new Parquet version which already fixed this issue, I prefer to keep things as-is in branch-3.3.

[LuciferYang]

@yujhe Oops, I totally forgot this one...

Your analysis makes sense, I took a look at the parquet non-vectorized reading code path, injecting such a workaround does not clean as we did in this PR. As this is not a regression, and the new Spark version adopts new Parquet version which already fixed this issue, I prefer to keep things as-is in branch-3.3.

+1, It seems there's no elegant way to fix the path where the native parquet-mr initializes the reader.

[yujhe-li]

Totally understand, we will try to apply patch PARQUET-2160 to parquet-hadoop-1.12.2 to our cluster.