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 = 0 while (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的大小 默认是 1MB
   parquet.block.size                  128*1024*1024       -- rowgroup的大小 默认是 128MB
   parquet.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到磁盘。