refactor DeltaTable interface

Author: harperjiang

connectors/flink/v1.20/src/main/java/io/delta/flink/DeltaTable.java

@@ -0,0 +1,50 @@
+package io.delta.flink;
+
+import io.delta.kernel.data.FilteredColumnarBatch;
+import io.delta.kernel.data.Row;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.List;
+import java.util.Map;
+
+import io.delta.kernel.expressions.Literal;
+import io.delta.kernel.types.StructType;
+import io.delta.kernel.utils.CloseableIterable;
+import io.delta.kernel.utils.CloseableIterator;
+
+/**
+ * This is an interface to access the Delta table.
+ * It contains table info such as table schema, partition columns, and also
+ * implements methods to read and write table content.
+ */
+public interface DeltaTable extends Serializable {
+
+  /**
+   * An id that can uniquely identify the table.
+   */
+  String getId();
+
+  StructType getSchema();
+
+  List<String> getPartitionColumns();
+
+  /**
+   * Commit a new version to the table.
+   * @param actions actions to be committed
+   */
+  void commit(CloseableIterable<Row> actions);
+
+  /**
+   * Write parquet files and create an associated addfile action.
+   *
+   * @param pathSuffix      string to be appended to the table path
+   * @param data            row data to be written
+   * @param partitionValues partition values
+   * @return a row representing the written action
+   */
+  CloseableIterator<Row> writeParquet(
+          String pathSuffix,
+          CloseableIterator<FilteredColumnarBatch> data,
+          Map<String, Literal> partitionValues) throws IOException;
+}

connectors/flink/v1.20/src/main/java/io/delta/flink/sink/Conversions.java

@@ -61,48 +61,49 @@ public static DataType dataType(LogicalType flinkType) {
     }
 
     public static Map<String, Literal> partitionValues(
-        RowType rowType, RowData rowData, Collection<String> partitionColNames) {
+        StructType rowType, Collection<String> partitionColNames, RowData rowData) {
       return partitionColNames.stream()
           .map(
               name -> {
-                final int partitionValueColIdx = rowType.getFieldIndex(name);
+                final int partitionValueColIdx = rowType.indexOf(name);
                 return new Object[] {
                   name, Conversions.FlinkToDelta.data(rowType, rowData, partitionValueColIdx)
                 };
               })
           .collect(Collectors.toMap(o -> (String) o[0], o -> (Literal) o[1]));
     }
 
