Spark中写parquet文件是怎么实现的

本文主要是介绍Spark中写parquet文件是怎么实现的，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

背景

本文基于 Spark 3.5.0
写本篇文章的目的是在于能够配合spark.sql.maxConcurrentOutputFileWriters参数来加速写parquet文件的速度，为此研究一下Spark写parquet的时候会占用内存的大小，便于配置spark.sql.maxConcurrentOutputFileWriters的值，从而保证任务的稳定性

结论

一个spark parquet writer可能会占用128MB的内存（也就是parquet.block.size的大小）。所有在调整spark.sql.maxConcurrentOutputFileWriters的时候得注意不能调整过大，否则会导致OOM，但是如果在最后写文件的时候加入合并小文件的功能（AQE+Rebalance的方式），也可以适当的调整大一点，因为这个时候的Task 不像没有shuffle一样，可能还会涉及到sort以及aggregate等消耗内存的操作，（这个时候就是一个task纯写parquet文件）
大家也可以参考Parquet文件是怎么被写入的-Row Groups，Pages，需要的内存，以及flush操作

分析

还是得从InsertIntoHadoopFsRelationCommand类中说起，涉及到写parquet的数据流如下:

InsertIntoHadoopFsRelationCommand.run||\/
FileFormatWriter.write||\/
fileFormat.prepareWrite||\/
executeWrite => planForWrites.executeWrite ||\/WriteFilesExec.doExecuteWrite||\/FileFormatWriter.executeTask||\/dataWriter.writeWithIterator||\/dataWriter.writeWithMetrics||\/DynamicPartitionDataConcurrentWriter.write||\/writeRecord||\/ParquetOutputWriter.write||\/recordWriter.write

其中fileFormat.prepareWrite 涉及到 spark这一层级有关parquet的设置，并返回一个生成ParquetOutputWriter实例的工厂类实例OutputWriterFactory
主要设置如 parquet.compression 压缩格式,一般是 zstd ,也可以通过 spark.sql.parquet.compression.codec设置
parquet.write.support.class 为ParquetWriteSupport,该类的作用为Spark把内部IternalRow转为parquet message
DynamicPartitionDataConcurrentWriter.write 涉及到了InternalRow到UnsafeRow代码生成
这里不讨论这部分的细节，只说一下getPartitionValues和renewCurrentWriter 方法中的 getPartitionPath这两部分
- getPartitionValues
  这个是InternalRow => UnsafeRow转换，为什么这么做，是因为对于UnsafeRow这种数据结构来说，能够很好管理内存和避免GC问题
```
    val proj = UnsafeProjection.create(description.partitionColumns, description.allColumns)row => proj(row)
```
  我们以UnsafeProjection的子类InterpretedUnsafeProjection，该类不是代码生成的类（这样便于分析），
```
  override def apply(row: InternalRow): UnsafeRow = {if (subExprEliminationEnabled) {runtime.setInput(row)}// Put the expression results in the intermediate row.var i = 0while (i < numFields) {values(i) = exprs(i).eval(row)i += 1}// Write the intermediate row to an unsafe row.rowWriter.reset()writer(intermediate)rowWriter.getRow()}
```
  - 首先是消除公共子表达式
  - 用values数组保存每个表达式计算出来的结果
  - rowWriter.reset() 用来对齐cursor,便于对于String类型的写入，这可以参考UnsafeRow内存布局和代码优化
  - unsafeWriter按照不同的类型写入到unsaferow不同的位置上去，这里的offset在cursor的内部的，也就是说cursor的值要大于offset
  - 返回UnsafeRow类型
    通过这种方式完成了InternalRow => UnsafeRow转换
- getPartitionPath
  这个是通过表达式的方式获取partition的函数，从而完成InternalRow => String的转换，涉及的代码如下：
```
  private lazy val partitionPathExpression: Expression = Concat(description.partitionColumns.zipWithIndex.flatMap { case (c, i) =>val partitionName = ScalaUDF(ExternalCatalogUtils.getPartitionPathString _,StringType,Seq(Literal(c.name), Cast(c, StringType, Option(description.timeZoneId))))if (i == 0) Seq(partitionName) else Seq(Literal(Path.SEPARATOR), partitionName)})private lazy val getPartitionPath: InternalRow => String = {val proj = UnsafeProjection.create(Seq(partitionPathExpression), description.partitionColumns)row => proj(row).getString(0)}
```
  UnsafeProjection.create 上面已经说了怎么实现的了,重点说partitionPathExpression 生成partition的表达式,
  该表达式主要通过UDF中getPartitionPathString来生成，关键的一点是,传入的参数：Literal(c.name)和Cast(c, StringType, Option(description.timeZoneId))))
  这里的Literal(c.name)表示的是partition名字的常量
  Cast(c, StringType, Option(description.timeZoneId)))表示的是c这个变量所代表的值，
  为什么这么说，因为在ScalaUDF的内部计算方法中有：
```
  override def eval(input: InternalRow): Any = {val result = try {f(input)} catch {case e: Exception =>throw QueryExecutionErrors.failedExecuteUserDefinedFunctionError(functionName, inputTypesString, outputType, e)}resultConverter(result)}
```
  这里的f中会对传入的每个参数都会调用eval(InernalRow)，对于Literal来说就是常亮，而对于Cast(Attribute)来说就是属性的值(通过BindReferences.bindReference方法)。
