This is an automated email from the ASF dual-hosted git repository.
JingsongLi 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 f99d00837d [spark] Support parallel sorted index building across
Paimon partitions (#9491)
f99d00837d is described below
commit f99d00837d840325d5f9091cbcec0118a173cfef
Author: liangjie <[email protected]>
AuthorDate: Mon Aug 31 17:25:45 2026 +0800
[spark] Support parallel sorted index building across Paimon partitions
(#9491)
---
.../globalindex/sorted/SortedIndexTopoBuilder.java | 200 +++++++++++++++------
.../sorted/SortedIndexTopoBuilderTest.java | 118 ++++++++++++
.../procedure/CreateGlobalIndexProcedureTest.scala | 33 +++-
3 files changed, 293 insertions(+), 58 deletions(-)
diff --git
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilder.java
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilder.java
index 62bfb157d3..0ee778de3e 100644
---
a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilder.java
+++
b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilder.java
@@ -61,6 +61,7 @@ import
org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation;
import org.apache.spark.sql.functions;
import java.io.IOException;
+import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -144,22 +145,23 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
SortedGlobalIndexWriter indexWriter =
new SortedGlobalIndexWriter(table, indexType, options)
.withIndexField(indexField.name());
- for (Map.Entry<BinaryRow, Map<Range, List<Split>>> partitionEntry :
- partitionRangeSplits.entrySet()) {
- for (Map.Entry<Range, List<Split>> entry :
partitionEntry.getValue().entrySet()) {
- Range range = entry.getKey();
- List<Split> rangeSplits = entry.getValue();
- if (rangeSplits.isEmpty()) {
- continue;
- }
-
- final byte[] serializedWriter =
InstantiationUtil.serializeObject(indexWriter);
- final byte[] partitionBytes =
+ final byte[] serializedWriter =
InstantiationUtil.serializeObject(indexWriter);
+ if (keyExtractor.isIdentity()) {
+ List<SortedBuildTask> buildTasks = new ArrayList<>();
+ List<Dataset<Row>> taskInputs = new ArrayList<>();
+ for (Map.Entry<BinaryRow, Map<Range, List<Split>>> partitionEntry :
+ partitionRangeSplits.entrySet()) {
+ byte[] partitionBytes =
binaryRowSerializer.serializeToBytes(partitionEntry.getKey());
- if (keyExtractor.isIdentity()) {
- int partitionNum = Math.max((int) (range.count() /
recordsPerRange), 1);
- partitionNum = Math.min(partitionNum, maxParallelism);
+ for (Map.Entry<Range, List<Split>> entry :
partitionEntry.getValue().entrySet()) {
+ Range range = entry.getKey();
+ List<Split> rangeSplits = entry.getValue();
+ if (rangeSplits.isEmpty()) {
+ continue;
+ }
+ long taskId = buildTasks.size();
+ buildTasks.add(new SortedBuildTask(taskId, range,
partitionBytes));
Dataset<Row> source =
PaimonUtils.createDataset(
spark,
@@ -167,35 +169,59 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
rangeSplits.toArray(new Split[0]),
relation));
Dataset<Row> selected =
source.select(
- readType.getFieldNames().stream()
- .map(functions::col)
- .toArray(Column[]::new));
- Column[] sortFields =
- new Column[] {
- functions.col(indexField.name()),
- functions.col(SpecialFields.ROW_ID.name())
- };
- Dataset<Row> partitioned =
- selected.repartitionByRange(partitionNum,
sortFields)
- .sortWithinPartitions(sortFields);
- JavaRDD<byte[]> written =
- partitioned
- .javaRDD()
- .map(row -> (InternalRow) (new
SparkRow(readType, row)))
- .mapPartitions(
-
(FlatMapFunction<Iterator<InternalRow>, byte[]>)
- iter ->
- buildSortedIndex(
- iter,
-
serializedWriter,
- range,
-
partitionKeyNum,
-
partitionBytes,
-
scanSnapshotId));
-
allMessages.addAll(CommitMessageSerializer.deserializeAll(written.collect()));
+ readType.getFieldNames().stream()
+ .map(functions::col)
+ .toArray(Column[]::new))
+ .withColumn(taskIdField,
functions.lit(taskId).cast("long"));
+ taskInputs.add(
+ selected.select(
+ functions.col(taskIdField),
+ functions.col(indexField.name()),
+
functions.col(SpecialFields.ROW_ID.name())));
+ }
+ }
+
+ if (!buildTasks.isEmpty()) {
+ int partitionNum =
+ calculateParallelism(buildTasks, recordsPerRange,
maxParallelism);
+ Dataset<Row> partitioned =
+ combineAndSortBuildTaskInputs(
+ taskInputs, partitionNum, taskIdField,
indexField.name());
+ Map<Long, SortedBuildTask> buildTasksById = new HashMap<>();
+ for (SortedBuildTask task : buildTasks) {
+ buildTasksById.put(task.taskId, task);
+ }
+ JavaRDD<byte[]> written =
+ partitioned
+ .javaRDD()
+ .map(row -> (InternalRow) (new
SparkRow(normalizedReadType, row)))
+ .mapPartitions(
+
(FlatMapFunction<Iterator<InternalRow>, byte[]>)
+ iter ->
+ buildSortedIndexes(
+ iter,
+
serializedWriter,
+ buildTasksById,
+
partitionKeyNum,
+
scanSnapshotId));
+
allMessages.addAll(CommitMessageSerializer.deserializeAll(written.collect()));
+ }
+ addDeletedIndexMessages(allMessages,
scanResult.deletedIndexEntries());
+ return allMessages;
+ }
+
+ for (Map.Entry<BinaryRow, Map<Range, List<Split>>> partitionEntry :
+ partitionRangeSplits.entrySet()) {
+ for (Map.Entry<Range, List<Split>> entry :
partitionEntry.getValue().entrySet()) {
+ Range range = entry.getKey();
+ List<Split> rangeSplits = entry.getValue();
+ if (rangeSplits.isEmpty()) {
continue;
}
+ final byte[] partitionBytes =
+
binaryRowSerializer.serializeToBytes(partitionEntry.getKey());
+
Map<Range, List<Split>> shardedSplits =
shardSplitsByRowRange(
Collections.singletonMap(range, rangeSplits),
recordsPerRange);
@@ -256,7 +282,7 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
.mapPartitions(
(FlatMapFunction<Iterator<InternalRow>, byte[]>)
iter ->
- buildSortedIndexes(
+
buildShardedSortedIndexes(
iter,
serializedWriter,
taskRanges,
@@ -267,8 +293,34 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
allMessages.addAll(CommitMessageSerializer.deserializeAll(commitBytes));
}
}
- for (IndexManifestEntry entry : scanResult.deletedIndexEntries()) {
- allMessages.add(
+ addDeletedIndexMessages(allMessages, scanResult.deletedIndexEntries());
+ return allMessages;
+ }
+
+ static Dataset<Row> combineAndSortBuildTaskInputs(
+ List<Dataset<Row>> taskInputs,
+ int partitionNum,
+ String taskIdField,
+ String indexField) {
+ Dataset<Row> combined = taskInputs.get(0);
+ for (int i = 1; i < taskInputs.size(); i++) {
+ combined = combined.union(taskInputs.get(i));
+ }
+
+ Column[] sortFields =
+ new Column[] {
+ functions.col(taskIdField),
+ functions.col(indexField),
+ functions.col(SpecialFields.ROW_ID.name())
+ };
+ return combined.repartitionByRange(partitionNum, sortFields)
+ .sortWithinPartitions(sortFields);
+ }
+
+ private static void addDeletedIndexMessages(
+ List<CommitMessage> messages, List<IndexManifestEntry>
deletedIndexEntries) {
+ for (IndexManifestEntry entry : deletedIndexEntries) {
+ messages.add(
new CommitMessageImpl(
entry.partition(),
entry.bucket(),
@@ -277,28 +329,40 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
Collections.singletonList(entry.indexFile())),
CompactIncrement.emptyIncrement()));
}
- return allMessages;
}
- private static Iterator<byte[]> buildSortedIndex(
+ private static Iterator<byte[]> buildSortedIndexes(
Iterator<InternalRow> input,
byte[] serializedWriter,
- Range range,
+ Map<Long, SortedBuildTask> buildTasksById,
int partitionKeyNum,
- byte[] partitionBytes,
long scanSnapshotId)
throws IOException, ClassNotFoundException {
final BinaryRowSerializer binaryRowSerializer = new
BinaryRowSerializer(partitionKeyNum);
- BinaryRow partition =
binaryRowSerializer.deserializeFromBytes(partitionBytes);
SortedGlobalIndexWriter writer =
InstantiationUtil.deserializeObject(
serializedWriter,
SortedGlobalIndexWriter.class.getClassLoader());
- return CommitMessageSerializer.serializeAll(
- writer.buildForSinglePartition(range, partition,
input, scanSnapshotId))
- .iterator();
+ SortedTaskInput taskInput = new SortedTaskInput(input,
writer.keyExtractor().keyType());
+ List<byte[]> results = new ArrayList<>();
+ while (taskInput.hasTask()) {
+ long taskId = taskInput.taskId();
+ SortedBuildTask task = buildTasksById.get(taskId);
+ if (task == null) {
+ throw new IllegalArgumentException("Unknown sorted index build
task id: " + taskId);
+ }
+ BinaryRow partition =
binaryRowSerializer.deserializeFromBytes(task.partition);
+ results.addAll(
+ CommitMessageSerializer.serializeAll(
+ writer.buildForSinglePartition(
+ task.rowRange,
+ partition,
+ taskInput.consumeTask(taskId),
+ scanSnapshotId)));
+ }
+ return results.iterator();
}
- private static Iterator<byte[]> buildSortedIndexes(
+ private static Iterator<byte[]> buildShardedSortedIndexes(
Iterator<InternalRow> input,
byte[] serializedWriter,
Map<Long, Range> taskRanges,
@@ -330,6 +394,22 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
return results.iterator();
}
+ static int calculateParallelism(
+ List<SortedBuildTask> buildTasks, long recordsPerRange, int
maxParallelism) {
+ long totalRecords = 0;
+ for (SortedBuildTask task : buildTasks) {
+ long count = task.rowRange.count();
+ if (Long.MAX_VALUE - totalRecords < count) {
+ totalRecords = Long.MAX_VALUE;
+ } else {
+ totalRecords += count;
+ }
+ }
+
+ long parallelism = Math.max(totalRecords / recordsPerRange, 1);
+ return (int) Math.min(parallelism, maxParallelism);
+ }
+
private static String buildTaskIdFieldName(RowType readType) {
String fieldName = BUILD_TASK_ID_FIELD;
while (readType.containsField(fieldName)) {
@@ -349,6 +429,22 @@ public class SortedIndexTopoBuilder implements
GlobalIndexTopologyBuilder {
readType.getField(SpecialFields.ROW_ID.name()));
}
+ /** Metadata for one sorted index build range. */
+ static class SortedBuildTask implements Serializable {
+
+ private static final long serialVersionUID = 1L;
+
+ private final long taskId;
+ private final Range rowRange;
+ private final byte[] partition;
+
+ SortedBuildTask(long taskId, Range rowRange, byte[] partition) {
+ this.taskId = taskId;
+ this.rowRange = rowRange;
+ this.partition = partition;
+ }
+ }
+
/** Groups a partition already sorted by task id, normalized key and row
id. */
private static class SortedTaskInput {
diff --git
a/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilderTest.java
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilderTest.java
new file mode 100644
index 0000000000..bddd63bee2
--- /dev/null
+++
b/paimon-spark/paimon-spark-common/src/test/java/org/apache/paimon/spark/globalindex/sorted/SortedIndexTopoBuilderTest.java
@@ -0,0 +1,118 @@
+/*
+ * 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.globalindex.sorted;
+
+import
org.apache.paimon.spark.globalindex.sorted.SortedIndexTopoBuilder.SortedBuildTask;
+import org.apache.paimon.table.SpecialFields;
+import org.apache.paimon.utils.Range;
+
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.catalyst.plans.logical.RepartitionByExpression;
+import org.apache.spark.sql.catalyst.plans.logical.Sort;
+import org.apache.spark.sql.catalyst.plans.logical.Union;
+import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning;
+import org.apache.spark.sql.functions;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link SortedIndexTopoBuilder}. */
+public class SortedIndexTopoBuilderTest {
+
+ private static final String TASK_ID = "task_id";
+ private static final String INDEX_KEY = "index_key";
+
+ @Test
+ void testBuildTopologyAcrossPartitions() {
+ SparkSession spark =
+ SparkSession.builder()
+ .master("local[1]")
+ .appName("sorted-index-topology-test")
+ .config("spark.ui.enabled", "false")
+ .getOrCreate();
+ try {
+ List<Dataset<Row>> partitionInputs =
+ Arrays.asList(taskInput(spark, 0), taskInput(spark, 1),
taskInput(spark, 2));
+
+ Dataset<Row> topology =
+ SortedIndexTopoBuilder.combineAndSortBuildTaskInputs(
+ partitionInputs, 3, TASK_ID, INDEX_KEY);
+
+
assertThat(topology.queryExecution().logical()).isInstanceOf(Sort.class);
+ Sort sort = (Sort) topology.queryExecution().logical();
+ assertThat(sort.global()).isFalse();
+ assertThat(sort.order().size()).isEqualTo(3);
+
assertThat(sort.child()).isInstanceOf(RepartitionByExpression.class);
+
+ RepartitionByExpression repartition = (RepartitionByExpression)
sort.child();
+ assertThat(repartition.numPartitions()).isEqualTo(3);
+ assertThat(repartition.shuffle()).isTrue();
+
assertThat(repartition.partitioning()).isInstanceOf(RangePartitioning.class);
+ assertThat(repartition.partitionExpressions().size()).isEqualTo(3);
+ assertThat(repartition.child()).isInstanceOf(Union.class);
+ assertThat(((Union)
repartition.child()).children().size()).isEqualTo(3);
+ } finally {
+ spark.stop();
+ SparkSession.clearActiveSession();
+ SparkSession.clearDefaultSession();
+ }
+ }
+
+ @Test
+ void testCalculateParallelismByTotalRowsInsteadOfRangeCount() {
+ List<SortedBuildTask> tasks = new ArrayList<>();
+ for (int i = 0; i < 100; i++) {
+ tasks.add(new SortedBuildTask(i, new Range(i * 10L, i * 10L + 9),
new byte[0]));
+ }
+
+ assertThat(SortedIndexTopoBuilder.calculateParallelism(tasks, 1000L,
4096)).isEqualTo(1);
+ }
+
+ @Test
+ void testCalculateParallelismHonorsMaxParallelism() {
+ List<SortedBuildTask> tasks = new ArrayList<>();
+ for (int i = 0; i < 100; i++) {
+ tasks.add(new SortedBuildTask(i, new Range(i * 1000L, i * 1000L +
999), new byte[0]));
+ }
+
+ assertThat(SortedIndexTopoBuilder.calculateParallelism(tasks, 1000L,
16)).isEqualTo(16);
+ }
+
+ @Test
+ void testCalculateParallelismKeepsSingleRangeBehavior() {
+ List<SortedBuildTask> tasks = new ArrayList<>();
+ tasks.add(new SortedBuildTask(0, new Range(0, 1499), new byte[0]));
+
+ assertThat(SortedIndexTopoBuilder.calculateParallelism(tasks, 1000L,
16)).isEqualTo(1);
+ }
+
+ private static Dataset<Row> taskInput(SparkSession spark, long taskId) {
+ return spark.range(1)
+ .select(
+ functions.lit(taskId).cast("long").alias(TASK_ID),
+ functions.lit(taskId).alias(INDEX_KEY),
+
functions.col("id").alias(SpecialFields.ROW_ID.name()));
+ }
+}
diff --git
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/CreateGlobalIndexProcedureTest.scala
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/CreateGlobalIndexProcedureTest.scala
index 94d836a613..90c1eda5cb 100644
---
a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/CreateGlobalIndexProcedureTest.scala
+++
b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/CreateGlobalIndexProcedureTest.scala
@@ -26,6 +26,7 @@ import org.apache.paimon.spark.PaimonSparkTestBase
import org.apache.paimon.types.VarCharType
import org.apache.paimon.utils.Range
+import org.apache.spark.scheduler.{SparkListener, SparkListenerStageSubmitted}
import org.apache.spark.sql.Row
import org.apache.spark.sql.paimon.Utils
import org.apache.spark.sql.streaming.StreamTest
@@ -34,6 +35,7 @@ import java.io.File
import scala.collection.JavaConverters._
import scala.collection.immutable
+import scala.collection.mutable
class CreateGlobalIndexProcedureTest extends PaimonSparkTestBase with
StreamTest {
@@ -322,16 +324,35 @@ class CreateGlobalIndexProcedureTest extends
PaimonSparkTestBase with StreamTest
values = (0 until 33333).map(i => s"($i, 'name_$i', 'p2')").mkString(",")
spark.sql(s"INSERT INTO T VALUES $values")
+ val submittedStageTasks = mutable.ListBuffer.empty[Int]
+ val listener = new SparkListener {
+ override def onStageSubmitted(stageSubmitted:
SparkListenerStageSubmitted): Unit = {
+ submittedStageTasks += stageSubmitted.stageInfo.numTasks
+ }
+ }
val output =
- spark
- .sql(
- "CALL sys.create_global_index(table => 'test.T', index_column =>
'name', index_type => 'btree'," +
- " options => 'btree-index.records-per-range=1000')")
- .collect()
- .head
+ try {
+ spark.sparkContext.addSparkListener(listener)
+ spark
+ .sql(
+ "CALL sys.create_global_index(table => 'test.T', index_column =>
'name', index_type => 'btree'," +
+ " options => 'btree-index.records-per-range=1000')")
+ .collect()
+ .head
+ } finally {
+ Utils.waitUntilEventEmpty(spark)
+ spark.sparkContext.removeSparkListener(listener)
+ }
assert(output.getBoolean(0))
+ val expectedBuildParallelism = (189088L / 1000).toInt
+ assert(
+ submittedStageTasks.count(_ == expectedBuildParallelism) == 1,
+ s"Expected one global build stage with $expectedBuildParallelism
tasks, " +
+ s"but observed stages with ${submittedStageTasks.mkString(", ")}
tasks"
+ )
+
assertMultiplePartitionsResult("T", 189088L, 3)
}
}
