com.clickzetta.platform.bulkload.ParquetRowAppender Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation

The java SDK for clickzetta's Lakehouse

There is a newer version: 2.0.0

package com.clickzetta.platform.bulkload;

import com.clickzetta.platform.arrow.ArrowTable;
import com.clickzetta.platform.client.api.Row;
import cz.proto.ingestion.v2.IngestionV2;
import org.apache.iceberg.Schema;
import org.apache.iceberg.TableProperties;
import org.apache.iceberg.data.GenericRecord;
import org.apache.iceberg.data.parquet.GenericParquetWriter;
import org.apache.iceberg.io.FileAppender;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.parquet.Parquet;

import java.io.IOException;

/**
 * A RowApplier implementation to write data to a single file.
 * It assumes the input row has the same schema as the target table.
 */
public final class ParquetRowAppender implements RowAppender {

  private IngestionV2.StreamSchema schema;
  private FileAppender appender;
  private GenericRecord record;
  private long closedLength = 0;
  private long recordCount = 0;

  public ParquetRowAppender(OutputFile outputFile, ArrowTable table) throws IOException {
    this.schema = table.getStreamSchema();
    Schema icebergSchema = SchemaConverter.convertToIcebergSchema(schema);
    record = GenericRecord.create(icebergSchema);
    appender =
      Parquet.write(outputFile)
        .schema(icebergSchema)
        .set(TableProperties.PARQUET_COMPRESSION, "zstd")
        .set(TableProperties.PARQUET_COMPRESSION_LEVEL, "1")
        .createWriterFunc(GenericParquetWriter::buildWriter)
        .overwrite(false)
        .build();
  }

  @Override
  public void append(Row... rows) {
    for (Row row : rows) {
      RecordConverter.convertToRecord(schema, row, record);
      appender.add(record);
    }
    recordCount += rows.length;
  }

  @Override
  public long length() {
    if (appender == null) {
      return closedLength;
    }
    return appender.length();
  }

  @Override
  public long recordCount() {
    return recordCount;
  }

  @Override
  public void close() throws IOException {
    if (appender != null) {
      appender.close();
      closedLength = appender.length();
      appender = null;
      record = null;
    }
  }

}