-    public static Literal data(RowType rowType, RowData rowData, int colIdx) {
-      final LogicalType flinkType = rowType.getTypeAt(colIdx);
-      final LogicalTypeRoot typeRoot = flinkType.getTypeRoot();
-      switch (typeRoot) {
-        case INTEGER:
-          return Literal.ofInt(rowData.getInt(colIdx));
-        case BIGINT:
-          return Literal.ofLong(rowData.getLong(colIdx));
-        case VARCHAR:
-        case CHAR:
-          return Literal.ofString(rowData.getString(colIdx).toString());
-        case DOUBLE:
-          return Literal.ofDouble(rowData.getDouble(colIdx));
-        case FLOAT:
-          return Literal.ofFloat(rowData.getFloat(colIdx));
-        case DECIMAL:
-          DecimalType decimalType = (DecimalType) flinkType;
-          int precision = decimalType.getPrecision();
-          int scale = decimalType.getScale();
-          return Literal.ofDecimal(
-              rowData.getDecimal(colIdx, precision, scale).toBigDecimal(), precision, scale);
-        case DATE:
-          return Literal.ofDate(rowData.getInt(colIdx));
-        case TIMESTAMP_WITH_TIME_ZONE:
-          return Literal.ofTimestamp(rowData.getLong(colIdx));
-        case TIMESTAMP_WITHOUT_TIME_ZONE:
-          return Literal.ofTimestampNtz(rowData.getLong(colIdx));
-        default:
-          throw new UnsupportedOperationException(
-              String.format("Type not supported: %s", flinkType));
+    public static Literal data(StructType rowType, RowData rowData, int colIdx) {
+      final StructField field = rowType.at(colIdx);
+      final DataType dataType = field.getDataType();
+      if (dataType.equivalent(IntegerType.INTEGER)) {
+        return Literal.ofInt(rowData.getInt(colIdx));
+      } else if (dataType.equivalent(LongType.LONG)) {
+        return Literal.ofLong(rowData.getLong(colIdx));
+      } else if (dataType.equivalent(StringType.STRING)) {
+        return Literal.ofString(rowData.getString(colIdx).toString());
+      } else if (dataType.equivalent(DoubleType.DOUBLE)) {
+        return Literal.ofDouble(rowData.getDouble(colIdx));
+      } else if (dataType.equivalent(FloatType.FLOAT)) {
+        return Literal.ofFloat(rowData.getFloat(colIdx));
+      } else if (dataType instanceof io.delta.kernel.types.DecimalType) {
+        io.delta.kernel.types.DecimalType decimalType =
+                (io.delta.kernel.types.DecimalType) dataType;
+        int precision = decimalType.getPrecision();
+        int scale = decimalType.getScale();
+        return Literal.ofDecimal(
+                rowData.getDecimal(colIdx, precision, scale).toBigDecimal(),
+                precision,
+                scale
+        );
+      } else if (dataType.equivalent(DateType.DATE)) {
+        return Literal.ofDate(rowData.getInt(colIdx));
+      } else if (dataType.equivalent(TimestampType.TIMESTAMP)) {
+        return Literal.ofTimestamp(rowData.getLong(colIdx));
+      } else if (dataType.equivalent(TimestampNTZType.TIMESTAMP_NTZ)) {
+        return Literal.ofTimestampNtz(rowData.getLong(colIdx));
+      } else {
+        throw new UnsupportedOperationException("Unsupported data type: " + dataType);
       }
     }
   }

connectors/flink/v1.20/src/main/java/io/delta/flink/sink/DeltaCommittable.java

@@ -4,11 +4,12 @@
 import io.delta.kernel.defaults.internal.json.JsonUtils;
 import io.delta.kernel.internal.actions.SingleAction;
 import io.delta.kernel.internal.util.Preconditions;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+
 import java.io.*;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
-import org.apache.flink.core.io.SimpleVersionedSerializer;
 
 public class DeltaCommittable {
   private final String jobId;

connectors/flink/v1.20/src/main/java/io/delta/flink/sink/DeltaCommitter.java

@@ -1,24 +1,23 @@
 package io.delta.flink.sink;
 
-import io.delta.kernel.*;
+import io.delta.flink.DeltaTable;
 import io.delta.kernel.data.Row;
-import io.delta.kernel.defaults.internal.json.JsonUtils;
 import io.delta.kernel.engine.Engine;
-import io.delta.kernel.exceptions.TableNotFoundException;
-import io.delta.kernel.internal.data.TransactionStateRow;
-import io.delta.kernel.internal.util.Preconditions;
 import io.delta.kernel.internal.util.Utils;
-import io.delta.kernel.types.StructType;
 import io.delta.kernel.utils.CloseableIterable;
 import io.delta.kernel.utils.CloseableIterator;
-import java.io.IOException;
-import java.util.*;
-import java.util.stream.Collectors;
 import org.apache.flink.api.connector.sink2.Committer;
 import org.apache.flink.metrics.groups.SinkCommitterMetricGroup;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.Objects;
+import java.util.TreeMap;
+import java.util.stream.Collectors;
+
 /**
  * The Committer is responsible for committing the data staged by the CommittingSinkWriter in the
  * second step of a two-phase commit protocol.
@@ -51,38 +50,23 @@ public class DeltaCommitter implements Committer<DeltaCommittable> {
   // All committables should have the same job id as the committer.
   // For simplicity, we get the job id from constructor.
   private String jobId;
-  private Engine engine;
-  private Table table;
-  private final Row committerContext;
+  private DeltaTable deltaTable;
 
   private SinkCommitterMetricGroup metricGroup;
 
-  private boolean creatingNewTable;
-
   private DeltaCommitter(
-      String jobId,
-      Engine engine,
-      Table table,
-      Row committerContext,
-      SinkCommitterMetricGroup metricGroup) {
+          String jobId,
+          DeltaTable deltaTable,
+          SinkCommitterMetricGroup metricGroup) {
     this.jobId = jobId;
-    this.engine = engine;
-    this.table = table;
-    this.committerContext = committerContext;
-
+    this.deltaTable = deltaTable;
     this.metricGroup = metricGroup;
   }
 
   @Override
   public void commit(Collection<CommitRequest<DeltaCommittable>> committables)
       throws IOException, InterruptedException {
     LOG.debug("Starting commit");
-    try {
-      table.getLatestSnapshot(engine);
-      creatingNewTable = false;
-    } catch (TableNotFoundException e) {
-      creatingNewTable = true;
-    }
     sortCommittablesByCheckpointId(committables).forEach(this::commitForSingleCheckpointId);
   }
 
@@ -93,44 +77,6 @@ private void commitForSingleCheckpointId(
       long checkpointId, List<CommitRequest<DeltaCommittable>> committables) {
     LOG.debug("Committing {} committables on checkpoint {}", committables.size(), checkpointId);
 
-    TransactionBuilder txnBuilder =
-        table
-            .createTransactionBuilder(
-                engine,
-                "DeltaSink/Kernel",
-                creatingNewTable ? Operation.CREATE_TABLE : Operation.WRITE)
-            .withTransactionId(engine, jobId, checkpointId);
-
-    if (creatingNewTable) {
-      // For a new table set the table schema in the transaction builder
-      txnBuilder =
-          txnBuilder
-              .withSchema(engine, TransactionStateRow.getLogicalSchema(committerContext))
-              .withPartitionColumns(
-                  engine, TransactionStateRow.getPartitionColumnsList(committerContext));
-    }
-    final Transaction txn = txnBuilder.build(engine);
-
-    // We check the table's latest schema is still the same as committer schema.
-    // The check is delayed here to detect external modification to the table schema.
-    if (!creatingNewTable) {
-      final Snapshot readSnapshot = table.getSnapshotAsOfVersion(engine, txn.getReadTableVersion());
-      final StructType tableSchema = txn.getSchema(engine);
-      final StructType committerSchema = TransactionStateRow.getLogicalSchema(committerContext);
-      Preconditions.checkArgument(
-          readSnapshot.getPath().equals(TransactionStateRow.getTablePath(this.committerContext)),
-          String.format(
-              "Committer path does not match the latest table path."
-                  + "Table path: %s, Committer path: %s",
-              readSnapshot.getPath(), TransactionStateRow.getTablePath(this.committerContext)));
-      Preconditions.checkArgument(
-          committerSchema.equivalent(tableSchema),
-          String.format(
-              "DeltaSink does not support schema evolution. "
-                  + "Table schema: %s, Committer schema: %s",
-              tableSchema, committerSchema));
-    }
-
     final CloseableIterable<Row> dataActions =
         new CloseableIterable<Row>() {
           @Override
@@ -147,8 +93,7 @@ public void close() throws IOException {
           }
         };
 
-    txn.commit(engine, dataActions);
-    creatingNewTable = false;
+    deltaTable.commit(dataActions);
   }
 
   private TreeMap<Long, List<CommitRequest<DeltaCommittable>>> sortCommittablesByCheckpointId(
@@ -163,9 +108,7 @@ private TreeMap<Long, List<CommitRequest<DeltaCommittable>>> sortCommittablesByC
 
   public static final class Builder {
     private String jobId;
-    private Engine engine;
-    private Table table;
-    private Row committerContext;
+    private DeltaTable deltaTable;
     private SinkCommitterMetricGroup metricGroup;
 
     public Builder() {}
@@ -175,24 +118,8 @@ public Builder withJobId(String jobId) {
       return this;
     }
 
-    public Builder withEngine(Engine engine) {
-      this.engine = engine;
-      return this;
-    }
-
-    public Builder withTable(Table table) {
-      this.table = table;
-      return this;
-    }
-
-    public Builder withCommitterContext(Row committerContext) {
-      this.committerContext = committerContext;
-      return this;
-    }
-
-    public Builder withCommitterContext(String committerContextJson) {
-      this.committerContext =
-          JsonUtils.rowFromJson(committerContextJson, TransactionStateRow.SCHEMA);
+    public Builder withDeltaTable(DeltaTable deltaTable) {
+      this.deltaTable = deltaTable;
       return this;
     }
 
@@ -203,12 +130,10 @@ public Builder withMetricGroup(SinkCommitterMetricGroup metricGroup) {
 
     public DeltaCommitter build() {
       Objects.requireNonNull(jobId, "jobId must not be null");
-      Objects.requireNonNull(engine, "engine must not be null");
-      Objects.requireNonNull(table, "table must not be null");
-      Objects.requireNonNull(committerContext, "committerContext must not be null");
+      Objects.requireNonNull(deltaTable, "tableLoader must not be null");
       Objects.requireNonNull(metricGroup, "metricGroup must not be null");
 
-      return new DeltaCommitter(jobId, engine, table, committerContext, metricGroup);
+      return new DeltaCommitter(jobId, deltaTable, metricGroup);
     }
   }
 }