Skip to content

Commit 233092a

Browse files
platinumhamburgplatinumhamburg
authored
[kv] Bound prefix lookups with an exclusive upper limit (#4744)
* [kv] Bound prefix lookups with an exclusive upper limit * [kv] Remove unused byte prefix comparison
1 parent 26717b0 commit 233092a

3 files changed

Lines changed: 139 additions & 110 deletions

File tree

‎fluss-common/src/main/java/org/apache/fluss/utils/BytesUtils.java‎

Lines changed: 0 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,8 @@
1818
package org.apache.fluss.utils;
1919

2020
import org.apache.fluss.annotation.Internal;
21-
import org.apache.fluss.utils.crc.Java;
2221

23-
import java.lang.reflect.Method;
2422
import java.nio.ByteBuffer;
25-
import java.util.Arrays;
2623

2724
/** Utils for bytes. */
2825
@Internal
@@ -62,100 +59,4 @@ public static byte[] toArray(ByteBuffer buffer, int offset, int size) {
6259
}
6360
return dest;
6461
}
65-
66-
/**
67-
* Check if the given first byte array ({@code prefix}) is a prefix of the second byte array
68-
* ({@code bytes}).
69-
*
70-
* @param prefix the prefix byte array
71-
* @param bytes the byte array to check if it has the prefix
72-
* @return true if the given bytes has the given prefix, false otherwise.
73-
*/
74-
public static boolean prefixEquals(byte[] prefix, byte[] bytes) {
75-
return BEST_EQUAL_COMPARER.prefixEquals(prefix, bytes);
76-
}
77-
78-
// -------------------------------------------------------------------------------------------
79-
80-
private static final BytesPrefixComparer BEST_EQUAL_COMPARER;
81-
82-
static {
83-
if (Java.IS_JAVA9_COMPATIBLE) {
84-
BEST_EQUAL_COMPARER = new Java9BytesPrefixComparer();
85-
} else {
86-
BEST_EQUAL_COMPARER = new PureJavaBytesPrefixComparer();
87-
}
88-
}
89-
90-
/** Compare two byte arrays for equality. */
91-
private interface BytesPrefixComparer {
92-
93-
/**
94-
* Check if the given first byte array ({@code prefix}) is a prefix of the second byte array
95-
* ({@code bytes}).
96-
*
97-
* @param prefix The prefix byte array
98-
* @param bytes The byte array to check if it has the prefix
99-
* @return true if the given bytes has the given prefix, false otherwise
100-
*/
101-
boolean prefixEquals(byte[] prefix, byte[] bytes);
102-
}
103-
104-
private static final class Java9BytesPrefixComparer implements BytesPrefixComparer {
105-
private static final Method EQUALS_METHOD;
106-
107-
static {
108-
try {
109-
EQUALS_METHOD =
110-
Class.forName(Arrays.class.getName())
111-
.getMethod(
112-
"equals",
113-
byte[].class,
114-
int.class,
115-
int.class,
116-
byte[].class,
117-
int.class,
118-
int.class);
119-
} catch (Exception e) {
120-
throw new RuntimeException("Failed to load Arrays.equals method", e);
121-
}
122-
}
123-
124-
@Override
125-
public boolean prefixEquals(byte[] prefix, byte[] bytes) {
126-
if (prefix.length > bytes.length) {
127-
return false;
128-
}
129-
try {
130-
int fromIndex = 0; // inclusive
131-
int toIndex = prefix.length; // exclusive
132-
return (boolean)
133-
EQUALS_METHOD.invoke(
134-
null, prefix, fromIndex, toIndex, bytes, fromIndex, toIndex);
135-
} catch (Throwable e) {
136-
// should never happen
137-
throw new RuntimeException(e);
138-
}
139-
}
140-
}
141-
142-
/**
143-
* A pure Java implementation of the {@link BytesPrefixComparer} that does not rely on Java 9
144-
* APIs and compares bytes one by one which is slower.
145-
*/
146-
private static final class PureJavaBytesPrefixComparer implements BytesPrefixComparer {
147-
148-
@Override
149-
public boolean prefixEquals(byte[] prefix, byte[] bytes) {
150-
if (prefix.length > bytes.length) {
151-
return false;
152-
}
153-
for (int i = 0; i < prefix.length; i++) {
154-
if (prefix[i] != bytes[i]) {
155-
return false;
156-
}
157-
}
158-
return true;
159-
}
160-
}
16162
}

‎fluss-server/src/main/java/org/apache/fluss/server/kv/rocksdb/RocksDBKv.java‎

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import org.apache.fluss.rocksdb.RocksDBOperationUtils;
2525
import org.apache.fluss.server.kv.KvCloseMode;
2626
import org.apache.fluss.server.utils.ResourceGuard;
27-
import org.apache.fluss.utils.BytesUtils;
2827
import org.apache.fluss.utils.IOUtils;
2928

