This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 5c0c74d5308 Optimize consuming MAP key access (#19168)
5c0c74d5308 is described below
commit 5c0c74d5308e237b0b5b0659118f62af8a186481
Author: Xiang Fu <[email protected]>
AuthorDate: Wed Aug 12 23:03:37 2026 -0700
Optimize consuming MAP key access (#19168)
Avoid deserializing complete consuming-segment MAP values for selective key
lookups by adding backward-compatible forward-index defaults and an optimized
mutable off-heap implementation. Cache encoded keys, validate malformed frames
before allocation, preserve fallback behavior, and benchmark the actual
forward-index path.
---
.../apache/pinot/perf/BenchmarkMapKeyAccess.java | 132 +++++++++++++++++++++
.../writer/impl/MutableOffHeapByteArrayStore.java | 25 ++++
.../impl/forward/VarByteSVMutableForwardIndex.java | 8 ++
.../local/segment/index/map/MapKeyIndexReader.java | 25 ++--
.../impl/MutableOffHeapByteArrayStoreTest.java | 102 ++++++++++++++++
.../mutable/VarByteSVMutableForwardIndexTest.java | 22 ++++
.../segment/index/map/MapKeyIndexReaderTest.java | 112 +++++++++++++++++
.../spi/index/reader/ForwardIndexReader.java | 20 ++++
.../java/org/apache/pinot/spi/utils/MapUtils.java | 109 +++++++++++++++++
.../org/apache/pinot/spi/utils/MapUtilsTest.java | 91 ++++++++++++++
10 files changed, 633 insertions(+), 13 deletions(-)
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkMapKeyAccess.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkMapKeyAccess.java
new file mode 100644
index 00000000000..2ba7aa04a20
--- /dev/null
+++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkMapKeyAccess.java
@@ -0,0 +1,132 @@
+/**
+ * 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.pinot.perf;
+
+import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager;
+import
org.apache.pinot.segment.local.realtime.impl.forward.VarByteSVMutableForwardIndex;
+import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.MapUtils;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.TearDown;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.options.ChainedOptionsBuilder;
+import org.openjdk.jmh.runner.options.CommandLineOptions;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/// Measures selective MAP key access against the current full-map read path.
+///
+/// Keys are fixed-length and share a common prefix, modelling dotted
OpenTelemetry-style attribute names
+/// (`k8s.workload.name`, `k8s.namespace.name`, ...). That is the honest case
for a scanning extractor: every entry
+/// clears the key-length check, so the key bytes are actually compared rather
than skipped on a length mismatch.
+/// Each benchmark thread owns its forward index and off-heap memory manager.
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Fork(value = 2)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+@State(Scope.Thread)
+public class BenchmarkMapKeyAccess {
+
+ public static void main(String[] args)
+ throws Exception {
+ // Inherit the command line so `-p`, `-prof`, `-f` and friends take
effect. Without the parent options they are
+ // parsed and then silently dropped, and the run quietly falls back to the
annotated defaults.
+ ChainedOptionsBuilder opt = new OptionsBuilder().parent(new
CommandLineOptions(args))
+ .include(BenchmarkMapKeyAccess.class.getSimpleName());
+ new Runner(opt.build()).run();
+ }
+
+ @Param({"4", "16", "64"})
+ private int _numEntries;
+
+ @Param({"first", "last"})
+ private String _targetPosition;
+
+ /// `flat` models scalar attribute maps (the common ingestion shape);
`nested` models maps whose values are
+ /// themselves objects, where the full-map path pays to materialize a
container per value.
+ @Param({"flat", "nested"})
+ private String _valueShape;
+
+ private String _targetKey;
+ private PreparedMapKey _targetMapKey;
+ private PinotDataBufferMemoryManager _memoryManager;
+ private VarByteSVMutableForwardIndex _forwardIndex;
+
+ @Setup(Level.Trial)
+ public void setUp() {
+ boolean nested = "nested".equals(_valueShape);
+ Map<String, Object> map = new LinkedHashMap<>();
+ for (int i = 0; i < _numEntries; i++) {
+ String value = "value-with-enough-bytes-to-exercise-json-parsing-" + i;
+ map.put(key(i), nested ? Map.of("value", value) : value);
+ }
+ _targetKey = key("first".equals(_targetPosition) ? 0 : _numEntries - 1);
+ _targetMapKey = new PreparedMapKey(_targetKey);
+ byte[] serialized = MapUtils.serializeMap(map, false);
+ _memoryManager = new
DirectMemoryManager(BenchmarkMapKeyAccess.class.getSimpleName());
+ _forwardIndex =
+ new VarByteSVMutableForwardIndex(DataType.MAP, _memoryManager,
"mapColumn", 1, serialized.length);
+ _forwardIndex.setBytes(0, serialized);
+ }
+
+ @TearDown(Level.Trial)
+ public void tearDown()
+ throws IOException {
+ try {
+ _forwardIndex.close();
+ } finally {
+ _memoryManager.close();
+ }
+ }
+
+ private static String key(int i) {
+ return String.format("k8s.attribute.%03d.name", i);
+ }
+
+ /// Runs the existing consuming path: copy the complete off-heap value,
deserialize every MAP entry, then select the
+ /// requested key.
+ @Benchmark
+ public Object fullMapValue() {
+ return _forwardIndex.getMap(0, null).get(_targetKey);
+ }
+
+ /// Runs the optimized consuming path, including store lookup and creation
of the read-only direct view.
+ @Benchmark
+ public Object selectiveMapValue() {
+ return _forwardIndex.getMapEntryValue(0, null, _targetMapKey);
+ }
+}
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
index aad235a2212..8a8d7c05009 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStore.java
@@ -21,6 +21,7 @@ package org.apache.pinot.segment.local.io.writer.impl;
import com.google.common.annotations.VisibleForTesting;
import java.io.Closeable;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.util.LinkedList;
import java.util.List;
import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
@@ -142,6 +143,17 @@ public class MutableOffHeapByteArrayStore implements
Closeable {
return value;
}
+ private ByteBuffer getByteBuffer(int index) {
+ int startOffset = _pinotDataBuffer.getInt(index * Integer.BYTES);
+ int endOffset;
+ if (index != 0) {
+ endOffset = _pinotDataBuffer.getInt((index - 1) * Integer.BYTES);
+ } else {
+ endOffset = _size;
+ }
+ return _pinotDataBuffer.toDirectByteBuffer(startOffset, endOffset -
startOffset);
+ }
+
private int getValueSize(int index) {
int startOffset = _pinotDataBuffer.getInt(index * Integer.BYTES);
int endOffset;
@@ -221,6 +233,19 @@ public class MutableOffHeapByteArrayStore implements
Closeable {
throw new RuntimeException("dictionary ID '" + index + "' too low");
}
+ /// Returns a read-only view of the value at the given index without copying
it.
+ /// The returned buffer must not be used after this store is closed.
+ public ByteBuffer getByteBuffer(int index) {
+ List<Buffer> bufList = _buffers;
+ for (int x = bufList.size() - 1; x >= 0; x--) {
+ Buffer buffer = bufList.get(x);
+ if (index >= buffer.getStartIndex()) {
+ return buffer.getByteBuffer(index -
buffer.getStartIndex()).asReadOnlyBuffer();
+ }
+ }
+ throw new RuntimeException("dictionary ID '" + index + "' too low");
+ }
+
public int getValueSize(int index) {
List<Buffer> bufList = _buffers;
for (int x = bufList.size() - 1; x >= 0; x--) {
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/VarByteSVMutableForwardIndex.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/VarByteSVMutableForwardIndex.java
index d6e29bd0496..6672e392e0f 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/VarByteSVMutableForwardIndex.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/forward/VarByteSVMutableForwardIndex.java
@@ -21,6 +21,7 @@ package org.apache.pinot.segment.local.realtime.impl.forward;
import java.io.IOException;
import java.math.BigDecimal;
import java.util.Map;
+import javax.annotation.Nullable;
import
org.apache.pinot.segment.local.io.writer.impl.MutableOffHeapByteArrayStore;
import org.apache.pinot.segment.spi.index.mutable.MutableForwardIndex;
import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
@@ -28,6 +29,7 @@ import
org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager;
import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.utils.BigDecimalUtils;
import org.apache.pinot.spi.utils.MapUtils;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
import org.apache.pinot.spi.utils.Utf8Utils;
import static java.nio.charset.StandardCharsets.UTF_8;
@@ -100,6 +102,12 @@ public class VarByteSVMutableForwardIndex implements
MutableForwardIndex {
return MapUtils.deserializeMap(getBytes(docId));
}
+ @Nullable
+ @Override
+ public Object getMapEntryValue(int docId, ForwardIndexReaderContext context,
PreparedMapKey key) {
+ return
MapUtils.deserializeMapEntryValue(_byteArrayStore.getByteBuffer(docId), key);
+ }
+
@Override
public void setBigDecimal(int docId, BigDecimal value) {
setBytes(docId, BigDecimalUtils.serialize(value));
diff --git
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReader.java
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReader.java
index 3d987b5c103..dca7a8d2aed 100644
---
a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReader.java
+++
b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReader.java
@@ -20,23 +20,23 @@ package org.apache.pinot.segment.local.segment.index.map;
import java.io.IOException;
import java.math.BigDecimal;
-import java.util.Map;
import javax.annotation.Nullable;
import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
import org.apache.pinot.spi.data.FieldSpec;
import org.apache.pinot.spi.utils.BigDecimalUtils;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
public class MapKeyIndexReader implements ForwardIndexReader {
private final ForwardIndexReader _forwardIndexReader;
private final FieldSpec _keyFieldSpec;
- private final String _keyName;
+ private final PreparedMapKey _mapKey;
private final Object _defaultNullValue;
public MapKeyIndexReader(ForwardIndexReader forwardIndexReader, String
keyName, FieldSpec keyFieldSpec) {
_forwardIndexReader = forwardIndexReader;
- _keyName = keyName;
+ _mapKey = new PreparedMapKey(keyName);
_keyFieldSpec = keyFieldSpec;
_defaultNullValue = keyFieldSpec.getDefaultNullValue();
}
@@ -58,42 +58,41 @@ public class MapKeyIndexReader implements
ForwardIndexReader {
@Override
public int getInt(int docId, ForwardIndexReaderContext context) {
- return Integer.parseInt(extractMapValue(docId, context,
_keyName).toString());
+ return Integer.parseInt(extractMapValue(docId, context).toString());
}
@Override
public long getLong(int docId, ForwardIndexReaderContext context) {
- return Long.parseLong(extractMapValue(docId, context,
_keyName).toString());
+ return Long.parseLong(extractMapValue(docId, context).toString());
}
@Override
public float getFloat(int docId, ForwardIndexReaderContext context) {
- return Float.parseFloat(extractMapValue(docId, context,
_keyName).toString());
+ return Float.parseFloat(extractMapValue(docId, context).toString());
}
@Override
public double getDouble(int docId, ForwardIndexReaderContext context) {
- return Double.parseDouble(extractMapValue(docId, context,
_keyName).toString());
+ return Double.parseDouble(extractMapValue(docId, context).toString());
}
@Override
public String getString(int docId, ForwardIndexReaderContext context) {
- return extractMapValue(docId, context, _keyName).toString();
+ return extractMapValue(docId, context).toString();
}
@Override
public byte[] getBytes(int docId, ForwardIndexReaderContext context) {
- return (byte[]) extractMapValue(docId, context, _keyName);
+ return (byte[]) extractMapValue(docId, context);
}
@Override
public BigDecimal getBigDecimal(int docId, ForwardIndexReaderContext
context) {
- return BigDecimalUtils.deserialize((byte[]) extractMapValue(docId,
context, _keyName));
+ return BigDecimalUtils.deserialize((byte[]) extractMapValue(docId,
context));
}
- private Object extractMapValue(int docId, ForwardIndexReaderContext context,
String key) {
- Map map = _forwardIndexReader.getMap(docId, context);
- Object object = map.get(key);
+ private Object extractMapValue(int docId, ForwardIndexReaderContext context)
{
+ Object object = _forwardIndexReader.getMapEntryValue(docId, context,
_mapKey);
if (object == null) {
return _defaultNullValue;
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
index c316992c034..e817598ce3b 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/io/writer/impl/MutableOffHeapByteArrayStoreTest.java
@@ -18,7 +18,15 @@
*/
package org.apache.pinot.segment.local.io.writer.impl;
+import java.nio.ByteBuffer;
+import java.nio.ReadOnlyBufferException;
import java.util.Arrays;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule;
import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager;
import org.testng.Assert;
@@ -61,6 +69,100 @@ public class MutableOffHeapByteArrayStoreTest implements
PinotBuffersAfterClassC
}
}
+ @Test
+ public void byteBufferTest()
+ throws Exception {
+ try (MutableOffHeapByteArrayStore store =
+ new MutableOffHeapByteArrayStore(_memoryManager, "bytesColumn", 1, 1))
{
+ byte[] firstValue = {1};
+ byte[] secondValue = {2, 3, 4};
+ int firstIndex = store.add(firstValue);
+ int secondIndex = store.add(secondValue);
+
+ ByteBuffer firstBuffer = store.getByteBuffer(firstIndex);
+ byte[] firstResult = new byte[firstBuffer.remaining()];
+ firstBuffer.get(firstResult);
+ Assert.assertEquals(firstResult, firstValue);
+
+ ByteBuffer secondBuffer = store.getByteBuffer(secondIndex);
+ Assert.assertTrue(secondBuffer.isReadOnly());
+ byte[] secondResult = new byte[secondBuffer.remaining()];
+ secondBuffer.get(secondResult);
+ Assert.assertEquals(secondResult, secondValue);
+ Assert.assertThrows(ReadOnlyBufferException.class, () ->
store.getByteBuffer(secondIndex).put((byte) 0));
+ }
+ }
+
+ @Test
+ public void getByteBufferDuringConcurrentAppendTest()
+ throws Exception {
+ int numReaders = 4;
+ int numValues = 2_048;
+ ExecutorService executor = Executors.newFixedThreadPool(numReaders + 1);
+ AtomicInteger publishedCount = new AtomicInteger();
+ CountDownLatch start = new CountDownLatch(1);
+ CountDownLatch firstValueRead = new CountDownLatch(numReaders);
+ try (MutableOffHeapByteArrayStore store =
+ new MutableOffHeapByteArrayStore(_memoryManager,
"concurrentBytesColumn", 1, 1)) {
+ Future<?> writer = executor.submit(() -> {
+ await(start);
+ Assert.assertEquals(store.add(valueForIndex(0)), 0);
+ publishedCount.set(1);
+ await(firstValueRead);
+ for (int i = 1; i < numValues; i++) {
+ Assert.assertEquals(store.add(valueForIndex(i)), i);
+ // Publish only after the value and any expanded buffer are visible
to readers.
+ publishedCount.set(i + 1);
+ }
+ });
+
+ Future<?>[] readers = new Future<?>[numReaders];
+ for (int i = 0; i < numReaders; i++) {
+ readers[i] = executor.submit(() -> {
+ await(start);
+ int nextIndex = 0;
+ while (nextIndex < numValues) {
+ int readableCount = publishedCount.get();
+ while (nextIndex < readableCount) {
+ ByteBuffer byteBuffer = store.getByteBuffer(nextIndex);
+ Assert.assertTrue(byteBuffer.isReadOnly());
+ byte[] actual = new byte[byteBuffer.remaining()];
+ byteBuffer.get(actual);
+ Assert.assertEquals(actual, valueForIndex(nextIndex));
+ nextIndex++;
+ if (nextIndex == 1) {
+ firstValueRead.countDown();
+ }
+ }
+ Thread.yield();
+ }
+ });
+ }
+
+ start.countDown();
+ writer.get(30, TimeUnit.SECONDS);
+ for (Future<?> reader : readers) {
+ reader.get(30, TimeUnit.SECONDS);
+ }
+ } finally {
+ executor.shutdownNow();
+ Assert.assertTrue(executor.awaitTermination(30, TimeUnit.SECONDS));
+ }
+ }
+
+ private static byte[] valueForIndex(int index) {
+ return ByteBuffer.allocate(2 *
Integer.BYTES).putInt(index).putInt(~index).array();
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ Assert.assertTrue(latch.await(30, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException(e);
+ }
+ }
+
@Test
public void startSizeTest() {
Assert.assertEquals(MutableOffHeapByteArrayStore.getStartSize(1, ONE_GB),
ONE_GB + 4);
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java
index 08561c19795..72d5feee947 100644
---
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/forward/mutable/VarByteSVMutableForwardIndexTest.java
@@ -19,6 +19,8 @@
package org.apache.pinot.segment.local.segment.index.forward.mutable;
import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.Map;
import java.util.Random;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.pinot.segment.local.PinotBuffersAfterClassCheckRule;
@@ -26,6 +28,8 @@ import
org.apache.pinot.segment.local.io.writer.impl.DirectMemoryManager;
import
org.apache.pinot.segment.local.realtime.impl.forward.VarByteSVMutableForwardIndex;
import org.apache.pinot.segment.spi.memory.PinotDataBufferMemoryManager;
import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.MapUtils;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
import org.testng.Assert;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
@@ -98,4 +102,22 @@ public class VarByteSVMutableForwardIndexTest implements
PinotBuffersAfterClassC
Assert.assertTrue(readerWriter.canAddMore());
}
}
+
+ @Test
+ public void testMapValue()
+ throws IOException {
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put("first", "value");
+ map.put("k8s.workload.name", "pinot-server");
+ map.put("last", 42);
+ try (VarByteSVMutableForwardIndex readerWriter = new
VarByteSVMutableForwardIndex(DataType.MAP, _memoryManager,
+ "MapColumn", 1, 64)) {
+ readerWriter.setBytes(0, MapUtils.serializeMap(map));
+
+ Assert.assertEquals(readerWriter.getMapEntryValue(0, null, new
PreparedMapKey("k8s.workload.name")),
+ "pinot-server");
+ Assert.assertNull(readerWriter.getMapEntryValue(0, null, new
PreparedMapKey("missing")));
+ Assert.assertEquals(readerWriter.getMap(0, null), map);
+ }
+ }
}
diff --git
a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReaderTest.java
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReaderTest.java
new file mode 100644
index 00000000000..f83f8c3457c
--- /dev/null
+++
b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/map/MapKeyIndexReaderTest.java
@@ -0,0 +1,112 @@
+/**
+ * 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.pinot.segment.local.segment.index.map;
+
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader;
+import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/// Covers [MapKeyIndexReader] over both shapes of underlying reader: one that
implements the selective
+/// [ForwardIndexReader#getMapEntryValue] override (as the mutable forward
index does), and one that only implements
+/// [ForwardIndexReader#getMap] and therefore falls through to the default.
Both must agree.
+public class MapKeyIndexReaderTest {
+ private static final Map<String, Object> MAP =
+ Map.of("k8s.workload.name", "pinot-server", "k8s.workload.replicas", 3);
+
+ @Test
+ public void testSelectiveReader() {
+ assertReaderBehavior(new SelectiveReader());
+ }
+
+ /// The immutable sparse-key path inherits the default `getMapEntryValue`.
It has to keep working unchanged.
+ @Test
+ public void testReaderWithoutSelectiveOverride() {
+ assertReaderBehavior(new FullMapOnlyReader());
+ }
+
+ @Test
+ public void testStringOverloadDelegatesToSelectiveLookup() {
+ assertEquals(new SelectiveReader().getMapEntryValue(0, null,
"k8s.workload.name"), "pinot-server");
+ }
+
+ private static void
assertReaderBehavior(ForwardIndexReader<ForwardIndexReaderContext> reader) {
+ FieldSpec stringSpec = new DimensionFieldSpec("value", DataType.STRING,
true);
+ assertEquals(new MapKeyIndexReader(reader, "k8s.workload.name",
stringSpec).getString(0, null), "pinot-server");
+
+ // A key that is absent from the map resolves to the field spec's default
null value, not to null.
+ assertEquals(new MapKeyIndexReader(reader, "missing",
stringSpec).getString(0, null),
+ stringSpec.getDefaultNullValue());
+
+ FieldSpec intSpec = new DimensionFieldSpec("value", DataType.INT, true);
+ assertEquals(new MapKeyIndexReader(reader, "k8s.workload.replicas",
intSpec).getInt(0, null), 3);
+ }
+
+ /// Mirrors the mutable forward index: answers a single key without
materializing the map.
+ private static class SelectiveReader extends BaseReader {
+ @Nullable
+ @Override
+ public Object getMapEntryValue(int docId, ForwardIndexReaderContext
context, PreparedMapKey key) {
+ return MAP.get(key.getKey());
+ }
+
+ @Override
+ public Map<String, Object> getMap(int docId, ForwardIndexReaderContext
context) {
+ throw new AssertionError("Selective lookup must not materialize the full
map");
+ }
+ }
+
+ /// Mirrors a reader that only knows how to hand back the whole map.
+ private static class FullMapOnlyReader extends BaseReader {
+ }
+
+ private abstract static class BaseReader implements
ForwardIndexReader<ForwardIndexReaderContext> {
+ @Override
+ public Map<String, Object> getMap(int docId, ForwardIndexReaderContext
context) {
+ return MAP;
+ }
+
+ @Override
+ public boolean isDictionaryEncoded() {
+ return false;
+ }
+
+ @Override
+ public boolean isSingleValue() {
+ return true;
+ }
+
+ @Override
+ public DataType getStoredType() {
+ return DataType.MAP;
+ }
+
+ @Override
+ public void close() {
+ }
+ }
+}
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/reader/ForwardIndexReader.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/reader/ForwardIndexReader.java
index 2beba6d0b1f..66749669d94 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/reader/ForwardIndexReader.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/reader/ForwardIndexReader.java
@@ -30,6 +30,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.spi.utils.BigDecimalUtils;
import org.apache.pinot.spi.utils.BytesUtils;
import org.apache.pinot.spi.utils.MapUtils;
+import org.apache.pinot.spi.utils.MapUtils.PreparedMapKey;
import org.apache.pinot.spi.utils.hash.MurmurHashFunctions;
@@ -466,6 +467,25 @@ public interface ForwardIndexReader<T extends
ForwardIndexReaderContext> extends
+ "ForwardIndexReader is being created to read this column.");
}
+ /// Reads a value for a key from a MAP type single-value column at the given
document id.
+ /// Implementations can override this method to avoid deserializing the
entire map when only one key is needed.
+ ///
+ /// @param docId Document id
+ /// @param context Reader context
+ /// @param key Map key
+ /// @return Value for the key, or `null` if the key is missing or its value
is null
+ @Nullable
+ default Object getMapEntryValue(int docId, T context, String key) {
+ return getMapEntryValue(docId, context, new PreparedMapKey(key));
+ }
+
+ /// Variant of [#getMapEntryValue(int, ForwardIndexReaderContext, String)]
that reuses a pre-encoded MAP key.
+ /// Implementations can override this method to avoid repeated key encoding
as well as full-map deserialization.
+ @Nullable
+ default Object getMapEntryValue(int docId, T context, PreparedMapKey key) {
+ return getMap(docId, context).get(key.getKey());
+ }
+
default int get32BitsMurmur3Hash(int docId, T context) {
return MurmurHashFunctions.murmurHash3X64Bit32(getBytes(docId, context),
0);
}
diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/MapUtils.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/MapUtils.java
index b7bf71bd08c..a6e5011e384 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/MapUtils.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/MapUtils.java
@@ -25,12 +25,15 @@ import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import java.io.IOException;
import java.io.OutputStream;
+import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
import java.util.Arrays;
import java.util.Collection;
import java.util.Map;
import java.util.Map.Entry;
import java.util.SortedMap;
+import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -48,6 +51,21 @@ public class MapUtils {
private MapUtils() {
}
+ /// Immutable MAP key with its UTF-8 representation cached for selective
lookup hot paths.
+ public static final class PreparedMapKey {
+ private final String _key;
+ private final byte[] _utf8Bytes;
+
+ public PreparedMapKey(String key) {
+ _key = key;
+ _utf8Bytes = Utf8Utils.encode(key);
+ }
+
+ public String getKey() {
+ return _key;
+ }
+ }
+
private static final Logger LOGGER = LoggerFactory.getLogger(MapUtils.class);
// Pinot's standard ObjectMapper config (JSR-310 / ISO-8601 for LocalDate /
LocalTime — see JsonUtils),
@@ -183,6 +201,91 @@ public class MapUtils {
return map;
}
+ /// Deserializes only the value for the requested key from a length-prefixed
MAP frame.
+ /// Non-matching keys and values are skipped without allocating byte arrays
or invoking Jackson.
+ ///
+ /// @param bytes Serialized MAP frame
+ /// @param key Key whose value should be deserialized
+ /// @return Deserialized value, or `null` if the key is missing, has a null
value, or its JSON value cannot be
+ /// deserialized
+ /// @throws BufferUnderflowException if the MAP frame is malformed or
truncated
+ @Nullable
+ public static Object deserializeMapEntryValue(byte[] bytes, String key) {
+ return deserializeMapEntryValue(ByteBuffer.wrap(bytes), new
PreparedMapKey(key));
+ }
+
+ /// Variant of [#deserializeMapEntryValue(byte[], String)] that reuses a
pre-encoded MAP key.
+ @Nullable
+ public static Object deserializeMapEntryValue(byte[] bytes, PreparedMapKey
key) {
+ return deserializeMapEntryValue(ByteBuffer.wrap(bytes), key);
+ }
+
+ /// Variant of [#deserializeMapEntryValue(byte[], String)] that reads from
the supplied buffer without copying the
+ /// complete MAP frame.
+ ///
+ /// Consumes the buffer from its current position and forces
[ByteOrder#BIG_ENDIAN] on it — the write path frames
+ /// lengths through a big-endian [ByteBuffer], while an off-heap view
inherits the platform's native order.
+ ///
+ /// @throws BufferUnderflowException if the MAP frame is malformed or
truncated
+ @Nullable
+ public static Object deserializeMapEntryValue(ByteBuffer byteBuffer, String
key) {
+ return deserializeMapEntryValue(byteBuffer, new PreparedMapKey(key));
+ }
+
+ /// Variant of [#deserializeMapEntryValue(ByteBuffer, String)] that reuses a
pre-encoded MAP key.
+ ///
+ /// @throws BufferUnderflowException if the MAP frame is malformed or
truncated
+ @Nullable
+ public static Object deserializeMapEntryValue(ByteBuffer byteBuffer,
PreparedMapKey key) {
+ byteBuffer.order(ByteOrder.BIG_ENDIAN);
+ int size = byteBuffer.getInt();
+ if (size < 0) {
+ throw new BufferUnderflowException();
+ }
+ if (size == 0) {
+ return null;
+ }
+ byte[] keyBytes = key._utf8Bytes;
+ int keyBytesLength = keyBytes.length;
+ for (int i = 0; i < size; i++) {
+ int keyLength = byteBuffer.getInt();
+ // Bounds-check up front so the absolute gets below are provably in
range, and so a truncated frame still
+ // surfaces as BufferUnderflowException rather than
IndexOutOfBoundsException.
+ checkLength(byteBuffer, keyLength);
+ // Compare through absolute gets so a length mismatch or a differing
byte skips the rest of the key outright,
+ // rather than walking it one relative get at a time just to advance the
position.
+ boolean matches = keyLength == keyBytesLength;
+ if (matches) {
+ int keyOffset = byteBuffer.position();
+ for (int j = 0; j < keyLength; j++) {
+ if (byteBuffer.get(keyOffset + j) != keyBytes[j]) {
+ matches = false;
+ break;
+ }
+ }
+ }
+ byteBuffer.position(byteBuffer.position() + keyLength);
+
+ int valueLength = byteBuffer.getInt();
+ checkLength(byteBuffer, valueLength);
+ if (!matches) {
+ byteBuffer.position(byteBuffer.position() + valueLength);
+ continue;
+ }
+ // Keys within a frame are unique - the write path iterates a Map - so
the first match is the only match and
+ // the remaining entries never need to be scanned.
+ byte[] valueBytes = new byte[valueLength];
+ byteBuffer.get(valueBytes);
+ try {
+ return JsonUtils.bytesToObject(valueBytes, Object.class);
+ } catch (IOException e) {
+ LOGGER.error("Caught exception while deserializing value for key: {}",
key.getKey(), e);
+ return null;
+ }
+ }
+ return null;
+ }
+
private static byte[] readLengthPrefixed(ByteBuffer byteBuffer) {
int length = byteBuffer.getInt();
byte[] bytes = new byte[length];
@@ -190,6 +293,12 @@ public class MapUtils {
return bytes;
}
+ private static void checkLength(ByteBuffer byteBuffer, int length) {
+ if (length < 0 || length > byteBuffer.remaining()) {
+ throw new BufferUnderflowException();
+ }
+ }
+
public static String toString(Map<String, Object> map) {
return toString(map, true);
}
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/MapUtilsTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/MapUtilsTest.java
index a2ee66ef2a3..4b90ba8ee97 100644
--- a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/MapUtilsTest.java
+++ b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/MapUtilsTest.java
@@ -19,6 +19,8 @@
package org.apache.pinot.spi.utils;
import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
import java.time.LocalDate;
import java.time.LocalTime;
import java.util.HashMap;
@@ -66,6 +68,95 @@ public class MapUtilsTest {
assertNull(deserialized.get("nullValue"), "Null value should be
preserved");
}
+ @Test
+ void testDeserializeMapEntryValue() {
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put("first", Map.of("nested", List.of(1, 2, 3)));
+ map.put("k8s.workload.name", "pinot-server");
+ map.put("nullValue", null);
+ byte[] serialized = MapUtils.serializeMap(map);
+
+ assertEquals(MapUtils.deserializeMapEntryValue(serialized,
"k8s.workload.name"), "pinot-server");
+
assertEquals(MapUtils.deserializeMapEntryValue(ByteBuffer.wrap(serialized),
"first"),
+ Map.of("nested", List.of(1, 2, 3)));
+ assertNull(MapUtils.deserializeMapEntryValue(serialized, "missing"));
+ assertNull(MapUtils.deserializeMapEntryValue(serialized, "nullValue"));
+
assertNull(MapUtils.deserializeMapEntryValue(MapUtils.serializeMap(Map.of()),
"any"));
+ }
+
+ /// Keys that are the same length and differ only in their trailing bytes
are the case the scanning extractor can
+ /// get wrong: the length check passes for every entry, so the match has to
come from the byte comparison alone.
+ @Test
+ void testDeserializeMapEntryValueWithCollidingKeyShapes() {
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put("k8s.workload.name", "workload");
+ map.put("k8s.workload.kind", "kind");
+ map.put("k8s.namespace.nam", "namespace");
+ map.put("k8s.workload", "prefix-of-another-key");
+ map.put("k8s.workload.name.suffixed", "longer-than-another-key");
+ byte[] serialized = MapUtils.serializeMap(map, false);
+
+ for (Map.Entry<String, Object> entry : map.entrySet()) {
+ assertEquals(MapUtils.deserializeMapEntryValue(serialized,
entry.getKey()), entry.getValue(),
+ "Value should match for key: " + entry.getKey());
+ }
+ assertNull(MapUtils.deserializeMapEntryValue(serialized,
"k8s.workload.nam"));
+ assertNull(MapUtils.deserializeMapEntryValue(serialized,
"k8s.workload.names"));
+ }
+
+ /// The extractor matches on encoded UTF-8 bytes rather than decoding each
key, so multi-byte keys - and keys whose
+ /// character count differs from their byte count - have to resolve
correctly.
+ @Test
+ void testDeserializeMapEntryValueWithNonAsciiKeys() {
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put("hôte", "host");
+ map.put("hote", "ascii-host");
+ map.put("命名空间", "namespace");
+ byte[] serialized = MapUtils.serializeMap(map, false);
+
+ for (Map.Entry<String, Object> entry : map.entrySet()) {
+ assertEquals(MapUtils.deserializeMapEntryValue(serialized,
entry.getKey()), entry.getValue(),
+ "Value should match for key: " + entry.getKey());
+ }
+ assertNull(MapUtils.deserializeMapEntryValue(serialized, "命名"));
+ }
+
+ /// An off-heap forward-index view inherits the platform's native byte
order, while the frame is always written
+ /// big-endian. The extractor has to force the order rather than trust the
incoming buffer.
+ @Test
+ void testDeserializeMapEntryValueForcesBigEndian() {
+ byte[] serialized = MapUtils.serializeMap(Map.of("k8s.workload.name",
"pinot-server"));
+ ByteBuffer littleEndian =
ByteBuffer.wrap(serialized).order(ByteOrder.LITTLE_ENDIAN);
+
+ assertEquals(MapUtils.deserializeMapEntryValue(littleEndian,
"k8s.workload.name"), "pinot-server");
+ }
+
+ @Test
+ void testDeserializeMapEntryValueRejectsInvalidValueLength() {
+ String key = "key";
+ byte[] serialized = MapUtils.serializeMap(Map.of(key, "value"));
+ ByteBuffer frame = ByteBuffer.wrap(serialized);
+ frame.position(Integer.BYTES);
+ int keyLength = frame.getInt();
+ int valueLengthOffset = frame.position() + keyLength;
+
+ byte[] negativeLength = serialized.clone();
+ ByteBuffer.wrap(negativeLength).putInt(valueLengthOffset, -1);
+ assertThrows(BufferUnderflowException.class, () ->
MapUtils.deserializeMapEntryValue(negativeLength, key));
+ assertThrows(BufferUnderflowException.class, () ->
MapUtils.deserializeMapEntryValue(negativeLength, "missing"));
+
+ byte[] truncatedValue = serialized.clone();
+ ByteBuffer.wrap(truncatedValue).putInt(valueLengthOffset,
serialized.length);
+ assertThrows(BufferUnderflowException.class, () ->
MapUtils.deserializeMapEntryValue(truncatedValue, key));
+ }
+
+ @Test
+ void testDeserializeMapEntryValueRejectsNegativeMapSize() {
+ byte[] negativeSize =
ByteBuffer.allocate(Integer.BYTES).putInt(-1).array();
+
+ assertThrows(BufferUnderflowException.class, () ->
MapUtils.deserializeMapEntryValue(negativeSize, "key"));
+ }
+
@Test
void testSerializeAndDeserializeWithSpecialCharacters() {
Map<String, Object> map = new HashMap<>();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
