com.netease.arctic.spark.io.ArcticSparkChangeTaskWriter Maven / Gradle / Ivy

Go to download
package com.netease.arctic.spark.io;

import com.netease.arctic.data.ChangeAction;
import com.netease.arctic.io.ArcticFileIO;
import com.netease.arctic.io.writer.ChangeTaskWriter;
import com.netease.arctic.io.writer.OutputFileFactory;
import com.netease.arctic.spark.SparkInternalRowCastWrapper;
import com.netease.arctic.spark.SparkInternalRowWrapper;
import com.netease.arctic.table.PrimaryKeySpec;
import com.netease.arctic.shade.org.apache.iceberg.FileFormat;
import com.netease.arctic.shade.org.apache.iceberg.PartitionSpec;
import com.netease.arctic.shade.org.apache.iceberg.Schema;
import com.netease.arctic.shade.org.apache.iceberg.StructLike;
import com.netease.arctic.shade.org.apache.iceberg.io.FileAppenderFactory;
import com.netease.arctic.shade.org.apache.iceberg.spark.SparkSchemaUtil;
import com.netease.arctic.shade.org.apache.spark.sql.catalyst.InternalRow;

/**
 * change task writer
 */
public class ArcticSparkChangeTaskWriter extends ChangeTaskWriter {
  private final Schema schema;

  protected ArcticSparkChangeTaskWriter(
      FileFormat format,
      FileAppenderFactory appenderFactory,
      OutputFileFactory outputFileFactory,
      ArcticFileIO io,
      long targetFileSize,
      long mask,
      Schema schema,
      PartitionSpec spec,
      PrimaryKeySpec primaryKeySpec,
      boolean orderedWriter
  ) {
    super(format, appenderFactory, outputFileFactory, io,
        targetFileSize, mask, schema, spec, primaryKeySpec, orderedWriter);
    this.schema = schema;
  }

  @Override
  protected StructLike asStructLike(InternalRow data) {
    return new SparkInternalRowWrapper(SparkSchemaUtil.convert(schema)).wrap(data);
  }

  @Override
  protected InternalRow appendMetaColumns(InternalRow data, Long fileOffset) {
    SparkInternalRowCastWrapper row = (SparkInternalRowCastWrapper) data;
    return row.setFileOffset(fileOffset);
  }

  @Override
  protected ChangeAction action(InternalRow data) {
    SparkInternalRowCastWrapper row = (SparkInternalRowCastWrapper) data;
    return row.getChangeAction();
  }

}