3029
import io.github.fluss_contrib.rocksdb.Cache;
@@ -35,6 +34,7 @@
3534
import io.github.fluss_contrib.rocksdb.RocksDB;
3635
import io.github.fluss_contrib.rocksdb.RocksDBException;
3736
import io.github.fluss_contrib.rocksdb.RocksIterator;
37+
import io.github.fluss_contrib.rocksdb.Slice;
3838
import io.github.fluss_contrib.rocksdb.Statistics;
3939
import io.github.fluss_contrib.rocksdb.WriteOptions;
4040
import org.slf4j.Logger;
@@ -44,6 +44,7 @@
4444

4545
import java.io.IOException;
4646
import java.util.ArrayList;
47+
import java.util.Arrays;
4748
import java.util.List;
4849
import java.util.Objects;
4950
import java.util.concurrent.TimeUnit;
@@ -211,17 +212,20 @@ public List<byte[]> multiGet(List<byte[]> keys) throws IOException {
211212

212213
public List<byte[]> prefixLookup(byte[] prefixKey) {
213214
List<byte[]> pkList = new ArrayList<>();
214-
ReadOptions readOptions = new ReadOptions();
215-
RocksIterator iterator = db.newIterator(defaultColumnFamilyHandle, readOptions);
216-
try {
217-
iterator.seek(prefixKey);
218-
while (iterator.isValid() && BytesUtils.prefixEquals(prefixKey, iterator.key())) {
219-
pkList.add(iterator.value());
220-
iterator.next();
215+
byte[] upperBound = prefixUpperBound(prefixKey);
216+
// The bound must outlive both the read options and the iterator that reference it.
217+
try (Slice upperBoundSlice = upperBound == null ? null : new Slice(upperBound);
218+
ReadOptions readOptions = new ReadOptions().setTotalOrderSeek(true)) {
219+
if (upperBoundSlice != null) {
220+
readOptions.setIterateUpperBound(upperBoundSlice);
221+
}
222+
try (RocksIterator iterator = db.newIterator(defaultColumnFamilyHandle, readOptions)) {
223+
iterator.seek(prefixKey);
224+
while (iterator.isValid()) {
225+
pkList.add(iterator.value());
226+
iterator.next();
227+
}
221228
}
222-
} finally {
223-
readOptions.close();
224-
iterator.close();
225229
}
226230

227231
return pkList;
@@ -482,4 +486,18 @@ private float l0Pressure(long l0Files) {
482486
long offset = Math.min(l0Files - flussL0SlowdownTrigger, (long) window - 1);
483487
return (float) offset / window;
484488
}
489+
490+
/** Returns the exclusive bytewise bound, or null if the prefix has no finite upper bound. */
491+
@Nullable
492+
private static byte[] prefixUpperBound(byte[] prefix) {
493+
for (int i = prefix.length - 1; i >= 0; i--) {
494+
if ((prefix[i] & 0xff) != 0xff) {
495+
byte[] upperBound = Arrays.copyOf(prefix, i + 1);
496+
upperBound[i]++;
497+
return upperBound;
498+
}
499+
}
500+
// Empty and all-0xff prefixes match every key at or after the seek position.
501+
return null;
502+
}
485503
}

‎fluss-server/src/test/java/org/apache/fluss/server/kv/rocksdb/RocksDBKvTest.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,20 @@
2121
import org.apache.fluss.exception.FlussRuntimeException;
2222
import org.apache.fluss.metrics.SimpleCounter;
2323
import org.apache.fluss.metrics.util.TestHistogram;
24+
import org.apache.fluss.row.BinaryString;
25+
import org.apache.fluss.row.GenericRow;
26+
import org.apache.fluss.row.encode.CompactedKeyEncoder;
2427
import org.apache.fluss.server.kv.KvCloseMode;
28+
import org.apache.fluss.types.DataTypes;
29+
import org.apache.fluss.types.RowType;
2530

2631
import io.github.fluss_contrib.rocksdb.ColumnFamilyOptions;
2732
import io.github.fluss_contrib.rocksdb.FlushOptions;
2833
import io.github.fluss_contrib.rocksdb.RocksDBException;
2934
import org.junit.jupiter.api.Test;
3035
import org.junit.jupiter.api.io.TempDir;
36+
import org.junit.jupiter.params.ParameterizedTest;
37+
import org.junit.jupiter.params.provider.ValueSource;
3138

3239
import java.io.File;
3340
import java.nio.file.Path;
@@ -68,6 +75,109 @@ void testRocksDbKv(@TempDir Path tempDir) throws Exception {
6875
}
6976
}
7077

78+
@ParameterizedTest
79+
@ValueSource(booleans = {false, true})
80+
void testPrefixLookupByteBoundaries(boolean flush, @TempDir Path tempDir) throws Exception {
81+
// Already in unsigned bytewise order, including exact prefixes and their upper bounds.
82+
byte[][] keys = {
83+
{},
84+
{0},
85+
{0, 0},
86+
{0, (byte) 0xff},
87+
{1},
88+
{1, 0},
89+
{1, (byte) 0xfe},
90+
{1, (byte) 0xff},
91+
{1, (byte) 0xff, 0},
92+
{1, (byte) 0xff, (byte) 0xff},
93+
{2},
94+
{0x7f},
95+
{0x7f, (byte) 0xff},
96+
{(byte) 0x80},
97+
{(byte) 0x80, 0},
98+
{(byte) 0xfe},
99+
{(byte) 0xff},
100+
{(byte) 0xff, 0},
101+
{(byte) 0xff, (byte) 0xff},
102+
{(byte) 0xff, (byte) 0xff, (byte) 0xff}
103+
};
104+
byte[][] prefixes = {
105+
{},
106+
{0},
107+
{0, (byte) 0xff},
108+
{1},
109+
{1, (byte) 0xff},
110+
{1, (byte) 0xff, (byte) 0xff},
111+
{1, (byte) 0xff, (byte) 0xff, 0},
112+
{2},
113+
{3},
114+
{0x7f},
115+
{0x7f, (byte) 0xff},
116+
{(byte) 0x80},
117+
{(byte) 0xff},
118+
{(byte) 0xff, (byte) 0xff},
119+
{(byte) 0xff, (byte) 0xff, (byte) 0xff, (byte) 0xff}
120+
};
121+
try (RocksDBKv kv = buildRocksDBKv(tempDir.toFile())) {
122+
for (byte[] key : keys) {
123+
kv.put(key, key);
124+
}
125+
if (flush) {
126+
try (FlushOptions options = new FlushOptions().setWaitForFlush(true)) {
127+
kv.db.flush(options);
128+
}
129+
}
130+
for (byte[] prefix : prefixes) {
131+
byte[][] expected =
132+
Arrays.stream(keys)
133+
.filter(
134+
key ->
135+
key.length >= prefix.length
136+
&& Arrays.equals(
137+
Arrays.copyOf(key, prefix.length),
138+
prefix))
139+
.toArray(byte[][]::new);
140+
assertThat(kv.prefixLookup(prefix))
141+
.as("prefix %s (flushed=%s)", Arrays.toString(prefix), flush)
142+
.containsExactly(expected);
143+
}
144+
}
145+
}
146+
147+
@Test
148+
void testPrefixLookupWithEncodedKeysAndMixedSstAndMemtable(@TempDir Path tempDir)
149+
throws Exception {
150+
RowType keyType = RowType.of(DataTypes.STRING(), DataTypes.INT());
151+
CompactedKeyEncoder keyEncoder = new CompactedKeyEncoder(keyType);
152+
CompactedKeyEncoder prefixEncoder = new CompactedKeyEncoder(keyType, new int[] {0});
153+
byte[] prefix = prefixEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ab"), 0));
154+
byte[] key1 = keyEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ab"), 1));
155+
byte[] key2 = keyEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ab"), 2));
156+
byte[] key3 = keyEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ab"), 3));
157+
byte[] key4 = keyEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ab"), 4));
158+
byte[] other = keyEncoder.encodeKey(GenericRow.of(BinaryString.fromString("ac"), 1));
159+
try (RocksDBKv kv = buildRocksDBKv(tempDir.toFile())) {
160+
kv.put(key1, new byte[] {1});
161+
kv.put(key2, new byte[] {2});
162+
kv.put(key3, new byte[] {3});
163+
kv.put(other, new byte[] {9});
164+
try (FlushOptions options = new FlushOptions().setWaitForFlush(true)) {
165+
kv.db.flush(options);
166+
kv.put(key1, new byte[] {5});
167+
kv.delete(key2);
168+
kv.put(key4, new byte[] {4});
169+
assertThat(kv.prefixLookup(prefix))
170+
.containsExactly(new byte[] {5}, new byte[] {3}, new byte[] {4});
171+
kv.db.flush(options);
172+
assertThat(kv.prefixLookup(prefix))
173+
.containsExactly(new byte[] {5}, new byte[] {3}, new byte[] {4});
174+
kv.db.compactRange();
175+
assertThat(kv.prefixLookup(prefix))
176+
.containsExactly(new byte[] {5}, new byte[] {3}, new byte[] {4});
177+
}
178+
}
179+
}
180+
71181
@Test
72182
void testLiveSstFilesSize(@TempDir Path tempDir) throws Exception {
73183
try (RocksDBKv rocksDBKv = buildRocksDBKv(tempDir.toFile())) {

0 commit comments

Comments
 (0)