recordWriter.write涉及到 ParquetOutputFormat.getRecordWriter方法，该方法中涉及到parquet中的一些原生参数设置：

public RecordWriter<Void, T> getRecordWriter(Configuration conf, Path file, CompressionCodecName codec, Mode mode)throws IOException, InterruptedException {final WriteSupport<T> writeSupport = getWriteSupport(conf);ParquetProperties.Builder propsBuilder = ParquetProperties.builder().withPageSize(getPageSize(conf)).withDictionaryPageSize(getDictionaryPageSize(conf)).withDictionaryEncoding(getEnableDictionary(conf)).withWriterVersion(getWriterVersion(conf)).estimateRowCountForPageSizeCheck(getEstimatePageSizeCheck(conf)).withMinRowCountForPageSizeCheck(getMinRowCountForPageSizeCheck(conf)).withMaxRowCountForPageSizeCheck(getMaxRowCountForPageSizeCheck(conf)).withColumnIndexTruncateLength(getColumnIndexTruncateLength(conf)).withStatisticsTruncateLength(getStatisticsTruncateLength(conf)).withMaxBloomFilterBytes(getBloomFilterMaxBytes(conf)).withBloomFilterEnabled(getBloomFilterEnabled(conf)).withPageRowCountLimit(getPageRowCountLimit(conf)).withPageWriteChecksumEnabled(getPageWriteChecksumEnabled(conf));new ColumnConfigParser().withColumnConfig(ENABLE_DICTIONARY, key -> conf.getBoolean(key, false), propsBuilder::withDictionaryEncoding).withColumnConfig(BLOOM_FILTER_ENABLED, key -> conf.getBoolean(key, false),propsBuilder::withBloomFilterEnabled).withColumnConfig(BLOOM_FILTER_EXPECTED_NDV, key -> conf.getLong(key, -1L), propsBuilder::withBloomFilterNDV).withColumnConfig(BLOOM_FILTER_FPP, key -> conf.getDouble(key, ParquetProperties.DEFAULT_BLOOM_FILTER_FPP),propsBuilder::withBloomFilterFPP).parseConfig(conf);ParquetProperties props = propsBuilder.build();long blockSize = getLongBlockSize(conf);int maxPaddingSize = getMaxPaddingSize(conf);boolean validating = getValidation(conf);...WriteContext fileWriteContext = writeSupport.init(conf);FileEncryptionProperties encryptionProperties = createEncryptionProperties(conf, file, fileWriteContext);ParquetFileWriter w = new ParquetFileWriter(HadoopOutputFile.fromPath(file, conf),fileWriteContext.getSchema(), mode, blockSize, maxPaddingSize, props.getColumnIndexTruncateLength(),props.getStatisticsTruncateLength(), props.getPageWriteChecksumEnabled(), encryptionProperties);w.start();...return new ParquetRecordWriter<T>(w,writeSupport,fileWriteContext.getSchema(),fileWriteContext.getExtraMetaData(),blockSize,codec,validating,props,memoryManager,conf);}

这里涉及到的关键的几个参数是：

   parquet.page.size                   1*1024*1024         -- page的大小 默认是 1MBparquet.block.size                  128*1024*1024       -- rowgroup的大小 默认是 128MBparquet.page.size.row.check.min     100                 -- page检查是否达到page size的最小行数parquet.page.size.row.check.max     10000               -- page检查是否达到page size的最大行数parquet.page.row.count.limit        20_000              -- page检查是否达到page size的行数极限行数

parquet.page.size.row.check.min parquet.page.size.row.check.max parquet.page.row.count.limit 这三个配置项存在着相互制约的关系，总的目标就是检查当行数达到了一定的阈值以后，来检查是否能够flush到内存page中，具体的可以查看ColumnWriteStoreBase类中的方法

接下来就是真正写操作了,调用的是InternalParquetRecordWriter.write方法，如下：

 private void initStore() {ColumnChunkPageWriteStore columnChunkPageWriteStore = new ColumnChunkPageWriteStore(compressor,schema, props.getAllocator(), props.getColumnIndexTruncateLength(), props.getPageWriteChecksumEnabled(),fileEncryptor, rowGroupOrdinal);pageStore = columnChunkPageWriteStore;bloomFilterWriteStore = columnChunkPageWriteStore;columnStore = props.newColumnWriteStore(schema, pageStore, bloomFilterWriteStore);MessageColumnIO columnIO = new ColumnIOFactory(validating).getColumnIO(schema);this.recordConsumer = columnIO.getRecordWriter(columnStore);writeSupport.prepareForWrite(recordConsumer);}public void write(T value) throws IOException, InterruptedException {writeSupport.write(value);++ recordCount;checkBlockSizeReached();}

initStore主要是初始化 pageStore和columnStore
具体的spark interalRow怎么转换为parquet message，主要在writeSupport.write中的rootFieldWriters中
接下来就是checkBlockSizeReached,这里主要就是flush rowgroup到磁盘了，
具体的读者可以看代码：
对于flush到page可以看checkBlockSizeReached中columnStore.flush()
对于flush rowroup到磁盘可以看checkBlockSizeReached中pageStore.flushToFileWriter(parquetFileWriter)
总结出来就是一个spark parquet writer可能会占用128MB的内存（也就是parquet.block.size的大小），
因为只有在满足了rowgroup的大小以后，才会真正的flush到磁盘。