Spark写Parquet源码分析

/**
   * Interface for saving the content of the non-streaming Dataset out into external storage.
   *
   * @group basic
   * @since 1.6.0
   */
  def write: DataFrameWriter[T] = {
    if (isStreaming) {
      logicalPlan.failAnalysis(
        "'write' can not be called on streaming Dataset/DataFrame")
    }
    new DataFrameWriter[T](this)
  }

def parquet(path: String): Unit = {
  format("parquet").save(path)
}

/**
 * Writes the given [[DataFrame]] out to this [[DataSource]].
 */
def write(mode: SaveMode, data: DataFrame): Unit = {
  if (data.schema.map(_.dataType).exists(_.isInstanceOf[CalendarIntervalType])) {
    throw new AnalysisException("Cannot save interval data type into external storage.")
  }
  providingClass.newInstance() match {
    case dataSource: CreatableRelationProvider =>
      dataSource.createRelation(sparkSession.sqlContext, mode, caseInsensitiveOptions, data)
    case format: FileFormat =>
      writeInFileFormat(format, mode, data)
    case _ =>
      sys.error(s"${providingClass.getCanonicalName} does not allow create table as select.")
  }
}

/**
 * Writes the given [[DataFrame]] out in this [[FileFormat]].
 */
private def writeInFileFormat(format: FileFormat, mode: SaveMode, data: DataFrame): Unit = {
  ...
  ...
  val plan =
      InsertIntoHadoopFsRelationCommand(
        outputPath = outputPath,
        staticPartitionKeys = Map.empty,
        customPartitionLocations = Map.empty,
        partitionColumns = columns,
        bucketSpec = bucketSpec,
        fileFormat = format,
        refreshFunction = _ => Unit, // No existing table needs to be refreshed.
        options = options,
        query = data.logicalPlan,
        mode = mode,
        catalogTable = catalogTable)
    sparkSession.sessionState.executePlan(plan).toRdd
}

FileFormatWriter.write(
        sparkSession = sparkSession,
        queryExecution = Dataset.ofRows(sparkSession, query).queryExecution,
        fileFormat = fileFormat,
        committer = committer,
        outputSpec = FileFormatWriter.OutputSpec(
          qualifiedOutputPath.toString, customPartitionLocations),
        hadoopConf = hadoopConf,
        partitionColumns = partitionColumns,
        bucketSpec = bucketSpec,
        refreshFunction = refreshFunction,
        options = options)

def write(
      sparkSession: SparkSession,
      queryExecution: QueryExecution,
      fileFormat: FileFormat,
      committer: FileCommitProtocol,
      outputSpec: OutputSpec,
      hadoopConf: Configuration,
      partitionColumns: Seq[Attribute],
      bucketSpec: Option[BucketSpec],
      refreshFunction: (Seq[TablePartitionSpec]) => Unit,
      options: Map[String, String]): Unit = {
    ...
    ...
    val outputWriterFactory =
      fileFormat.prepareWrite(sparkSession, job, options, dataColumns.toStructType)
}

override def prepareWrite(
      sparkSession: SparkSession,
      job: Job,
      options: Map[String, String],
      dataSchema: StructType): OutputWriterFactory = {
    new OutputWriterFactory {
      // This OutputWriterFactory instance is deserialized when writing Parquet files on the
      // executor side without constructing or deserializing ParquetFileFormat. Therefore, we hold
      // another reference to ParquetLogRedirector.INSTANCE here to ensure the latter class is
      // initialized.
      private val parquetLogRedirector = ParquetLogRedirector.INSTANCE

        override def newInstance(
          path: String,
          dataSchema: StructType,
          context: TaskAttemptContext): OutputWriter = {
        new ParquetOutputWriter(path, context)
      }

      override def getFileExtension(context: TaskAttemptContext): String = {
        CodecConfig.from(context).getCodec.getExtension + ".parquet"
      }
    }
}

private def consumeField(field: String, index: Int)(f: => Unit): Unit = {
    recordConsumer.startField(field, index)
    f
    recordConsumer.endField(field, index)
  }

// NOTE: This class is instantiated and used on executor side only, no need to be serializable.
private[parquet] class ParquetOutputWriter(path: String, context: TaskAttemptContext)
  extends OutputWriter {

  private val recordWriter: RecordWriter[Void, InternalRow] = {
    new ParquetOutputFormat[InternalRow]() {
      override def getDefaultWorkFile(context: TaskAttemptContext, extension: String): Path = {
        new Path(path)
      }
    }.getRecordWriter(context)
  }

  override def write(row: Row): Unit = throw new UnsupportedOperationException("call writeInternal")

  override def writeInternal(row: InternalRow): Unit = recordWriter.write(null, row)

  override def close(): Unit = recordWriter.close(context)
}

