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 59003b520f [core] Support streaming LocalKvDb bulk load (#8870)
59003b520f is described below
commit 59003b520f370c4df4d5f71bf2fdd507fc092718
Author: Jingsong Lee <[email protected]>
AuthorDate: Tue Jul 28 08:20:30 2026 +0800
[core] Support streaming LocalKvDb bulk load (#8870)
---
.../apache/paimon/lookup/sort/db/LocalKvDb.java | 240 +++++++++++++++------
.../paimon/lookup/sort/db/LocalKvDbTest.java | 43 ++++
2 files changed, 222 insertions(+), 61 deletions(-)
diff --git
a/paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java
b/paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java
index f1a0f43960..2c73a87a95 100644
---
a/paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java
+++
b/paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java
@@ -37,6 +37,7 @@ import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
@@ -123,6 +124,7 @@ public class LocalKvDb implements Closeable {
private final Map<File, SortLookupStoreReader> readerCache;
private final AtomicLong fileSequence;
+ @Nullable private BulkLoadWriter activeBulkLoadWriter;
private boolean closed;
private LocalKvDb(
@@ -147,6 +149,7 @@ public class LocalKvDb implements Closeable {
this.levels = new LsmLevels(MAX_LEVELS);
this.readerCache = new HashMap<>();
this.fileSequence = new AtomicLong();
+ this.activeBulkLoadWriter = null;
this.closed = false;
LsmCompactor.CompactorFactory compactorFactory =
fileDeleter ->
@@ -214,6 +217,7 @@ public class LocalKvDb implements Closeable {
*/
public void put(byte[] key, byte[] value) throws IOException {
ensureOpen();
+ ensureNoBulkLoad();
checkCompactionFailure();
if (value.length == 0) {
throw new IllegalArgumentException(
@@ -239,6 +243,7 @@ public class LocalKvDb implements Closeable {
*/
public void delete(byte[] key) throws IOException {
ensureOpen();
+ ensureNoBulkLoad();
checkCompactionFailure();
MemorySlice wrappedKey = MemorySlice.wrap(key);
byte[] oldValue = memTable.put(wrappedKey, TOMBSTONE);
@@ -262,115 +267,212 @@ public class LocalKvDb implements Closeable {
*/
public void bulkLoad(Iterator<Map.Entry<byte[], byte[]>> sortedEntries,
long numEntries)
throws IOException {
+ checkArgument(numEntries >= 0, "numEntries must be non-negative.");
+ try (BulkLoadWriter writer = createBulkLoadWriter(numEntries)) {
+ long loadedEntries = 0;
+ while (sortedEntries.hasNext()) {
+ checkArgument(
+ loadedEntries < numEntries,
+ "The iterator contains more entries than numEntries
(%s).",
+ numEntries);
+ Map.Entry<byte[], byte[]> entry = sortedEntries.next();
+ writer.put(entry.getKey(), entry.getValue());
+ loadedEntries++;
+ }
+ checkArgument(
+ loadedEntries == numEntries,
+ "The iterator contains %s entries, but numEntries is %s.",
+ loadedEntries,
+ numEntries);
+ writer.finish();
+ }
+ }
+
+ /**
+ * Create a streaming bulk-load writer. Keys must be written in strictly
increasing order
+ * according to the configured key comparator. The database must be empty.
+ */
+ public BulkLoadWriter createBulkLoadWriter() throws IOException {
+ return createBulkLoadWriter(UNKNOWN_NUM_ENTRIES);
+ }
+
+ private BulkLoadWriter createBulkLoadWriter(long expectedEntries) throws
IOException {
ensureOpen();
checkCompactionFailure();
- checkArgument(numEntries >= 0, "numEntries must be non-negative.");
+ if (activeBulkLoadWriter != null) {
+ throw new IllegalStateException("Another bulk load is already in
progress.");
+ }
if (!memTable.isEmpty() || getSstFileCount() > 0) {
throw new IllegalStateException(
"bulkLoad requires an empty database (no memTable entries
and no SST files)");
}
+ BulkLoadWriter writer = new BulkLoadWriter(expectedEntries);
+ activeBulkLoadWriter = writer;
+ return writer;
+ }
- int targetLevel = MAX_LEVELS - 1;
- List<SstFileMetadata> bulkLoadFiles = new ArrayList<>();
-
- SortLookupStoreWriter currentWriter = null;
- File currentSstFile = null;
- MemorySlice currentFileMinKey = null;
- MemorySlice currentFileMaxKey = null;
- MemorySlice previousFileMaxKey = null;
- long currentBatchSize = 0;
- long loadedEntries = 0;
- long loadedBytes = 0;
-
- try {
- while (sortedEntries.hasNext()) {
+ /** A streaming writer which atomically publishes bulk-loaded SST files on
{@link #finish()}. */
+ public final class BulkLoadWriter implements Closeable {
+
+ private final long expectedEntries;
+ private final int targetLevel;
+ private final List<SstFileMetadata> bulkLoadFiles;
+
+ @Nullable private SortLookupStoreWriter currentWriter;
+ @Nullable private File currentSstFile;
+ @Nullable private MemorySlice currentFileMinKey;
+ @Nullable private MemorySlice currentFileMaxKey;
+ @Nullable private MemorySlice previousFileMaxKey;
+ @Nullable private MemorySlice previousKey;
+ private long currentBatchSize;
+ private long loadedEntries;
+ private long loadedBytes;
+ private boolean active;
+
+ private BulkLoadWriter(long expectedEntries) {
+ this.expectedEntries = expectedEntries;
+ this.targetLevel = MAX_LEVELS - 1;
+ this.bulkLoadFiles = new ArrayList<>();
+ this.active = true;
+ }
+
+ /** Write one key-value pair. */
+ public void put(byte[] key, byte[] value) throws IOException {
+ ensureActive();
+ try {
checkArgument(
- loadedEntries < numEntries,
- "The iterator contains more entries than numEntries
(%s).",
- numEntries);
- Map.Entry<byte[], byte[]> entry = sortedEntries.next();
- byte[] key = entry.getKey();
- byte[] value = entry.getValue();
+ expectedEntries < 0 || loadedEntries < expectedEntries,
+ "The bulk load contains more entries than expected
(%s).",
+ expectedEntries);
+ checkArgument(
+ value.length > 0,
+ "Value must not be empty, which is reserved as the
tombstone marker.");
MemorySlice currentKey = MemorySlice.wrap(key);
+ checkArgument(
+ previousKey == null ||
keyComparator.compare(previousKey, currentKey) < 0,
+ "bulkLoad requires entries sorted in strictly
increasing order according "
+ + "to the configured key comparator; generated
SST key ranges "
+ + "must be ordered.");
long entrySize = (long) key.length + value.length;
if (currentWriter == null) {
currentSstFile = newSstFile();
long expectedEntries =
- estimateBulkLoadSstEntries(
- numEntries - loadedEntries,
- loadedEntries,
- loadedBytes,
- entrySize);
+ this.expectedEntries < 0
+ ? UNKNOWN_NUM_ENTRIES
+ : estimateBulkLoadSstEntries(
+ this.expectedEntries -
loadedEntries,
+ loadedEntries,
+ loadedBytes,
+ entrySize);
currentWriter =
storeFactory.createWriter(
currentSstFile,
bloomFilterBuilderFactory.apply(expectedEntries));
- currentFileMinKey = currentKey;
+ currentFileMinKey = copyKey(key);
currentBatchSize = 0;
}
currentWriter.put(key, value);
- currentFileMaxKey = currentKey;
+ currentFileMaxKey = copyKey(key);
+ previousKey = currentFileMaxKey;
currentBatchSize += entrySize;
loadedEntries++;
loadedBytes += entrySize;
if (currentBatchSize >= maxSstFileSize) {
- currentWriter.close();
- currentWriter = null;
- previousFileMaxKey =
- addBulkLoadSstFile(
- bulkLoadFiles,
- currentSstFile,
- currentFileMinKey,
- currentFileMaxKey,
- previousFileMaxKey,
- targetLevel);
- currentSstFile = null;
- currentFileMinKey = null;
- currentFileMaxKey = null;
+ closeCurrentFile();
}
+ } catch (IOException | RuntimeException e) {
+ abort(e);
+ throw e;
}
+ }
+
+ /** Finish writing and publish all generated SST files. */
+ public void finish() throws IOException {
+ ensureActive();
+ try {
+ if (currentWriter != null) {
+ closeCurrentFile();
+ }
+ checkArgument(
+ expectedEntries < 0 || loadedEntries ==
expectedEntries,
+ "The bulk load contains %s entries, but expected %s.",
+ loadedEntries,
+ expectedEntries);
+ levels.addFiles(targetLevel, bulkLoadFiles);
+ active = false;
+ release();
+ } catch (IOException | RuntimeException e) {
+ abort(e);
+ throw e;
+ }
+
+ LOG.info(
+ "Bulk-loaded {} SST files directly to level {}",
+ bulkLoadFiles.size(),
+ targetLevel);
+ }
+ private void closeCurrentFile() throws IOException {
if (currentWriter != null) {
currentWriter.close();
currentWriter = null;
- addBulkLoadSstFile(
- bulkLoadFiles,
- currentSstFile,
- currentFileMinKey,
- currentFileMaxKey,
- previousFileMaxKey,
- targetLevel);
}
-
- checkArgument(
- loadedEntries == numEntries,
- "The iterator contains %s entries, but numEntries is %s.",
- loadedEntries,
- numEntries);
- } catch (IOException | RuntimeException e) {
+ previousFileMaxKey =
+ addBulkLoadSstFile(
+ bulkLoadFiles,
+ currentSstFile,
+ currentFileMinKey,
+ currentFileMaxKey,
+ previousFileMaxKey,
+ targetLevel);
+ currentSstFile = null;
+ currentFileMinKey = null;
+ currentFileMaxKey = null;
+ }
+
+ private void abort(Throwable cause) {
+ if (!active) {
+ return;
+ }
if (currentWriter != null) {
try {
currentWriter.close();
- } catch (IOException suppressed) {
- e.addSuppressed(suppressed);
+ } catch (IOException e) {
+ cause.addSuppressed(e);
}
+ currentWriter = null;
}
if (currentSstFile != null) {
deleteFileQuietly(currentSstFile);
+ currentSstFile = null;
}
for (SstFileMetadata metadata : bulkLoadFiles) {
deleteFileQuietly(metadata.getFile());
}
- throw e;
+ bulkLoadFiles.clear();
+ active = false;
+ release();
}
- levels.addFiles(targetLevel, bulkLoadFiles);
+ private void release() {
+ if (activeBulkLoadWriter == this) {
+ activeBulkLoadWriter = null;
+ }
+ }
- LOG.info(
- "Bulk-loaded {} SST files directly to level {}",
bulkLoadFiles.size(), targetLevel);
+ private void ensureActive() {
+ if (!active) {
+ throw new IllegalStateException("Bulk-load writer is already
closed.");
+ }
+ }
+
+ @Override
+ public void close() {
+ abort(new IOException("Bulk load was closed before finish."));
+ }
}
private long estimateBulkLoadSstEntries(
@@ -432,6 +534,7 @@ public class LocalKvDb implements Closeable {
@Nullable
public byte[] get(byte[] key) throws IOException {
ensureOpen();
+ ensureNoBulkLoad();
checkCompactionFailure();
// 1. Search MemTable first (newest data)
@@ -459,6 +562,7 @@ public class LocalKvDb implements Closeable {
*/
public void flush() throws IOException {
ensureOpen();
+ ensureNoBulkLoad();
checkCompactionFailure();
if (memTable.isEmpty()) {
return;
@@ -489,6 +593,7 @@ public class LocalKvDb implements Closeable {
*/
public void compact() throws IOException {
ensureOpen();
+ ensureNoBulkLoad();
compaction.fullCompact();
}
@@ -504,6 +609,9 @@ public class LocalKvDb implements Closeable {
closed = true;
IOException failure = null;
+ if (activeBulkLoadWriter != null) {
+ activeBulkLoadWriter.close();
+ }
try {
checkCompactionFailure();
} catch (IOException e) {
@@ -653,6 +761,16 @@ public class LocalKvDb implements Closeable {
}
}
+ private void ensureNoBulkLoad() {
+ if (activeBulkLoadWriter != null) {
+ throw new IllegalStateException("A bulk load is in progress.");
+ }
+ }
+
+ private static MemorySlice copyKey(byte[] key) {
+ return MemorySlice.wrap(Arrays.copyOf(key, key.length));
+ }
+
//
-------------------------------------------------------------------------
// Builder
//
-------------------------------------------------------------------------
diff --git
a/paimon-common/src/test/java/org/apache/paimon/lookup/sort/db/LocalKvDbTest.java
b/paimon-common/src/test/java/org/apache/paimon/lookup/sort/db/LocalKvDbTest.java
index f3be0c8047..66174ccf5a 100644
---
a/paimon-common/src/test/java/org/apache/paimon/lookup/sort/db/LocalKvDbTest.java
+++
b/paimon-common/src/test/java/org/apache/paimon/lookup/sort/db/LocalKvDbTest.java
@@ -1492,6 +1492,49 @@ public class LocalKvDbTest {
}
}
+ @Test
+ public void testStreamingBulkLoad() throws IOException {
+ File directory = new File(tempDir.toFile(), "streaming-bulk-load");
+ try (LocalKvDb db =
+ LocalKvDb.builder(directory)
+ .maxSstFileSize(16)
+ .blockSize(128)
+ .compressOptions(new CompressOptions("none", 1))
+ .build()) {
+ LocalKvDb.BulkLoadWriter writer = db.createBulkLoadWriter();
+ writer.put("key-1".getBytes(UTF_8), "value-1".getBytes(UTF_8));
+ writer.put("key-2".getBytes(UTF_8), "value-2".getBytes(UTF_8));
+ writer.put("key-3".getBytes(UTF_8), "value-3".getBytes(UTF_8));
+ writer.finish();
+
+ Assertions.assertEquals("value-1", getString(db, "key-1"));
+ Assertions.assertEquals("value-2", getString(db, "key-2"));
+ Assertions.assertEquals("value-3", getString(db, "key-3"));
+ Assertions.assertTrue(db.getLevelFileCount(LocalKvDb.MAX_LEVELS -
1) > 1);
+ }
+ }
+
+ @Test
+ public void testStreamingBulkLoadRejectsDuplicateKeysAndCleansFiles()
throws IOException {
+ File directory = new File(tempDir.toFile(),
"streaming-bulk-load-duplicate");
+ try (LocalKvDb db =
+ LocalKvDb.builder(directory)
+ .blockSize(128)
+ .compressOptions(new CompressOptions("none", 1))
+ .build()) {
+ LocalKvDb.BulkLoadWriter writer = db.createBulkLoadWriter();
+ writer.put("key".getBytes(UTF_8), "value-1".getBytes(UTF_8));
+
+ IllegalArgumentException exception =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> writer.put("key".getBytes(UTF_8),
"value-2".getBytes(UTF_8)));
+ Assertions.assertTrue(exception.getMessage().contains("strictly
increasing"));
+ Assertions.assertEquals(0, db.getSstFileCount());
+ assertNoSstFiles(directory);
+ }
+ }
+
@Test
public void testBulkLoadValidatesEntryCount() throws IOException {
List<Map.Entry<byte[], byte[]>> entries = new ArrayList<>();