Skip to content

[core] Support streaming LocalKvDb bulk load #8870

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JingsongLi merged 1 commit into from
Jul 28, 2026
+222 −61
Merged

[core] Support streaming LocalKvDb bulk load #8870

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
240 changes: 179 additions & 61 deletions paimon-common/src/main/java/org/apache/paimon/lookup/sort/db/LocalKvDb.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
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;
Expand Down Expand Up @@ -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(
Expand All @@ -147,6 +149,7 @@ private LocalKvDb(
this.levels = new LsmLevels(MAX_LEVELS);
this.readerCache = new HashMap<>();
this.fileSequence = new AtomicLong();
this.activeBulkLoadWriter = null;
this.closed = false;
LsmCompactor.CompactorFactory compactorFactory =
fileDeleter ->
Expand Down Expand Up @@ -214,6 +217,7 @@ public static Builder builder(File dataDirectory) {
*/
public void put(byte[] key, byte[] value) throws IOException {
ensureOpen();
ensureNoBulkLoad();
checkCompactionFailure();
if (value.length == 0) {
throw new IllegalArgumentException(
Expand All @@ -239,6 +243,7 @@ public void put(byte[] key, byte[] value) throws IOException {
*/
public void delete(byte[] key) throws IOException {
ensureOpen();
ensureNoBulkLoad();
checkCompactionFailure();
MemorySlice wrappedKey = MemorySlice.wrap(key);
byte[] oldValue = memTable.put(wrappedKey, TOMBSTONE);
Expand All @@ -262,115 +267,212 @@ public void delete(byte[] key) throws IOException {
*/
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(
Expand Down Expand Up @@ -432,6 +534,7 @@ private MemorySlice addBulkLoadSstFile(
@Nullable
public byte[] get(byte[] key) throws IOException {
ensureOpen();
ensureNoBulkLoad();
checkCompactionFailure();

// 1. Search MemTable first (newest data)
Expand Down Expand Up @@ -459,6 +562,7 @@ public byte[] get(byte[] key) throws IOException {
*/
public void flush() throws IOException {
ensureOpen();
ensureNoBulkLoad();
checkCompactionFailure();
if (memTable.isEmpty()) {
return;
Expand Down Expand Up @@ -491,6 +595,7 @@ private void flushMemTable() throws IOException {
*/
public void compact() throws IOException {
ensureOpen();
ensureNoBulkLoad();
compaction.fullCompact();
}

Expand All @@ -506,6 +611,9 @@ public void close() throws IOException {
closed = true;

IOException failure = null;
if (activeBulkLoadWriter != null) {
activeBulkLoadWriter.close();
}
try {
checkCompactionFailure();
} catch (IOException e) {
Expand Down Expand Up @@ -643,6 +751,16 @@ private void ensureOpen() {
}
}

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
// -------------------------------------------------------------------------
Expand Down
Loading
Loading