public RecordWriter<Void, T> getRecordWriter(Configuration conf, Path file, CompressionCodecName codec)
        throws IOException, InterruptedException {
    final WriteSupport<T> writeSupport = getWriteSupport(conf);

    CodecFactory codecFactory = new CodecFactory(conf);
    long blockSize = getLongBlockSize(conf);
    if (INFO) LOG.info("Parquet block size to " + blockSize);
    int pageSize = getPageSize(conf);
    if (INFO) LOG.info("Parquet page size to " + pageSize);
    int dictionaryPageSize = getDictionaryPageSize(conf);
    if (INFO) LOG.info("Parquet dictionary page size to " + dictionaryPageSize);
    boolean enableDictionary = getEnableDictionary(conf);
    if (INFO) LOG.info("Dictionary is " + (enableDictionary ? "on" : "off"));
    boolean validating = getValidation(conf);
    if (INFO) LOG.info("Validation is " + (validating ? "on" : "off"));
    WriterVersion writerVersion = getWriterVersion(conf);
    if (INFO) LOG.info("Writer version is: " + writerVersion);
    int maxPaddingSize = getMaxPaddingSize(conf);
    if (INFO) LOG.info("Maximum row group padding size is " + maxPaddingSize + " bytes");

    WriteContext init = writeSupport.init(conf);
    ParquetFileWriter w = new ParquetFileWriter(
        conf, init.getSchema(), file, Mode.CREATE, blockSize, maxPaddingSize);
    w.start();

    float maxLoad = conf.getFloat(ParquetOutputFormat.MEMORY_POOL_RATIO,
        MemoryManager.DEFAULT_MEMORY_POOL_RATIO);
    long minAllocation = conf.getLong(ParquetOutputFormat.MIN_MEMORY_ALLOCATION,
        MemoryManager.DEFAULT_MIN_MEMORY_ALLOCATION);
    if (memoryManager == null) {
      memoryManager = new MemoryManager(maxLoad, minAllocation);
    } else if (memoryManager.getMemoryPoolRatio() != maxLoad) {
      LOG.warn("The configuration " + MEMORY_POOL_RATIO + " has been set. It should not " +
          "be reset by the new value: " + maxLoad);
    }

    return new ParquetRecordWriter<T>(
        w,
        writeSupport,
        init.getSchema(),
        init.getExtraMetaData(),
        blockSize, pageSize,
        codecFactory.getCompressor(codec, pageSize),
        dictionaryPageSize,
        enableDictionary,
        validating,
        writerVersion,
        memoryManager);
  }

public ParquetRecordWriter(
      ParquetFileWriter w,
      WriteSupport<T> writeSupport,
      MessageType schema,
      Map<String, String> extraMetaData,
      long blockSize, int pageSize,
      BytesCompressor compressor,
      int dictionaryPageSize,
      boolean enableDictionary,
      boolean validating,
      WriterVersion writerVersion,
      MemoryManager memoryManager) {
    internalWriter = new InternalParquetRecordWriter<T>(w, writeSupport, schema,
        extraMetaData, blockSize, pageSize, compressor, dictionaryPageSize, enableDictionary,
        validating, writerVersion);
    this.memoryManager = checkNotNull(memoryManager, "memoryManager");
    memoryManager.addWriter(internalWriter, blockSize);
}
@Override
public void write(Void key, T value) throws IOException, InterruptedException {
    internalWriter.write(value);
}
@Override
public void close(TaskAttemptContext context) throws IOException, InterruptedException {
    internalWriter.close();
    if (memoryManager != null) {
      memoryManager.removeWriter(internalWriter);
    }
}

public InternalParquetRecordWriter(
    ParquetFileWriter parquetFileWriter,
    WriteSupport<T> writeSupport,
    MessageType schema,
    Map<String, String> extraMetaData,
    long rowGroupSize,
    int pageSize,
    BytesCompressor compressor,
    int dictionaryPageSize,
    boolean enableDictionary,
    boolean validating,
    WriterVersion writerVersion) {
  this.parquetFileWriter = parquetFileWriter;
  this.writeSupport = checkNotNull(writeSupport, "writeSupport");
  this.schema = schema;
  this.extraMetaData = extraMetaData;
  this.rowGroupSize = rowGroupSize;
  this.rowGroupSizeThreshold = rowGroupSize;
  this.nextRowGroupSize = rowGroupSizeThreshold;
  this.pageSize = pageSize;
  this.compressor = compressor;
  this.validating = validating;
  this.parquetProperties = new ParquetProperties(dictionaryPageSize, writerVersion, enableDictionary);
  initStore();
}

