This is an automated email from the ASF dual-hosted git repository.
danny0405 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new 2d73fe40dd8a feat(spark): support bucket index for LSM tables (#19727)
2d73fe40dd8a is described below
commit 2d73fe40dd8a106af58454139484cbff64aca899
Author: Shuo Cheng <[email protected]>
AuthorDate: Fri Aug 28 11:24:46 2026 +0800
feat(spark): support bucket index for LSM tables (#19727)
* feat(spark): support bucket index for LSM tables
---
.../table/BucketSortBulkInsertPartitioner.java | 6 ++
.../BucketIndexBulkInsertPartitionerWithRows.java | 37 +++++--
.../bulkinsert/RDDBucketIndexPartitioner.java | 5 +-
.../apache/spark/sql/BucketPartitionUtils.scala | 42 +++++++-
.../BaseDatasetBulkInsertCommitActionExecutor.java | 4 +-
.../DatasetBucketRescaleCommitActionExecutor.java | 4 +-
.../bulkinsert/TestLSMBulkInsertPartitioner.java | 40 +++++++
.../apache/hudi/functional/TestLSMDataSource.scala | 120 ++++++++++++++++++++-
8 files changed, 242 insertions(+), 16 deletions(-)
diff --git
a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/table/BucketSortBulkInsertPartitioner.java
b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/table/BucketSortBulkInsertPartitioner.java
index 9dbb2e96be36..d9b4cee7bfe8 100644
---
a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/table/BucketSortBulkInsertPartitioner.java
+++
b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/table/BucketSortBulkInsertPartitioner.java
@@ -19,6 +19,7 @@
package org.apache.hudi.table;
import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.execution.bulkinsert.BulkInsertSortMode;
/**
@@ -32,6 +33,11 @@ public abstract class BucketSortBulkInsertPartitioner<T>
implements BulkInsertPa
public BucketSortBulkInsertPartitioner(HoodieTable table, String sortString)
{
this.table = table;
+ if (table.getMetaClient().getTableConfig().isLSMTreeStorageLayout()
+ && !StringUtils.isNullOrEmpty(sortString)) {
+ throw new HoodieException("Custom sort columns are not supported for
bucket index on LSM tables because "
+ + "LSM files must be ordered by record key");
+ }
if (!StringUtils.isNullOrEmpty(sortString)) {
this.sortColumnNames = sortString.split(",");
} else {
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
index 520427dbd458..39e8b0e768ba 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
@@ -36,20 +36,44 @@ public class BucketIndexBulkInsertPartitionerWithRows
implements BulkInsertParti
private final String indexKeyFields;
private final NumBucketsFunction numBucketsFunction;
private final HoodieWriteConfig writeConfig;
+ private final boolean sortByRecordKey;
private FileSystemViewStorageConfig viewConfig;
- public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields,
HoodieWriteConfig writeConfig) {
- this(writeConfig, NumBucketsFunction.fromWriteConfig(writeConfig),
indexKeyFields);
+ public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields,
+ HoodieWriteConfig
writeConfig) {
+ this(indexKeyFields, writeConfig, false);
}
- public BucketIndexBulkInsertPartitionerWithRows(HoodieWriteConfig
writeConfig, String expressions, String rule, int bucketNumber) {
- this(writeConfig, new NumBucketsFunction(expressions, rule, bucketNumber),
writeConfig.getBucketIndexHashFieldWithDefault());
+ public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields,
+ HoodieWriteConfig
writeConfig,
+ boolean sortByRecordKey) {
+ this(writeConfig, NumBucketsFunction.fromWriteConfig(writeConfig),
indexKeyFields, sortByRecordKey);
}
- private BucketIndexBulkInsertPartitionerWithRows(HoodieWriteConfig
writeConfig, NumBucketsFunction numBucketsFunction, String indexKeyFields) {
+ public BucketIndexBulkInsertPartitionerWithRows(HoodieWriteConfig
writeConfig,
+ String expressions,
+ String rule,
+ int bucketNumber) {
+ this(writeConfig, expressions, rule, bucketNumber, false);
+ }
+
+ public BucketIndexBulkInsertPartitionerWithRows(HoodieWriteConfig
writeConfig,
+ String expressions,
+ String rule,
+ int bucketNumber,
+ boolean sortByRecordKey) {
+ this(writeConfig, new NumBucketsFunction(expressions, rule, bucketNumber),
+ writeConfig.getBucketIndexHashFieldWithDefault(), sortByRecordKey);
+ }
+
+ private BucketIndexBulkInsertPartitionerWithRows(HoodieWriteConfig
writeConfig,
+ NumBucketsFunction
numBucketsFunction,
+ String indexKeyFields,
+ boolean sortByRecordKey) {
this.indexKeyFields = indexKeyFields;
this.numBucketsFunction = numBucketsFunction;
this.writeConfig = writeConfig;
+ this.sortByRecordKey = sortByRecordKey;
if (writeConfig.isUsingRemotePartitioner()) {
this.viewConfig = writeConfig.getViewStorageConfig();
}
@@ -60,7 +84,8 @@ public class BucketIndexBulkInsertPartitionerWithRows
implements BulkInsertParti
Partitioner partitioner = writeConfig.isUsingRemotePartitioner() &&
writeConfig.isEmbeddedTimelineServerEnabled()
? BucketPartitionUtils$.MODULE$.getRemotePartitioner(viewConfig,
numBucketsFunction, outputPartitions)
:
BucketPartitionUtils$.MODULE$.getLocalePartitioner(numBucketsFunction,
outputPartitions);
- return BucketPartitionUtils$.MODULE$.createDataFrame(rows, indexKeyFields,
numBucketsFunction, partitioner);
+ return BucketPartitionUtils$.MODULE$.createDataFrame(
+ rows, indexKeyFields, numBucketsFunction, partitioner,
sortByRecordKey);
}
@Override
diff --git
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDBucketIndexPartitioner.java
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDBucketIndexPartitioner.java
index 5239522cc6aa..9eef249d6738 100644
---
a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDBucketIndexPartitioner.java
+++
b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDBucketIndexPartitioner.java
@@ -23,6 +23,7 @@ import org.apache.hudi.common.model.HoodieKey;
import org.apache.hudi.common.model.HoodieRecord;
import org.apache.hudi.common.schema.HoodieSchema;
import org.apache.hudi.common.schema.HoodieSchemaUtils;
+import org.apache.hudi.common.util.StringUtils;
import org.apache.hudi.common.util.collection.FlatLists;
import org.apache.hudi.table.BucketIndexBulkInsertPartitioner;
import org.apache.hudi.table.HoodieTable;
@@ -121,7 +122,9 @@ public abstract class RDDBucketIndexPartitioner<T> extends
BucketIndexBulkInsert
LOG.warn("Bucket index does not support global sort mode, the sort will
only be done within each data partition");
}
- Comparator<HoodieKey> comparator = (Comparator<HoodieKey> & Serializable)
(t1, t2) -> t1.getRecordKey().compareTo(t2.getRecordKey());
+ Comparator<HoodieKey> comparator =
table.getMetaClient().getTableConfig().isLSMTreeStorageLayout()
+ ? (Comparator<HoodieKey> & Serializable) (t1, t2) ->
StringUtils.compareUtf8Bytes(t1.getRecordKey(), t2.getRecordKey())
+ : (Comparator<HoodieKey> & Serializable) (t1, t2) ->
t1.getRecordKey().compareTo(t2.getRecordKey());
return records.mapToPair(record -> new Tuple2<>(record.getKey(), record))
.repartitionAndSortWithinPartitions(partitioner, comparator)
diff --git
a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
index 14aff571c238..805dfe8ed566 100644
---
a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
+++
b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
@@ -18,7 +18,7 @@
package org.apache.spark.sql
-import org.apache.hudi.SparkAdapterSupport
+import org.apache.hudi.{HoodieUTF8String, SparkAdapterSupport}
import org.apache.hudi.common.model.HoodieRecord
import org.apache.hudi.common.table.view.FileSystemViewStorageConfig
import org.apache.hudi.common.util.{Functions, RemotePartitionHelper}
@@ -32,6 +32,14 @@ import org.apache.spark.sql.catalyst.InternalRow
object BucketPartitionUtils extends SparkAdapterSupport {
def createDataFrame(df: DataFrame, indexKeyFields: String,
numBucketsFunction: NumBucketsFunction, partitioner: Partitioner): DataFrame = {
+ createDataFrame(df, indexKeyFields, numBucketsFunction, partitioner,
sortByRecordKey = false)
+ }
+
+ def createDataFrame(df: DataFrame,
+ indexKeyFields: String,
+ numBucketsFunction: NumBucketsFunction,
+ partitioner: Partitioner,
+ sortByRecordKey: Boolean): DataFrame = {
// parse the comma-separated config once outside the per-row closure; the
list is a
// serializable java.util.List, safe to capture
val indexKeyFieldList = KeyGenUtils.getIndexKeyFields(indexKeyFields)
@@ -49,10 +57,34 @@ object BucketPartitionUtils extends SparkAdapterSupport {
val getPartitionKey = getPartitionKeyExtractor()
// use internalRow to avoid extra convert.
- val reRdd = df.queryExecution.toRdd
- .keyBy(row => getPartitionKey(row))
- .repartitionAndSortWithinPartitions(partitioner)
- .values
+ val internalRows = df.queryExecution.toRdd
+ val reRdd = if (sortByRecordKey) {
+ val utf8StringFactory = sparkAdapter.getUTF8StringFactory
+ // Use (bucket route, record key) as the shuffle key. Tuple ordering
groups rows by
+ // (partition path, bucket id) first, then orders each bucket by the
full record key.
+ // Scala derives the record-key ordering from HoodieUTF8String's
Comparable implementation,
+ // which uses binary UTF-8 ordering for both Spark 3 and Spark 4.
+ val keyedRows = internalRows.keyBy(row => {
+ val recordKey = utf8StringFactory.wrapUTF8String(
+ row.getUTF8String(HoodieRecord.RECORD_KEY_META_FIELD_ORD))
+ (getPartitionKey(row), recordKey)
+ })
+ // The record key participates only in sorting; bucket routing remains
unchanged.
+ val bucketRoutePartitioner = new Partitioner {
+ override def numPartitions: Int = partitioner.numPartitions
+
+ override def getPartition(key: Any): Int = {
+ val bucketRoute = key.asInstanceOf[((String, Int),
HoodieUTF8String)]._1
+ partitioner.getPartition(bucketRoute)
+ }
+ }
+
keyedRows.repartitionAndSortWithinPartitions(bucketRoutePartitioner).values
+ } else {
+ internalRows
+ .keyBy(row => getPartitionKey(row))
+ .repartitionAndSortWithinPartitions(partitioner)
+ .values
+ }
sparkAdapter.internalCreateDataFrame(df.sparkSession, reRdd, df.schema)
}
diff --git
a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
index e591eb272768..bdc1d2bc7380 100644
---
a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
+++
b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
@@ -141,7 +141,9 @@ public abstract class
BaseDatasetBulkInsertCommitActionExecutor implements Seria
if (populateMetaFields) {
if (writeConfig.getIndexType() == HoodieIndex.IndexType.BUCKET) {
if (writeConfig.getBucketIndexEngineType() ==
HoodieIndex.BucketIndexEngineType.SIMPLE) {
- return new
BucketIndexBulkInsertPartitionerWithRows(writeConfig.getBucketIndexHashFieldWithDefault(),
table.getConfig());
+ return new BucketIndexBulkInsertPartitionerWithRows(
+ writeConfig.getBucketIndexHashFieldWithDefault(),
table.getConfig(),
+ table.getMetaClient().getTableConfig().isLSMTreeStorageLayout());
} else {
return new ConsistentBucketIndexBulkInsertPartitionerWithRows(table,
Collections.emptyMap(), true);
}
diff --git
a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
index 9f9e9ae9288f..34040e1f26c0 100644
---
a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
+++
b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
@@ -57,7 +57,9 @@ public class DatasetBucketRescaleCommitActionExecutor extends
DatasetBulkInsertO
*/
@Override
protected BulkInsertPartitioner<Dataset<Row>> getPartitioner(boolean
populateMetaFields, boolean isTablePartitioned) {
- return new
BucketIndexBulkInsertPartitionerWithRows(writeClient.getConfig(), expression,
rule, bucketNumber);
+ return new BucketIndexBulkInsertPartitionerWithRows(
+ writeClient.getConfig(), expression, rule, bucketNumber,
+ table.getMetaClient().getTableConfig().isLSMTreeStorageLayout());
}
/**
diff --git
a/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/execution/bulkinsert/TestLSMBulkInsertPartitioner.java
b/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/execution/bulkinsert/TestLSMBulkInsertPartitioner.java
index 4c56005b4e31..65661f813c07 100644
---
a/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/execution/bulkinsert/TestLSMBulkInsertPartitioner.java
+++
b/hudi-spark-datasource/hudi-spark/src/test/java/org/apache/hudi/execution/bulkinsert/TestLSMBulkInsertPartitioner.java
@@ -24,6 +24,7 @@ import org.apache.hudi.common.model.HoodieRecord;
import org.apache.hudi.common.table.HoodieTableConfig;
import org.apache.hudi.common.table.HoodieTableMetaClient;
import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.config.HoodieIndexConfig;
import org.apache.hudi.config.HoodieWriteConfig;
import org.apache.hudi.exception.HoodieException;
import org.apache.hudi.table.BulkInsertPartitioner;
@@ -132,6 +133,36 @@ public class TestLSMBulkInsertPartitioner extends
HoodieSparkClientTestHarness {
assertTrue(partitioner.arePartitionRecordsSorted());
}
+ @Test
+ void
testSimpleBucketRowPartitionerSortsByUtf8RecordKeyWithoutChangingSchema() {
+ StructType schema = new StructType()
+ .add(HoodieRecord.COMMIT_TIME_METADATA_FIELD, DataTypes.StringType,
false)
+ .add(HoodieRecord.COMMIT_SEQNO_METADATA_FIELD, DataTypes.StringType,
false)
+ .add(HoodieRecord.RECORD_KEY_METADATA_FIELD, DataTypes.StringType,
false)
+ .add(HoodieRecord.PARTITION_PATH_METADATA_FIELD, DataTypes.StringType,
false)
+ .add(HoodieRecord.FILENAME_METADATA_FIELD, DataTypes.StringType, false)
+ .add("value", DataTypes.IntegerType, false);
+ Dataset<Row> input = sqlContext.createDataFrame(
+ jsc.parallelize(createRowsWithMetaFields(), 3), schema);
+ HoodieWriteConfig config = createWriteConfig(BulkInsertSortMode.NONE,
true);
+ config.setValue(HoodieIndexConfig.BUCKET_INDEX_HASH_FIELD, "id");
+ config.setValue(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS, "1");
+
+ List<BulkInsertPartitioner<Dataset<Row>>> partitioners = Arrays.asList(
+ new BucketIndexBulkInsertPartitionerWithRows("id", config, true),
+ new BucketIndexBulkInsertPartitionerWithRows(config, "p1|p2|p3,1",
"regex", 1, true));
+
+ for (BulkInsertPartitioner<Dataset<Row>> partitioner : partitioners) {
+ Dataset<Row> actual = partitioner.repartitionRecords(input, 1);
+
+ assertEquals(schema, actual.schema(), "Sorting must not add temporary
columns");
+ assertEquals(1, actual.javaRDD().getNumPartitions());
+ assertSortedSparkPartitions(actual.javaRDD().glom().collect(), row ->
new Tuple2<>(
+ row.getAs(HoodieRecord.PARTITION_PATH_METADATA_FIELD),
+ row.getAs(HoodieRecord.RECORD_KEY_METADATA_FIELD)));
+ }
+ }
+
@ParameterizedTest
@EnumSource(value = BulkInsertSortMode.class, names = {
"GLOBAL_SORT", "PARTITION_SORT", "PARTITION_PATH_REPARTITION_AND_SORT"})
@@ -294,6 +325,15 @@ public class TestLSMBulkInsertPartitioner extends
HoodieSparkClientTestHarness {
return rows;
}
+ private List<Row> createRowsWithMetaFields() {
+ List<Row> rows = new ArrayList<>();
+ int value = 0;
+ for (Tuple2<String, String> key : createKeys()) {
+ rows.add(RowFactory.create("001", "001_0", key._2, key._1, "", value++));
+ }
+ return rows;
+ }
+
private List<Tuple2<String, String>> createKeys() {
String bmpPrivateUse = new String(Character.toChars(0xE000));
String supplementary = new String(Character.toChars(0x20000));
diff --git
a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestLSMDataSource.scala
b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestLSMDataSource.scala
index cac27b42c61f..92382c9a165b 100644
---
a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestLSMDataSource.scala
+++
b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestLSMDataSource.scala
@@ -20,14 +20,17 @@ package org.apache.hudi.functional
import org.apache.hudi.{DataSourceUtils, DataSourceWriteOptions}
import org.apache.hudi.client.SparkRDDWriteClient
import org.apache.hudi.common.config.HoodieStorageConfig
-import org.apache.hudi.common.model.{HoodieBaseFile, HoodieRecord,
HoodieRecordPayload, HoodieTableType, WriteOperationType}
+import org.apache.hudi.common.fs.{FileNameParser, FSUtils}
+import org.apache.hudi.common.model.{HoodieBaseFile,
HoodieConsistentHashingMetadata, HoodieRecord, HoodieRecordPayload,
HoodieTableType, WriteOperationType}
import org.apache.hudi.common.table.{HoodieTableConfig, HoodieTableMetaClient}
import org.apache.hudi.common.testutils.HoodieTestUtils
import org.apache.hudi.common.util.Option
import org.apache.hudi.common.util.StringUtils
-import org.apache.hudi.config.{HoodieCompactionConfig, HoodieWriteConfig}
+import org.apache.hudi.config.{HoodieCompactionConfig, HoodieIndexConfig,
HoodieWriteConfig}
import org.apache.hudi.exception.HoodieException
import org.apache.hudi.execution.bulkinsert.{BulkInsertSortMode,
RowCustomColumnsSortPartitioner}
+import org.apache.hudi.index.HoodieIndex
+import org.apache.hudi.index.bucket.{BucketIdentifier,
ConsistentBucketIdentifier}
import org.apache.hudi.testutils.{HoodieClientTestUtils,
SparkClientFunctionalTestHarness}
import
org.apache.hudi.testutils.SparkClientFunctionalTestHarness.getSparkSqlConf
@@ -175,6 +178,53 @@ class TestLSMDataSource extends
SparkClientFunctionalTestHarness {
"Ａ-row-p2" -> "v1"))
}
+ @ParameterizedTest
+ @MethodSource(Array("bucketBulkInsertParams"))
+ def testBucketIndexBulkInsert(
+ tableType: HoodieTableType,
+ bucketEngineType: HoodieIndex.BucketIndexEngineType,
+ enableRowWriter: Boolean): Unit = {
+ val numBuckets = 2
+ val tablePath =
s"${basePath}_${tableType.name.toLowerCase}_${bucketEngineType.name.toLowerCase}_$enableRowWriter"
+ val options = baseOptions(tableType) ++ Map(
+ DataSourceWriteOptions.ENABLE_ROW_WRITER.key -> enableRowWriter.toString,
+ HoodieWriteConfig.BULK_INSERT_SORT_MODE.key ->
BulkInsertSortMode.NONE.name,
+ HoodieIndexConfig.INDEX_TYPE.key -> HoodieIndex.IndexType.BUCKET.name,
+ HoodieIndexConfig.BUCKET_INDEX_ENGINE_TYPE.key -> bucketEngineType.name,
+ HoodieIndexConfig.BUCKET_INDEX_HASH_FIELD.key -> "id",
+ HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key -> numBuckets.toString)
+ val insertValues = Seq(
+ ("😀-bucket-p1", "v1", 1L, FirstPartition),
+ ("Ａ-bucket-p1", "v1", 1L, FirstPartition),
+ ("middle-bucket-p1", "v1", 1L, FirstPartition),
+ ("😀-bucket-p2", "v1", 1L, SecondPartition),
+ ("Ａ-bucket-p2", "v1", 1L, SecondPartition)) ++
+ routingRecords(bucketEngineType, numBuckets)
+ val inserts = rows(insertValues)
+
+ write(inserts, tablePath, options, WriteOperationType.BULK_INSERT,
SaveMode.Overwrite)
+
+ assertLatestBaseFilesSorted(tablePath, WriteOperationType.BULK_INSERT)
+ assertBucketFileGroupRouting(tablePath, bucketEngineType, numBuckets)
+ val initialFileIds = latestBaseFiles(tablePath).map(_.getFileId).toSet
+ assertEquals(numBuckets * 2, initialFileIds.size)
+ assertSnapshot(tablePath, insertValues.map(value => value._1 ->
value._2).toMap)
+
+ val upsertValues = Seq(
+ ("😀-bucket-p1", "v2", 2L, FirstPartition),
+ ("Ａ-bucket-p1", "v2", 2L, FirstPartition),
+ ("middle-bucket-p1", "v2", 2L, FirstPartition),
+ ("ascii-new-p1", "new", 2L, FirstPartition))
+ val upserts = rows(upsertValues)
+ write(upserts, tablePath, options, WriteOperationType.UPSERT)
+
+ assertChangedFilesSorted(tablePath, tableType, WriteOperationType.UPSERT,
"log")
+ assertEquals(initialFileIds,
latestBaseFiles(tablePath).map(_.getFileId).toSet)
+ assertSnapshot(tablePath,
+ insertValues.map(value => value._1 -> value._2).toMap ++
+ upsertValues.map(value => value._1 -> value._2).toMap)
+ }
+
@ParameterizedTest
@EnumSource(value = classOf[BulkInsertSortMode], names = Array(
"NONE", "PARTITION_PATH_REPARTITION"))
@@ -412,6 +462,63 @@ class TestLSMDataSource extends
SparkClientFunctionalTestHarness {
}
}
+ private def assertBucketFileGroupRouting(
+ tablePath: String,
+ bucketEngineType: HoodieIndex.BucketIndexEngineType,
+ numBuckets: Int): Unit = {
+ val actualFileIdsByPartition = spark.read.format("hudi").load(tablePath)
+ .select("id", HoodieRecord.PARTITION_PATH_METADATA_FIELD,
HoodieRecord.FILENAME_METADATA_FIELD)
+ .collect()
+ .map { row =>
+ val recordKey = row.getString(0)
+ val partitionPath = row.getString(1)
+ val fileId = FSUtils.getFileId(row.getString(2))
+ val actualFileGroup = bucketEngineType match {
+ case HoodieIndex.BucketIndexEngineType.SIMPLE =>
+
BucketIdentifier.bucketIdStr(BucketIdentifier.bucketIdFromFileId(fileId))
+ case HoodieIndex.BucketIndexEngineType.CONSISTENT_HASHING =>
+ FileNameParser.getFileIdPfxFromFileId(fileId)
+ }
+ assertEquals(expectedFileGroup(recordKey, partitionPath,
bucketEngineType, numBuckets), actualFileGroup,
+ s"Unexpected file group for record $recordKey in partition
$partitionPath")
+ partitionPath -> actualFileGroup
+ }
+ .groupBy(_._1)
+
+ Seq(FirstPartition, SecondPartition).foreach { partitionPath =>
+ assertEquals(numBuckets,
actualFileIdsByPartition(partitionPath).map(_._2).distinct.length,
+ s"Expected every bucket to be exercised in partition $partitionPath")
+ }
+ }
+
+ private def routingRecords(
+ bucketEngineType: HoodieIndex.BucketIndexEngineType,
+ numBuckets: Int): Seq[(String, String, Long, String)] = {
+ Seq(FirstPartition, SecondPartition).flatMap { partitionPath =>
+ val keysByFileGroup = (0 until 100)
+ .map(index => s"route-$index-$partitionPath")
+ .groupBy(recordKey => expectedFileGroup(recordKey, partitionPath,
bucketEngineType, numBuckets))
+ assertEquals(numBuckets, keysByFileGroup.size,
+ s"Could not generate a routing key for every bucket in partition
$partitionPath")
+ keysByFileGroup.values.map(_.head).toSeq
+ .map(recordKey => (recordKey, "v1", 1L, partitionPath))
+ }
+ }
+
+ private def expectedFileGroup(
+ recordKey: String,
+ partitionPath: String,
+ bucketEngineType: HoodieIndex.BucketIndexEngineType,
+ numBuckets: Int): String = {
+ bucketEngineType match {
+ case HoodieIndex.BucketIndexEngineType.SIMPLE =>
+ BucketIdentifier.bucketIdStr(BucketIdentifier.getBucketId(recordKey,
"id", numBuckets))
+ case HoodieIndex.BucketIndexEngineType.CONSISTENT_HASHING =>
+ new ConsistentBucketIdentifier(new
HoodieConsistentHashingMetadata(partitionPath, numBuckets))
+ .getBucket(recordKey, "id").getFileIdPrefix
+ }
+ }
+
private def assertSnapshot(tablePath: String, expected: Map[String,
String]): Unit = {
val actual = spark.read.format("hudi").load(tablePath)
.select("id", "value")
@@ -430,6 +537,15 @@ class TestLSMDataSource extends
SparkClientFunctionalTestHarness {
object TestLSMDataSource {
+ def bucketBulkInsertParams(): java.util.stream.Stream[Arguments] =
+ java.util.stream.Stream.of(
+ Arguments.of(HoodieTableType.COPY_ON_WRITE,
HoodieIndex.BucketIndexEngineType.SIMPLE, Boolean.box(false)),
+ Arguments.of(HoodieTableType.COPY_ON_WRITE,
HoodieIndex.BucketIndexEngineType.SIMPLE, Boolean.box(true)),
+ Arguments.of(HoodieTableType.MERGE_ON_READ,
HoodieIndex.BucketIndexEngineType.SIMPLE, Boolean.box(false)),
+ Arguments.of(HoodieTableType.MERGE_ON_READ,
HoodieIndex.BucketIndexEngineType.SIMPLE, Boolean.box(true)),
+ Arguments.of(HoodieTableType.MERGE_ON_READ,
HoodieIndex.BucketIndexEngineType.CONSISTENT_HASHING, Boolean.box(false)),
+ Arguments.of(HoodieTableType.MERGE_ON_READ,
HoodieIndex.BucketIndexEngineType.CONSISTENT_HASHING, Boolean.box(true)))
+
def bulkInsertWithHoodieRecordPathParams():
java.util.stream.Stream[Arguments] =
java.util.stream.Stream.of(
Arguments.of(HoodieTableType.COPY_ON_WRITE,
BulkInsertSortMode.GLOBAL_SORT),
