This is an automated email from the ASF dual-hosted git repository.
lzljs3620320 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 9efb58d82 [spark] Optimize Spark write newly dynamic bucket table
(#3349)
9efb58d82 is described below
commit 9efb58d82161226d53e643ead819b910924c6643
Author: Jingsong Lee <[email protected]>
AuthorDate: Mon May 20 22:02:24 2024 +0800
[spark] Optimize Spark write newly dynamic bucket table (#3349)
---
.../org/apache/paimon/index/BucketAssigner.java | 4 +
.../apache/paimon/index/HashBucketAssigner.java | 2 +-
.../paimon/index/SimpleHashBucketAssigner.java | 21 ++-
.../paimon/index/SimpleHashBucketAssignerTest.java | 19 ++
.../java/org/apache/paimon/spark/SparkRow.java | 2 +-
.../org/apache/paimon/spark/SparkTableWrite.java | 78 ++++++++
.../paimon/spark/commands/BucketProcessor.scala | 19 +-
.../paimon/spark/commands/PaimonSparkWriter.scala | 204 ++++++++++++---------
8 files changed, 235 insertions(+), 114 deletions(-)
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/BucketAssigner.java
b/paimon-core/src/main/java/org/apache/paimon/index/BucketAssigner.java
index 9690c66c9..20c3b3987 100644
--- a/paimon-core/src/main/java/org/apache/paimon/index/BucketAssigner.java
+++ b/paimon-core/src/main/java/org/apache/paimon/index/BucketAssigner.java
@@ -27,6 +27,10 @@ public interface BucketAssigner {
void prepareCommit(long commitIdentifier);
+ static boolean isMyBucket(int bucket, int numAssigners, int assignId) {
+ return bucket % numAssigners == assignId % numAssigners;
+ }
+
static int computeHashKey(int partitionHash, int keyHash, int numChannels,
int numAssigners) {
int start = Math.abs(partitionHash % numChannels);
int id = Math.abs(keyHash % numAssigners);
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
b/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
index 114b8d3df..60bade817 100644
--- a/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
+++ b/paimon-core/src/main/java/org/apache/paimon/index/HashBucketAssigner.java
@@ -159,7 +159,7 @@ public class HashBucketAssigner implements BucketAssigner {
}
private boolean isMyBucket(int bucket) {
- return bucket % numAssigners == assignId % numAssigners;
+ return BucketAssigner.isMyBucket(bucket, numAssigners, assignId);
}
private PartitionIndex loadIndex(BinaryRow partition, int partitionHash) {
diff --git
a/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
b/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
index 7094684fc..5f7599370 100644
---
a/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
+++
b/paimon-core/src/main/java/org/apache/paimon/index/SimpleHashBucketAssigner.java
@@ -18,11 +18,13 @@
package org.apache.paimon.index;
+import org.apache.paimon.annotation.VisibleForTesting;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.utils.Int2ShortHashMap;
import java.util.HashMap;
import java.util.Map;
+import java.util.Set;
/** When we need to overwrite the table, we should use this to avoid loading
index. */
public class SimpleHashBucketAssigner implements BucketAssigner {
@@ -42,8 +44,12 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
@Override
public int assign(BinaryRow partition, int hash) {
- SimplePartitionIndex index =
- this.partitionIndex.computeIfAbsent(partition, p -> new
SimplePartitionIndex());
+ SimplePartitionIndex index = partitionIndex.get(partition);
+ if (index == null) {
+ partition = partition.copy();
+ index = new SimplePartitionIndex();
+ this.partitionIndex.put(partition, index);
+ }
return index.assign(hash);
}
@@ -52,6 +58,11 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
// do nothing
}
+ @VisibleForTesting
+ Set<BinaryRow> currentPartitions() {
+ return partitionIndex.keySet();
+ }
+
/** Simple partition bucket hash assigner. */
private class SimplePartitionIndex {
@@ -81,7 +92,7 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
private void loadNewBucket() {
for (int i = 0; i < Short.MAX_VALUE; i++) {
- if (i % numAssigners == assignId &&
!bucketInformation.containsKey(i)) {
+ if (isMyBucket(i) && !bucketInformation.containsKey(i)) {
currentBucket = i;
return;
}
@@ -90,4 +101,8 @@ public class SimpleHashBucketAssigner implements
BucketAssigner {
"Can't find a suitable bucket to assign, all the bucket
are assigned?");
}
}
+
+ private boolean isMyBucket(int bucket) {
+ return BucketAssigner.isMyBucket(bucket, numAssigners, assignId);
+ }
}
diff --git
a/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
b/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
index d795b8926..d1b26019f 100644
---
a/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
+++
b/paimon-core/src/test/java/org/apache/paimon/index/SimpleHashBucketAssignerTest.java
@@ -23,6 +23,9 @@ import org.apache.paimon.data.BinaryRow;
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;
+import static org.apache.paimon.io.DataFileTestUtils.row;
+import static org.assertj.core.api.Assertions.assertThat;
+
/** Tests for {@link SimpleHashBucketAssigner}. */
public class SimpleHashBucketAssignerTest {
@@ -66,4 +69,20 @@ public class SimpleHashBucketAssignerTest {
Assertions.assertThat(bucket).isEqualTo(0);
}
}
+
+ @Test
+ public void testPartitionCopy() {
+ SimpleHashBucketAssigner assigner = new SimpleHashBucketAssigner(1, 0,
5);
+
+ BinaryRow partition = row(1);
+ assertThat(assigner.assign(partition, 0)).isEqualTo(0);
+ assertThat(assigner.assign(partition, 1)).isEqualTo(0);
+
+ partition.setInt(0, 2);
+ assertThat(assigner.assign(partition, 5)).isEqualTo(0);
+ assertThat(assigner.assign(partition, 6)).isEqualTo(0);
+
+ assertThat(assigner.currentPartitions()).contains(row(1));
+ assertThat(assigner.currentPartitions()).contains(row(2));
+ }
}
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkRow.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkRow.java
index 5ba1a346f..c650719b0 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkRow.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkRow.java
@@ -68,7 +68,7 @@ public class SparkRow implements InternalRow, Serializable {
@Override
public int getFieldCount() {
- return row.size();
+ return type.getFieldCount();
}
@Override
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTableWrite.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTableWrite.java
new file mode 100644
index 000000000..3e8926080
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkTableWrite.java
@@ -0,0 +1,78 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.spark;
+
+import org.apache.paimon.disk.IOManager;
+import org.apache.paimon.spark.util.SparkRowUtils;
+import org.apache.paimon.table.sink.BatchTableWrite;
+import org.apache.paimon.table.sink.BatchWriteBuilder;
+import org.apache.paimon.table.sink.CommitMessage;
+import org.apache.paimon.table.sink.CommitMessageSerializer;
+import org.apache.paimon.types.RowType;
+
+import org.apache.spark.sql.Row;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+
+/** An util class for {@link BatchTableWrite}. */
+public class SparkTableWrite implements AutoCloseable {
+
+ private final BatchTableWrite write;
+ private final IOManager ioManager;
+
+ private final RowType rowType;
+ private final int rowKindColIdx;
+
+ public SparkTableWrite(BatchWriteBuilder writeBuilder, RowType rowType,
int rowKindColIdx) {
+ this.write = writeBuilder.newWrite();
+ this.rowType = rowType;
+ this.rowKindColIdx = rowKindColIdx;
+ this.ioManager = SparkUtils.createIOManager();
+ write.withIOManager(ioManager);
+ }
+
+ public void write(Row row) throws Exception {
+ write.write(toPaimonRow(row));
+ }
+
+ public void write(Row row, int bucket) throws Exception {
+ write.write(toPaimonRow(row), bucket);
+ }
+
+ public Iterator<byte[]> finish() throws Exception {
+ CommitMessageSerializer serializer = new CommitMessageSerializer();
+ List<byte[]> commitMessages = new ArrayList<>();
+ for (CommitMessage message : write.prepareCommit()) {
+ commitMessages.add(serializer.serialize(message));
+ }
+ return commitMessages.iterator();
+ }
+
+ @Override
+ public void close() throws Exception {
+ write.close();
+ ioManager.close();
+ }
+
+ private SparkRow toPaimonRow(Row row) {
+ return new SparkRow(rowType, row, SparkRowUtils.getRowKind(row,
rowKindColIdx));
+ }
+}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/BucketProcessor.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/BucketProcessor.scala
index 4af2b10b1..4a3393497 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/BucketProcessor.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/BucketProcessor.scala
@@ -76,7 +76,7 @@ case class CommonBucketProcessor(
override def next(): Row = {
val row = rowIterator.next
val sparkInternalRow = encoderGroup.rowToInternal(row)
- sparkInternalRow.setInt(bucketColIndex, getBucketId((new
SparkRow(rowType, row))))
+ sparkInternalRow.setInt(bucketColIndex, getBucketId(new
SparkRow(rowType, row)))
encoderGroup.internalToRow(sparkInternalRow)
}
}
@@ -123,20 +123,3 @@ case class DynamicBucketProcessor(
}
}
}
-
-case class UnawareBucketProcessor(bucketColIndex: Int, encoderGroup:
EncoderSerDeGroup)
- extends BucketProcessor {
-
- def processPartition(rowIterator: Iterator[Row]): Iterator[Row] = {
- new Iterator[Row] {
- override def hasNext: Boolean = rowIterator.hasNext
-
- override def next(): Row = {
- val row = rowIterator.next
- val sparkInternalRow = encoderGroup.rowToInternal(row)
- sparkInternalRow.setInt(bucketColIndex, 0)
- encoderGroup.internalToRow(sparkInternalRow)
- }
- }
- }
-}
diff --git
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala
index ae7e67fc0..b9abbb02d 100644
---
a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala
+++
b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala
@@ -18,23 +18,22 @@
package org.apache.paimon.spark.commands
-import org.apache.paimon.CoreOptions
import org.apache.paimon.CoreOptions.WRITE_ONLY
-import org.apache.paimon.index.BucketAssigner
-import org.apache.paimon.spark.SparkRow
-import org.apache.paimon.spark.SparkUtils.createIOManager
+import org.apache.paimon.index.{BucketAssigner, SimpleHashBucketAssigner}
+import org.apache.paimon.spark.{SparkRow, SparkTableWrite}
import org.apache.paimon.spark.schema.SparkSystemColumns
import org.apache.paimon.spark.schema.SparkSystemColumns.{BUCKET_COL,
ROW_KIND_COL}
import org.apache.paimon.spark.util.SparkRowUtils
import org.apache.paimon.table.{BucketMode, FileStoreTable}
import org.apache.paimon.table.sink.{BatchWriteBuilder, CommitMessage,
CommitMessageSerializer, RowPartitionKeyExtractor}
+import org.apache.paimon.utils.Preconditions
+import org.apache.paimon.utils.Preconditions.checkArgument
-import org.apache.spark.Partitioner
+import org.apache.spark.{Partitioner, TaskContext}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.apache.spark.sql.functions._
import java.io.IOException
-import java.util.Collections
import java.util.Collections.singletonMap
import scala.collection.JavaConverters._
@@ -45,12 +44,7 @@ case class PaimonSparkWriter(table: FileStoreTable) {
private lazy val rowType = table.rowType()
- private lazy val bucketMode = table match {
- case fileStoreTable: FileStoreTable =>
- fileStoreTable.bucketMode
- case _ =>
- BucketMode.HASH_FIXED
- }
+ private lazy val bucketMode = table.bucketMode
private lazy val primaryKeyCols = tableSchema.trimmedPrimaryKeys().asScala
@@ -62,80 +56,78 @@ case class PaimonSparkWriter(table: FileStoreTable) {
PaimonSparkWriter(table.copy(singletonMap(WRITE_ONLY.key(), "true")))
}
- def write(data: Dataset[_]): Seq[CommitMessage] = {
+ def write(data: Dataset[Row]): Seq[CommitMessage] = {
val sparkSession = data.sparkSession
import sparkSession.implicits._
- val dataSchema = SparkSystemColumns.filterSparkSystemColumns(data.schema)
- val rowkindColIdx = SparkRowUtils.getFieldIndex(data.schema, ROW_KIND_COL)
+ val rowKindColIdx = SparkRowUtils.getFieldIndex(data.schema, ROW_KIND_COL)
+ assert(
+ rowKindColIdx == -1 || rowKindColIdx == data.schema.length - 1,
+ "Row kind column should be the last field.")
// append _bucket_ column as placeholder
val withInitBucketCol = data.withColumn(BUCKET_COL, lit(-1))
val bucketColIdx = withInitBucketCol.schema.size - 1
-
- val originEncoderGroup = EncoderSerDeGroup(dataSchema)
val encoderGroupWithBucketCol = EncoderSerDeGroup(withInitBucketCol.schema)
- val withBucketCol =
- assignBucketId(sparkSession, withInitBucketCol, bucketColIdx,
encoderGroupWithBucketCol)
+ def newWrite(): SparkTableWrite = new SparkTableWrite(writeBuilder,
rowType, rowKindColIdx)
- val commitMessages = withBucketCol
- .mapPartitions {
+ def writeWithoutBucket(): Dataset[Array[Byte]] = {
+ data.mapPartitions {
iter =>
- val ioManager = createIOManager
- val write = writeBuilder.newWrite()
- write.withIOManager(ioManager)
- try {
- iter.foreach {
- row =>
- val bucket = row.getInt(bucketColIdx)
- val bucketColDropped =
-
originEncoderGroup.internalToRow(encoderGroupWithBucketCol.rowToInternal(row))
- val sparkRow = new SparkRow(
- rowType,
- bucketColDropped,
- SparkRowUtils.getRowKind(row, rowkindColIdx))
- write.write(sparkRow, bucket)
+ {
+ val write = newWrite()
+ try {
+ iter.foreach(row => write.write(row))
+ write.finish().asScala
+ } finally {
+ write.close()
}
- val serializer = new CommitMessageSerializer
- write.prepareCommit().asScala.map(serializer.serialize).toIterator
-
- } finally {
- write.close()
- ioManager.close()
}
}
- .collect()
- .map(deserializeCommitMessage(serializer, _))
-
- commitMessages.toSeq
- }
-
- def commit(commitMessages: Seq[CommitMessage]): Unit = {
- val tableCommit = writeBuilder.newCommit()
- try {
- tableCommit.commit(commitMessages.toList.asJava)
- } catch {
- case e: Throwable => throw new RuntimeException(e);
- } finally {
- tableCommit.close()
}
- }
- /** assign a valid bucket id for each of record. */
- private def assignBucketId(
- sparkSession: SparkSession,
- withInitBucketCol: DataFrame,
- bucketColIdx: Int,
- encoderGroupWithBucketCol: EncoderSerDeGroup): Dataset[Row] = {
+ def writeWithBucketProcessor(
+ dataFrame: DataFrame,
+ processor: BucketProcessor): Dataset[Array[Byte]] = {
+ val repartitioned = repartitionByPartitionsAndBucket(
+
dataFrame.mapPartitions(processor.processPartition)(encoderGroupWithBucketCol.encoder))
+ repartitioned.mapPartitions {
+ iter =>
+ {
+ val write = newWrite()
+ try {
+ iter.foreach(row => write.write(row, row.getInt(bucketColIdx)))
+ write.finish().asScala
+ } finally {
+ write.close()
+ }
+ }
+ }
+ }
- val encoderWithBucketCOl = encoderGroupWithBucketCol.encoder
+ def writeWithBucketAssigner(
+ dataFrame: DataFrame,
+ funcFactory: () => Row => Int): Dataset[Array[Byte]] = {
+ dataFrame.mapPartitions {
+ iter =>
+ {
+ val assigner = funcFactory.apply()
+ val write = newWrite()
+ try {
+ iter.foreach(row => write.write(row, assigner.apply(row)))
+ write.finish().asScala
+ } finally {
+ write.close()
+ }
+ }
+ }
+ }
- bucketMode match {
+ val written: Dataset[Array[Byte]] = bucketMode match {
case BucketMode.HASH_DYNAMIC =>
assert(primaryKeyCols.nonEmpty, "Only primary-key table can support
dynamic bucket.")
- // Topology: input -> shuffle by special key & partition hash ->
bucket-assigner -> shuffle by partition & bucket
val numParallelism =
Option(table.coreOptions.dynamicBucketAssignerParallelism)
.map(_.toInt)
.getOrElse {
@@ -147,42 +139,72 @@ case class PaimonSparkWriter(table: FileStoreTable) {
.map(initialBuckets => Math.min(initialBuckets.toInt,
numParallelism))
.getOrElse(numParallelism)
- val partitioned =
+ def partitionByKey(): DataFrame = {
repartitionByKeyPartitionHash(
sparkSession,
withInitBucketCol,
numParallelism,
numAssigners)
- val dynamicBucketProcessor =
- DynamicBucketProcessor(
- table,
- bucketColIdx,
- numParallelism,
- numAssigners,
- encoderGroupWithBucketCol)
- repartitionByPartitionsAndBucket(
-
partitioned.mapPartitions(dynamicBucketProcessor.processPartition)(encoderWithBucketCOl))
-
+ }
+
+ if (table.snapshotManager().latestSnapshot() == null) {
+ // bootstrap mode
+ // Topology: input -> shuffle by special key & partition hash ->
bucket-assigner
+ writeWithBucketAssigner(
+ partitionByKey(),
+ () => {
+ val extractor = new RowPartitionKeyExtractor(table.schema)
+ val assigner =
+ new SimpleHashBucketAssigner(
+ numAssigners,
+ TaskContext.getPartitionId(),
+ table.coreOptions.dynamicBucketTargetRowNum)
+ row => {
+ val sparkRow = new SparkRow(rowType, row)
+ assigner.assign(
+ extractor.partition(sparkRow),
+ extractor.trimmedPrimaryKey(sparkRow).hashCode)
+ }
+ }
+ )
+ } else {
+ // Topology: input -> shuffle by special key & partition hash ->
bucket-assigner -> shuffle by partition & bucket
+ writeWithBucketProcessor(
+ partitionByKey(),
+ DynamicBucketProcessor(
+ table,
+ bucketColIdx,
+ numParallelism,
+ numAssigners,
+ encoderGroupWithBucketCol))
+ }
case BucketMode.BUCKET_UNAWARE =>
- assert(primaryKeyCols.isEmpty, "Only append table can support unaware
bucket.")
-
- // Topology: input -> bucket-assigner
- val unawareBucketProcessor = UnawareBucketProcessor(bucketColIdx,
encoderGroupWithBucketCol)
- withInitBucketCol
-
.mapPartitions(unawareBucketProcessor.processPartition)(encoderWithBucketCOl)
- .toDF()
-
+ // Topology: input ->
+ writeWithoutBucket()
case BucketMode.HASH_FIXED =>
// Topology: input -> bucket-assigner -> shuffle by partition & bucket
- val commonBucketProcessor =
- CommonBucketProcessor(table, bucketColIdx, encoderGroupWithBucketCol)
- repartitionByPartitionsAndBucket(
-
withInitBucketCol.mapPartitions(commonBucketProcessor.processPartition)(
- encoderWithBucketCOl))
-
+ writeWithBucketProcessor(
+ withInitBucketCol,
+ CommonBucketProcessor(table, bucketColIdx,
encoderGroupWithBucketCol))
case _ =>
throw new UnsupportedOperationException(s"Spark doesn't support
$bucketMode mode.")
}
+
+ written
+ .collect()
+ .map(deserializeCommitMessage(serializer, _))
+ .toSeq
+ }
+
+ def commit(commitMessages: Seq[CommitMessage]): Unit = {
+ val tableCommit = writeBuilder.newCommit()
+ try {
+ tableCommit.commit(commitMessages.toList.asJava)
+ } catch {
+ case e: Throwable => throw new RuntimeException(e);
+ } finally {
+ tableCommit.close()
+ }
}
/** Compute bucket id in dynamic bucket mode. */
@@ -232,7 +254,7 @@ case class PaimonSparkWriter(table: FileStoreTable) {
}
}
- case class ModPartitioner(partitions: Int) extends Partitioner {
+ private case class ModPartitioner(partitions: Int) extends Partitioner {
override def numPartitions: Int = partitions
override def getPartition(key: Any): Int = key.asInstanceOf[Int] %
numPartitions
}