private void initStore() {
  pageStore = new ColumnChunkPageWriteStore(compressor, schema, pageSize);
  columnStore = parquetProperties.newColumnWriteStore(
      schema,
      pageStore,
      pageSize);
  MessageColumnIO columnIO = new ColumnIOFactory(validating).getColumnIO(schema);
  writeSupport.prepareForWrite(columnIO.getRecordWriter(columnStore));
}

private final Map<ColumnDescriptor, ColumnChunkPageWriter> writers = new HashMap<ColumnDescriptor, ColumnChunkPageWriter>();
private final MessageType schema;

public ColumnChunkPageWriteStore(BytesCompressor compressor, MessageType schema, int pageSize) {
    this.schema = schema;
    for (ColumnDescriptor path : schema.getColumns()) {
      writers.put(path,  new ColumnChunkPageWriter(path, compressor, pageSize));
    }
}

public ColumnWriteStore newColumnWriteStore(
      MessageType schema,
      PageWriteStore pageStore,
      int pageSize) {
    switch (writerVersion) {
    case PARQUET_1_0:
      return new ColumnWriteStoreV1(
          pageStore,
          pageSize,
          dictionaryPageSizeThreshold,
          enableDictionary, writerVersion);
    case PARQUET_2_0:
      return new ColumnWriteStoreV2(
          schema,
          pageStore,
          pageSize,
          new ParquetProperties(dictionaryPageSizeThreshold, writerVersion, enableDictionary));
    default:
      throw new IllegalArgumentException("unknown version " + writerVersion);
    }
}

private final Map<ColumnDescriptor, ColumnWriterV1> columns = new TreeMap<ColumnDescriptor, ColumnWriterV1>();

public ColumnWriter getColumnWriter(ColumnDescriptor path) {
  ColumnWriterV1 column = columns.get(path);
  if (column == null) {
    column = newMemColumn(path);
    columns.put(path, column);
  }
  return column;
}

private ColumnWriterV1 newMemColumn(ColumnDescriptor path) {
  PageWriter pageWriter = pageWriteStore.getPageWriter(path);
  return new ColumnWriterV1(path, pageWriter, pageSizeThreshold, dictionaryPageSizeThreshold, enableDictionary, writerVersion);
}

@Override
public void addInteger(int value) {
  if (DEBUG) log("addInt(" + value + ")");
  emptyField = false;
  getColumnWriter().write(value, r[currentLevel], currentColumnIO.getDefinitionLevel());

  setRepetitionLevel();
  if (DEBUG) printState();
}

@Override
public void write(int value, int repetitionLevel, int definitionLevel) {
  if (DEBUG) log(value, repetitionLevel, definitionLevel);
  repetitionLevelColumn.writeInteger(repetitionLevel);
  definitionLeveolumn.writeInteger(definitionLevel);
  dataColumn.writeInteger(value);
  updateStatistics(value);
  accountForValueWritten();
}

/**
 * Counts how many values have been written and checks the memory usage to flush the page when we reach the page threshold.
 *
 * We measure the memory used when we reach the mid point toward our estimated count.
 * We then update the estimate and flush the page if we reached the threshold.
 *
 * That way we check the memory size log2(n) times.
 *
 */
private void accountForValueWritten() {
  ++ valueCount;
  if (valueCount > valueCountForNextSizeCheck) {
    // not checking the memory used for every value
    long memSize = repetitionLevelColumn.getBufferedSize()
        + definitionLevelColumn.getBufferedSize()
        + dataColumn.getBufferedSize();
    if (memSize > pageSizeThreshold) {
      // we will write the current page and check again the size at the predicted middle of next page
      valueCountForNextSizeCheck = valueCount / 2;
      writePage();
    } else {
      // not reached the threshold, will check again midway
      valueCountForNextSizeCheck = (int)(valueCount + ((float)valueCount * pageSizeThreshold / memSize)) / 2 + 1;
    }
  }
}

private void writePage() {
  if (DEBUG) LOG.debug("write page");
  try {
    pageWriter.writePage(
        concat(repetitionLevelColumn.getBytes(), definitionLevelColumn.getBytes(), dataColumn.getBytes()),
        valueCount,
        statistics,
        repetitionLevelColumn.getEncoding(),
        definitionLevelColumn.getEncoding(),
        dataColumn.getEncoding());
  } catch (IOException e) {
    throw new ParquetEncodingException("could not write page for " + path, e);
  }
  repetitionLevelColumn.reset();
  definitionLevelColumn.reset();
  dataColumn.reset();
  valueCount = 0;
  resetStatistics();
}

@Override
public void writePage(BytesInput bytes,
                      int valueCount,
                      Statistics statistics,
                      Encoding rlEncoding,
                      Encoding dlEncoding,
                      Encoding valuesEncoding) throws IOException {
  long uncompressedSize = bytes.size();
  if (uncompressedSize > Integer.MAX_VALUE) {
    throw new ParquetEncodingException(
        "Cannot write page larger than Integer.MAX_VALUE bytes: " +
        uncompressedSize);
  }
  BytesInput compressedBytes = compressor.compress(bytes);
  long compressedSize = compressedBytes.size();
  if (compressedSize > Integer.MAX_VALUE) {
    throw new ParquetEncodingException(
        "Cannot write compressed page larger than Integer.MAX_VALUE bytes: "
        + compressedSize);
  }
  tempOutputStream.reset();
  parquetMetadataConverter.writeDataPageHeader(
      (int)uncompressedSize,
      (int)compressedSize,
      valueCount,
      statistics,
      rlEncoding,
      dlEncoding,
      valuesEncoding,
      tempOutputStream);
  this.uncompressedLength += uncompressedSize;
  this.compressedLength += compressedSize;
  this.totalValueCount += valueCount;
  this.pageCount += 1;
  this.totalStatistics.mergeStatistics(statistics);
  // by concatenating before collecting instead of collecting twice,
  // we only allocate one buffer to copy into instead of multiple.
  buf.collect(BytesInput.concat(BytesInput.from(tempOutputStream), compressedBytes));
  encodings.add(rlEncoding);
  encodings.add(dlEncoding);
  encodings.add(valuesEncoding);
}

public void write(T value) throws IOException, InterruptedException {
  writeSupport.write(value);
  ++ recordCount;
  checkBlockSizeReached();
}

private void checkBlockSizeReached() throws IOException {
  if (recordCount >= recordCountForNextMemCheck) { // checking the memory size is relatively expensive, so let's not do it for every record.
    long memSize = columnStore.getBufferedSize();
    long recordSize = memSize / recordCount;
    // flush the row group if it is within ~2 records of the limit
    // it is much better to be slightly under size than to be over at all
    if (memSize > (nextRowGroupSize - 2 * recordSize)) {
      LOG.info(format("mem size %,d > %,d: flushing %,d records to disk.", memSize, nextRowGroupSize, recordCount));
      flushRowGroupToStore();
      initStore();
      recordCountForNextMemCheck = min(max(MINIMUM_RECORD_COUNT_FOR_CHECK, recordCount / 2), MAXIMUM_RECORD_COUNT_FOR_CHECK);
      this.lastRowGroupEndPos = parquetFileWriter.getPos();
    } else {
      recordCountForNextMemCheck = min(
          max(MINIMUM_RECORD_COUNT_FOR_CHECK, (recordCount + (long)(nextRowGroupSize / ((float)recordSize))) / 2), // will check halfway
          recordCount + MAXIMUM_RECORD_COUNT_FOR_CHECK // will not look more than max records ahead
          );
      if (DEBUG) LOG.debug(format("Checked mem at %,d will check again at: %,d ", recordCount, recordCountForNextMemCheck));
    }
  }
}

private void flushRowGroupToStore()
    throws IOException {
  LOG.info(format("Flushing mem columnStore to file. allocated memory: %,d", columnStore.getAllocatedSize()));
  if (columnStore.getAllocatedSize() > (3 * rowGroupSizeThreshold)) {
    LOG.warn("Too much memory used: " + columnStore.memUsageString());
  }

  if (recordCount > 0) {
    parquetFileWriter.startBlock(recordCount);
    columnStore.flush();
    pageStore.flushToFileWriter(parquetFileWriter);
    recordCount = 0;
    parquetFileWriter.endBlock();
    this.nextRowGroupSize = Math.min(
        parquetFileWriter.getNextRowGroupSize(),
        rowGroupSizeThreshold);
  }

  columnStore = null;
  pageStore = null;
}

/**
 * writes a number of pages at once
 * @param bytes bytes to be written including page headers
 * @param uncompressedTotalPageSize total uncompressed size (without page headers)
 * @param compressedTotalPageSize total compressed size (without page headers)
 * @throws IOException
 */
 void writeDataPages(BytesInput bytes,
                     long uncompressedTotalPageSize,
                     long compressedTotalPageSize,
                     Statistics totalStats,
                     List<Encoding> encodings) throws IOException {
  state = state.write();
  if (DEBUG) LOG.debug(out.getPos() + ": write data pages");
  long headersSize = bytes.size() - compressedTotalPageSize;
  this.uncompressedLength += uncompressedTotalPageSize + headersSize;
  this.compressedLength += compressedTotalPageSize + headersSize;
  if (DEBUG) LOG.debug(out.getPos() + ": write data pages content");
  bytes.writeAllTo(out);
  currentEncodings.addAll(encodings);
  currentStatistics = totalStats;
}

public void close() throws IOException, InterruptedException {
  flushRowGroupToStore();
  parquetFileWriter.end(extraMetaData);
}