This is an automated email from the ASF dual-hosted git repository.
gortiz 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 4480f3b523d Build v1 DataTable sections in heap pages to avoid
humongous response allocations (#19906)
4480f3b523d is described below
commit 4480f3b523d2e2f90183b3208ea4c4df31412928
Author: Gonzalo Ortiz Jaureguizar <[email protected]>
AuthorDate: Mon Oct 19 16:30:53 2026 +0200
Build v1 DataTable sections in heap pages to avoid humongous response
allocations (#19906)
---
.../apache/pinot/common/datatable/DataTable.java | 5 +
.../pinot/common/datatable/DataTableImplV4.java | 320 +++++++++++----------
.../common/datatable/BaseDataTableBuilder.java | 57 ++--
.../core/common/datatable/DataTableBuilderV4.java | 14 +-
.../core/common/datatable/DataTableSerDeTest.java | 90 ++++++
.../segment/spi/memory/CompoundDataBuffer.java | 27 +-
.../segment/spi/memory/CompoundDataBufferTest.java | 26 ++
7 files changed, 342 insertions(+), 197 deletions(-)
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
index 989b148405f..d8486e5a190 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTable.java
@@ -84,6 +84,11 @@ public interface DataTable {
@Nullable
Map<String, Object> getMap(int rowId, int colId);
+ /// Returns the custom object at the given cell, or `null` for a null value.
+ ///
+ /// The buffer of the returned object can be a read-only view, for example
on a table built in this process. Code
+ /// that reads it, such as
`AggregationFunction.deserializeIntermediateResult` implementations and
+ /// `ObjectSerDeUtils` deserializers, must read it through the `ByteBuffer`
API and must not call `array()` on it.
@Nullable
CustomObject getCustomObject(int rowId, int colId);
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
index d2ab8c43b23..8b0b58cd4ba 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/datatable/DataTableImplV4.java
@@ -34,6 +34,8 @@ import org.apache.pinot.common.datablock.DataBlockUtils;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.common.utils.RoaringBitmapUtils;
+import org.apache.pinot.segment.spi.memory.DataBuffer;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
import org.apache.pinot.spi.accounting.ThreadResourceSnapshot;
import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
import org.apache.pinot.spi.query.QueryThreadContext;
@@ -88,10 +90,12 @@ public class DataTableImplV4 implements DataTable {
protected int _fixDataSize;
protected DataSchema _dataSchema;
protected String[] _stringDictionary;
- protected byte[] _fixedSizeDataBytes;
- protected ByteBuffer _fixedSizeData;
- protected byte[] _variableSizeDataBytes;
- protected ByteBuffer _variableSizeData;
+ /// Both sections are read with absolute offsets only, so they can be split
in pages (as the server builds them) or
+ /// be a single array (as the broker deserializes them).
+ @Nullable
+ protected DataBuffer _fixedSizeData;
+ @Nullable
+ protected DataBuffer _variableSizeData;
protected Map<String, String> _metadata;
protected int[] _columnOffsets;
protected int _rowSizeInBytes;
@@ -102,9 +106,7 @@ public class DataTableImplV4 implements DataTable {
_fixDataSize = 0;
_dataSchema = null;
_stringDictionary = null;
- _fixedSizeDataBytes = null;
_fixedSizeData = null;
- _variableSizeDataBytes = null;
_variableSizeData = null;
_rowSizeInBytes = 0;
_columnOffsets = null;
@@ -112,17 +114,17 @@ public class DataTableImplV4 implements DataTable {
_errCodeToExceptionMap = new HashMap<>();
}
- public DataTableImplV4(int numRows, DataSchema dataSchema, String[]
stringDictionary, byte[] fixedSizeDataBytes,
- byte[] variableSizeDataBytes) {
+ /// Creates a table that reads its fixed and variable size sections from the
given big-endian buffers, without
+ /// copying them.
+ public DataTableImplV4(int numRows, DataSchema dataSchema, String[]
stringDictionary, DataBuffer fixedSizeData,
+ DataBuffer variableSizeData) {
_numRows = numRows;
_dataSchema = dataSchema;
_numColumns = dataSchema == null ? 0 : dataSchema.size();
_fixDataSize = 0;
_stringDictionary = stringDictionary;
- _fixedSizeDataBytes = fixedSizeDataBytes;
- _fixedSizeData = ByteBuffer.wrap(fixedSizeDataBytes);
- _variableSizeDataBytes = variableSizeDataBytes;
- _variableSizeData = ByteBuffer.wrap(variableSizeDataBytes);
+ _fixedSizeData = fixedSizeData;
+ _variableSizeData = variableSizeData;
_metadata = new HashMap<>();
_errCodeToExceptionMap = new HashMap<>();
computeBlockObjectConstants();
@@ -170,22 +172,21 @@ public class DataTableImplV4 implements DataTable {
// Read fixed size data.
if (fixedSizeDataLength != 0) {
- _fixedSizeDataBytes = new byte[fixedSizeDataLength];
+ byte[] fixedSizeDataBytes = new byte[fixedSizeDataLength];
byteBuffer.position(fixedSizeDataStart);
- byteBuffer.get(_fixedSizeDataBytes);
- _fixedSizeData = ByteBuffer.wrap(_fixedSizeDataBytes);
+ byteBuffer.get(fixedSizeDataBytes);
+ _fixedSizeData = PinotByteBuffer.wrap(fixedSizeDataBytes);
} else {
- _fixedSizeDataBytes = null;
_fixedSizeData = null;
}
// Read variable size data.
- _variableSizeDataBytes = new byte[variableSizeDataLength];
+ byte[] variableSizeDataBytes = new byte[variableSizeDataLength];
if (variableSizeDataLength != 0) {
byteBuffer.position(variableSizeDataStart);
- byteBuffer.get(_variableSizeDataBytes);
+ byteBuffer.get(variableSizeDataBytes);
}
- _variableSizeData = ByteBuffer.wrap(_variableSizeDataBytes);
+ _variableSizeData = PinotByteBuffer.wrap(variableSizeDataBytes);
// Read metadata.
int metadataLength = byteBuffer.getInt();
@@ -243,10 +244,8 @@ public class DataTableImplV4 implements DataTable {
@Override
public BigDecimal getBigDecimal(int rowId, int colId) {
- int size = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- ByteBuffer byteBuffer = _variableSizeData.slice();
- byteBuffer.limit(size);
- return BigDecimalUtils.deserialize(byteBuffer);
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ return BigDecimalUtils.deserialize(copyFromVariableBuffer(offsetInFixed));
}
@Override
@@ -256,64 +255,70 @@ public class DataTableImplV4 implements DataTable {
@Override
public ByteArray getBytes(int rowId, int colId) {
- int size = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- byte[] buffer = new byte[size];
- _variableSizeData.get(buffer);
- return new ByteArray(buffer);
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ return new ByteArray(copyFromVariableBuffer(offsetInFixed));
}
// --------------------------------------------------------------------------
// Variable sized element access.
// --------------------------------------------------------------------------
+ // These read with absolute offsets instead of a stream: a stream over a
paged buffer creates a view of every page
+ // up to the end of the section on each call.
@Override
public int[] getIntArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- int[] ints = new int[length];
- for (int i = 0; i < length; i++) {
- ints[i] = _variableSizeData.getInt();
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ int[] ints = new int[_fixedSizeData.getInt(offsetInFixed + Integer.BYTES)];
+ for (int i = 0; i < ints.length; i++) {
+ ints[i] = _variableSizeData.getInt(offsetInVar + (long) i *
Integer.BYTES);
}
return ints;
}
@Override
public long[] getLongArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- long[] longs = new long[length];
- for (int i = 0; i < length; i++) {
- longs[i] = _variableSizeData.getLong();
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ long[] longs = new long[_fixedSizeData.getInt(offsetInFixed +
Integer.BYTES)];
+ for (int i = 0; i < longs.length; i++) {
+ longs[i] = _variableSizeData.getLong(offsetInVar + (long) i *
Long.BYTES);
}
return longs;
}
@Override
public float[] getFloatArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- float[] floats = new float[length];
- for (int i = 0; i < length; i++) {
- floats[i] = _variableSizeData.getFloat();
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ float[] floats = new float[_fixedSizeData.getInt(offsetInFixed +
Integer.BYTES)];
+ for (int i = 0; i < floats.length; i++) {
+ floats[i] = _variableSizeData.getFloat(offsetInVar + (long) i *
Float.BYTES);
}
return floats;
}
@Override
public double[] getDoubleArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- double[] doubles = new double[length];
- for (int i = 0; i < length; i++) {
- doubles[i] = _variableSizeData.getDouble();
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ double[] doubles = new double[_fixedSizeData.getInt(offsetInFixed +
Integer.BYTES)];
+ for (int i = 0; i < doubles.length; i++) {
+ doubles[i] = _variableSizeData.getDouble(offsetInVar + (long) i *
Double.BYTES);
}
return doubles;
}
@Override
public BigDecimal[] getBigDecimalArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- BigDecimal[] bigDecimals = new BigDecimal[length];
- for (int i = 0; i < length; i++) {
- int byteLength = _variableSizeData.getInt();
- byte[] bytes = new byte[byteLength];
- _variableSizeData.get(bytes);
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ BigDecimal[] bigDecimals = new
BigDecimal[_fixedSizeData.getInt(offsetInFixed + Integer.BYTES)];
+ for (int i = 0; i < bigDecimals.length; i++) {
+ byte[] bytes = new byte[_variableSizeData.getInt(offsetInVar)];
+ offsetInVar += Integer.BYTES;
+ _variableSizeData.copyTo(offsetInVar, bytes, 0, bytes.length);
+ offsetInVar += bytes.length;
bigDecimals[i] = BigDecimalUtils.deserialize(bytes);
}
return bigDecimals;
@@ -321,22 +326,25 @@ public class DataTableImplV4 implements DataTable {
@Override
public String[] getStringArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- String[] strings = new String[length];
- for (int i = 0; i < length; i++) {
- strings[i] = _stringDictionary[_variableSizeData.getInt()];
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ String[] strings = new String[_fixedSizeData.getInt(offsetInFixed +
Integer.BYTES)];
+ for (int i = 0; i < strings.length; i++) {
+ strings[i] = _stringDictionary[_variableSizeData.getInt(offsetInVar +
(long) i * Integer.BYTES)];
}
return strings;
}
@Override
public ByteArray[] getBytesArray(int rowId, int colId) {
- int length = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- ByteArray[] bytes = new ByteArray[length];
- for (int i = 0; i < length; i++) {
- int byteLength = _variableSizeData.getInt();
- byte[] value = new byte[byteLength];
- _variableSizeData.get(value);
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ long offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ ByteArray[] bytes = new ByteArray[_fixedSizeData.getInt(offsetInFixed +
Integer.BYTES)];
+ for (int i = 0; i < bytes.length; i++) {
+ byte[] value = new byte[_variableSizeData.getInt(offsetInVar)];
+ offsetInVar += Integer.BYTES;
+ _variableSizeData.copyTo(offsetInVar, value, 0, value.length);
+ offsetInVar += value.length;
bytes[i] = new ByteArray(value);
}
return bytes;
@@ -345,27 +353,27 @@ public class DataTableImplV4 implements DataTable {
@Nullable
@Override
public Map<String, Object> getMap(int rowId, int colId) {
- int size = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- if (size == 0) {
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ if (_fixedSizeData.getInt(offsetInFixed + Integer.BYTES) == 0) {
return null;
}
- ByteBuffer buffer = _variableSizeData.slice();
- buffer.limit(size);
- return MapUtils.deserializeMap(buffer);
+ return MapUtils.deserializeMap(copyFromVariableBuffer(offsetInFixed));
}
@Nullable
@Override
public CustomObject getCustomObject(int rowId, int colId) {
- int size = positionOffsetInVariableBufferAndGetLength(rowId, colId);
- int type = _variableSizeData.getInt();
+ int offsetInFixed = getOffsetInFixedBuffer(rowId, colId);
+ int offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ int size = _fixedSizeData.getInt(offsetInFixed + Integer.BYTES);
+ int type = _variableSizeData.getInt(offsetInVar);
if (size == 0) {
assert type == CustomObject.NULL_TYPE_VALUE;
return null;
}
- ByteBuffer buffer = _variableSizeData.slice();
- buffer.limit(size);
- return new CustomObject(type, buffer);
+ // A view when the value fits in one page, otherwise a copy. On a table
built in this process the pages are
+ // read-only, so the view is read-only too. See DataTable.getCustomObject
for the contract.
+ return new CustomObject(type, _variableSizeData.copyOrView(offsetInVar +
Integer.BYTES, size));
}
@Nullable
@@ -373,16 +381,13 @@ public class DataTableImplV4 implements DataTable {
public RoaringBitmap getNullRowIds(int colId) {
// _fixedSizeData stores two ints per col's null bitmap: offset, and
length.
int position = _fixDataSize + colId * Integer.BYTES * 2;
- if (_fixedSizeData == null || position >= _fixedSizeData.limit()) {
+ if (_fixedSizeData == null || position >= _fixedSizeData.size()) {
return null;
}
- _fixedSizeData.position(position);
- int offset = _fixedSizeData.getInt();
- int bytesLength = _fixedSizeData.getInt();
+ int bytesLength = _fixedSizeData.getInt(position + Integer.BYTES);
if (bytesLength > 0) {
- _variableSizeData.position(offset);
byte[] nullBitmapBytes = new byte[bytesLength];
- _variableSizeData.get(nullBitmapBytes);
+ _variableSizeData.copyTo(_fixedSizeData.getInt(position),
nullBitmapBytes, 0, bytesLength);
return RoaringBitmapUtils.deserialize(nullBitmapBytes);
} else {
return null;
@@ -428,23 +433,44 @@ public class DataTableImplV4 implements DataTable {
return _errCodeToExceptionMap;
}
+ /// Serializes the table into one array of the exact final size, so the
fixed and variable size sections are copied
+ /// once and no buffer grows by doubling.
@Override
public byte[] toBytes()
throws IOException {
ThreadResourceSnapshot resourceSnapshot = new ThreadResourceSnapshot();
- ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
- DataOutputStream dataOutputStream = new
DataOutputStream(byteArrayOutputStream);
- writeLeadingSections(dataOutputStream);
-
// Add table serialization time and memory metadata when the corresponding
measurement is enabled.
// When CPU time/memory usage is not collectable, we omit these values
from the metadata.
- if (ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled()) {
- getMetadata().put(MetadataKey.RESPONSE_SER_CPU_TIME_NS.getName(),
- String.valueOf(resourceSnapshot.getCpuTimeNs()));
- }
- if (ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled()) {
- getMetadata().put(MetadataKey.RESPONSE_SER_MEM_ALLOCATED_BYTES.getName(),
+ // The entries are reserved before the array is sized and filled in after
the sections are copied. Both are LONG
+ // keys, whose encoded length does not depend on the value, so the
metadata length does not change.
+ boolean measureCpuTime =
ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled();
+ boolean measureAllocatedBytes =
ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled();
+ if (measureCpuTime) {
+ _metadata.put(MetadataKey.RESPONSE_SER_CPU_TIME_NS.getName(), "0");
+ }
+ if (measureAllocatedBytes) {
+ _metadata.put(MetadataKey.RESPONSE_SER_MEM_ALLOCATED_BYTES.getName(),
"0");
+ }
+
+ byte[] exceptionsBytes = serializeExceptions();
+ byte[] dictionaryBytes = _stringDictionary != null ?
serializeStringDictionary() : null;
+ byte[] dataSchemaBytes = _dataSchema != null ? _dataSchema.toBytes() :
null;
+ int metadataLength = serializeMetadata().length;
+ long totalLength = HEADER_SIZE + exceptionsBytes.length +
lengthOf(dictionaryBytes) + lengthOf(dataSchemaBytes)
+ + sizeOf(_fixedSizeData) + sizeOf(_variableSizeData) + Integer.BYTES +
metadataLength;
+ if (totalLength > Integer.MAX_VALUE) {
+ throw new IllegalStateException("Serialized data table size " +
totalLength + " exceeds the 2 GB limit");
+ }
+ byte[] bytes = new byte[(int) totalLength];
+ ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
+ writeLeadingSections(byteBuffer, exceptionsBytes, dictionaryBytes,
dataSchemaBytes);
+
+ if (measureCpuTime) {
+ _metadata.put(MetadataKey.RESPONSE_SER_CPU_TIME_NS.getName(),
String.valueOf(resourceSnapshot.getCpuTimeNs()));
+ }
+ if (measureAllocatedBytes) {
+ _metadata.put(MetadataKey.RESPONSE_SER_MEM_ALLOCATED_BYTES.getName(),
String.valueOf(resourceSnapshot.getAllocatedBytes()));
}
@@ -452,83 +478,74 @@ public class DataTableImplV4 implements DataTable {
// NOTE: We ignore metadata serialization time in
"responseSerializationCpuTimeNs" as it's negligible while
// considering it will bring a lot code complexity.
byte[] metadataBytes = serializeMetadata();
- dataOutputStream.writeInt(metadataBytes.length);
- dataOutputStream.write(metadataBytes);
+ if (metadataBytes.length != metadataLength) {
+ throw new IllegalStateException(
+ "Metadata length changed during serialization from " +
metadataLength + " to " + metadataBytes.length);
+ }
+ byteBuffer.putInt(metadataBytes.length);
+ byteBuffer.put(metadataBytes);
+ assert !byteBuffer.hasRemaining();
+ return bytes;
+ }
- return byteArrayOutputStream.toByteArray();
+ private static int lengthOf(@Nullable byte[] bytes) {
+ return bytes != null ? bytes.length : 0;
}
- private void writeLeadingSections(DataOutputStream dataOutputStream)
- throws IOException {
- dataOutputStream.writeInt(getVersion());
- dataOutputStream.writeInt(_numRows);
- dataOutputStream.writeInt(_numColumns);
+ private static long sizeOf(@Nullable DataBuffer dataBuffer) {
+ return dataBuffer != null ? dataBuffer.size() : 0;
+ }
+
+ private void writeLeadingSections(ByteBuffer byteBuffer, byte[]
exceptionsBytes, @Nullable byte[] dictionaryBytes,
+ @Nullable byte[] dataSchemaBytes) {
+ int fixedSizeDataLength = (int) sizeOf(_fixedSizeData);
+ int variableSizeDataLength = (int) sizeOf(_variableSizeData);
+
+ byteBuffer.putInt(getVersion());
+ byteBuffer.putInt(_numRows);
+ byteBuffer.putInt(_numColumns);
int dataOffset = HEADER_SIZE;
// Write exceptions section offset(START|SIZE).
- dataOutputStream.writeInt(dataOffset);
- byte[] exceptionsBytes;
- exceptionsBytes = serializeExceptions();
- dataOutputStream.writeInt(exceptionsBytes.length);
+ byteBuffer.putInt(dataOffset);
+ byteBuffer.putInt(exceptionsBytes.length);
dataOffset += exceptionsBytes.length;
// Write dictionary map section offset(START|SIZE).
- dataOutputStream.writeInt(dataOffset);
- byte[] dictionaryBytes = null;
- if (_stringDictionary != null) {
- dictionaryBytes = serializeStringDictionary();
- dataOutputStream.writeInt(dictionaryBytes.length);
- dataOffset += dictionaryBytes.length;
- } else {
- dataOutputStream.writeInt(0);
- }
+ byteBuffer.putInt(dataOffset);
+ byteBuffer.putInt(lengthOf(dictionaryBytes));
+ dataOffset += lengthOf(dictionaryBytes);
// Write data schema section offset(START|SIZE).
- dataOutputStream.writeInt(dataOffset);
- byte[] dataSchemaBytes = null;
- if (_dataSchema != null) {
- dataSchemaBytes = _dataSchema.toBytes();
- dataOutputStream.writeInt(dataSchemaBytes.length);
- dataOffset += dataSchemaBytes.length;
- } else {
- dataOutputStream.writeInt(0);
- }
+ byteBuffer.putInt(dataOffset);
+ byteBuffer.putInt(lengthOf(dataSchemaBytes));
+ dataOffset += lengthOf(dataSchemaBytes);
// Write fixed size data section offset(START|SIZE).
- dataOutputStream.writeInt(dataOffset);
- if (_fixedSizeDataBytes != null) {
- dataOutputStream.writeInt(_fixedSizeDataBytes.length);
- dataOffset += _fixedSizeDataBytes.length;
- } else {
- dataOutputStream.writeInt(0);
- }
+ byteBuffer.putInt(dataOffset);
+ byteBuffer.putInt(fixedSizeDataLength);
+ dataOffset += fixedSizeDataLength;
// Write variable size data section offset(START|SIZE).
- dataOutputStream.writeInt(dataOffset);
- if (_variableSizeDataBytes != null) {
- dataOutputStream.writeInt(_variableSizeDataBytes.length);
- } else {
- dataOutputStream.writeInt(0);
- }
+ byteBuffer.putInt(dataOffset);
+ byteBuffer.putInt(variableSizeDataLength);
// Write actual data.
- // Write exceptions bytes.
- dataOutputStream.write(exceptionsBytes);
- // Write dictionary map bytes.
+ byteBuffer.put(exceptionsBytes);
if (dictionaryBytes != null) {
- dataOutputStream.write(dictionaryBytes);
+ byteBuffer.put(dictionaryBytes);
}
- // Write data schema bytes.
if (dataSchemaBytes != null) {
- dataOutputStream.write(dataSchemaBytes);
+ byteBuffer.put(dataSchemaBytes);
}
- // Write fixed size data bytes.
- if (_fixedSizeDataBytes != null) {
- dataOutputStream.write(_fixedSizeDataBytes);
- }
- // Write variable size data bytes.
- if (_variableSizeDataBytes != null) {
- dataOutputStream.write(_variableSizeDataBytes);
+ copySection(_fixedSizeData, byteBuffer, fixedSizeDataLength);
+ copySection(_variableSizeData, byteBuffer, variableSizeDataLength);
+ }
+
+ private static void copySection(@Nullable DataBuffer section, ByteBuffer
byteBuffer, int length) {
+ if (length > 0) {
+ section.copyTo(0, byteBuffer.array(), byteBuffer.position(), length);
+ byteBuffer.position(byteBuffer.position() + length);
}
}
@@ -648,7 +665,7 @@ public class DataTableImplV4 implements DataTable {
}
}
- /// return the offset in `_fixedSizeDataBytes` of the row/column ID.
+ /// Returns the offset in the fixed size buffer of the row/column ID.
/// @param rowId row ID
/// @param colId column ID
/// @return the offset in the fixed size buffer for the row/columnID.
@@ -656,16 +673,13 @@ public class DataTableImplV4 implements DataTable {
return rowId * _rowSizeInBytes + _columnOffsets[colId];
}
- /// position the `_variableSizeDataBytes` to the corresponding row/column
ID. and return the
- /// length of bytes to extract from the variable size buffer.
- ///
- /// @param rowId row ID
- /// @param colId column ID
- /// @return the length to extract from variable size buffer.
- protected int positionOffsetInVariableBufferAndGetLength(int rowId, int
colId) {
- int offset = getOffsetInFixedBuffer(rowId, colId);
- _variableSizeData.position(_fixedSizeData.getInt(offset));
- return _fixedSizeData.getInt(offset + 4);
+ /// Copies the variable size value referenced by the (offset, length) pair
at `offsetInFixed` into a new array.
+ private byte[] copyFromVariableBuffer(int offsetInFixed) {
+ int offsetInVar = _fixedSizeData.getInt(offsetInFixed);
+ int length = _fixedSizeData.getInt(offsetInFixed + Integer.BYTES);
+ byte[] bytes = new byte[length];
+ _variableSizeData.copyTo(offsetInVar, bytes, 0, length);
+ return bytes;
}
@Override
@@ -678,7 +692,7 @@ public class DataTableImplV4 implements DataTable {
@Override
public DataTable toDataOnlyDataTable() {
- return new DataTableImplV4(_numRows, _dataSchema, _stringDictionary,
_fixedSizeDataBytes, _variableSizeDataBytes);
+ return new DataTableImplV4(_numRows, _dataSchema, _stringDictionary,
_fixedSizeData, _variableSizeData);
}
public int getRowSizeInBytes() {
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
index ca62f32c723..41f054027ff 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/BaseDataTableBuilder.java
@@ -18,8 +18,6 @@
*/
package org.apache.pinot.core.common.datatable;
-import java.io.ByteArrayOutputStream;
-import java.io.DataOutputStream;
import java.io.IOException;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
@@ -32,6 +30,7 @@ import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.common.utils.RoaringBitmapUtils;
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.segment.spi.memory.PagedPinotOutputStream;
import org.apache.pinot.spi.utils.BigDecimalUtils;
import org.apache.pinot.spi.utils.ByteArray;
import org.apache.pinot.spi.utils.MapUtils;
@@ -60,17 +59,19 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
/// decodes as garbage in an 8-byte `LONG` / `DOUBLE` one.
private static final EnumSet<ColumnDataType> IN_BAND_NULL_TYPES =
EnumSet.of(ColumnDataType.OBJECT, ColumnDataType.UNKNOWN,
ColumnDataType.MAP);
+ /// Heap pages need no release, so one stateless allocator serves every
builder.
+ private static final PagedPinotOutputStream.PageAllocator PAGE_ALLOCATOR =
+ PagedPinotOutputStream.HeapPageAllocator.createSmall();
protected final DataSchema _dataSchema;
protected final int _version;
protected final int[] _columnOffsets;
protected final ByteBuffer _currentRowDataByteBuffer;
- protected final ByteArrayOutputStream _fixedSizeDataByteArrayOutputStream =
new ByteArrayOutputStream();
- protected final DataOutputStream _fixedSizeDataOutputStream =
- new DataOutputStream(_fixedSizeDataByteArrayOutputStream);
- protected final ByteArrayOutputStream _variableSizeDataByteArrayOutputStream
= new ByteArrayOutputStream();
- protected final DataOutputStream _variableSizeDataOutputStream =
- new DataOutputStream(_variableSizeDataByteArrayOutputStream);
+ /// Both sections are written in heap pages of
[PagedPinotOutputStream.PageAllocator#MIN_RECOMMENDED_PAGE_SIZE]
+ /// bytes and handed to the built table as is, without copying them into one
contiguous array. A large response
+ /// therefore never needs a humongous allocation while it is built.
+ protected final PagedPinotOutputStream _fixedSizeDataOutputStream = new
PagedPinotOutputStream(PAGE_ALLOCATOR);
+ protected final PagedPinotOutputStream _variableSizeDataOutputStream = new
PagedPinotOutputStream(PAGE_ALLOCATOR);
private final ColumnDataType[] _storedColumnDataTypes;
/// Entries are allocated lazily so a table without nulls pays nothing.
@@ -94,6 +95,12 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
_nullBitmaps = new RoaringBitmap[dataSchema.size()];
}
+ /// Returns the offset of the next byte in the variable size section. The
wire format stores it as an int. The
+ /// section had the same 2 GB limit when it was a `ByteArrayOutputStream`,
which failed with an `OutOfMemoryError`.
+ protected int variableSizeDataOffset() {
+ return Math.toIntExact(_variableSizeDataOutputStream.getCurrentOffset());
+ }
+
@Override
public void startRow() {
_numRows++;
@@ -127,23 +134,23 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, BigDecimal value)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
byte[] bytes = BigDecimalUtils.serialize(value);
_currentRowDataByteBuffer.putInt(bytes.length);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _variableSizeDataOutputStream.write(bytes);
}
@Override
public void setColumn(int colId, @Nullable Map<String, Object> value)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
if (value == null) {
_currentRowDataByteBuffer.putInt(0);
} else {
byte[] bytes = MapUtils.serializeMap(value, false);
_currentRowDataByteBuffer.putInt(bytes.length);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _variableSizeDataOutputStream.write(bytes);
}
}
@@ -151,7 +158,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, int[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (int value : values) {
_variableSizeDataOutputStream.writeInt(value);
@@ -162,7 +169,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, long[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (long value : values) {
_variableSizeDataOutputStream.writeLong(value);
@@ -173,7 +180,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, float[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (float value : values) {
_variableSizeDataOutputStream.writeFloat(value);
@@ -184,7 +191,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, double[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (double value : values) {
_variableSizeDataOutputStream.writeDouble(value);
@@ -195,12 +202,12 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, BigDecimal[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (BigDecimal value : values) {
byte[] bytes = BigDecimalUtils.serialize(value);
_variableSizeDataOutputStream.writeInt(bytes.length);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _variableSizeDataOutputStream.write(bytes);
}
}
@@ -208,12 +215,12 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId, ByteArray[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (ByteArray value : values) {
byte[] bytes = value.getBytes();
_variableSizeDataOutputStream.writeInt(bytes.length);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _variableSizeDataOutputStream.write(bytes);
}
}
@@ -221,7 +228,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
public void setColumn(int colId,
AggregationFunction.SerializedIntermediateResult value)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(value.getSize());
_variableSizeDataOutputStream.writeInt(value.getType());
value.writeTo(_variableSizeDataOutputStream);
@@ -233,7 +240,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
ColumnDataType storedColumnDataType = _storedColumnDataTypes[colId];
if (IN_BAND_NULL_TYPES.contains(storedColumnDataType)) {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(0);
_variableSizeDataOutputStream.writeInt(CustomObject.NULL_TYPE_VALUE);
return;
@@ -256,7 +263,7 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
@Override
public void finishRow()
throws IOException {
-
_fixedSizeDataByteArrayOutputStream.write(_currentRowDataByteBuffer.array());
+ _fixedSizeDataOutputStream.write(_currentRowDataByteBuffer.array());
}
@Override
@@ -289,13 +296,13 @@ public abstract class BaseDataTableBuilder implements
DataTableBuilder {
}
_nullBitmapsWritten = true;
for (RoaringBitmap nullBitmap : _nullBitmaps) {
-
_fixedSizeDataOutputStream.writeInt(_variableSizeDataByteArrayOutputStream.size());
+ _fixedSizeDataOutputStream.writeInt(variableSizeDataOffset());
if (nullBitmap == null) {
_fixedSizeDataOutputStream.writeInt(0);
} else {
byte[] bitmapBytes = RoaringBitmapUtils.serialize(nullBitmap);
_fixedSizeDataOutputStream.writeInt(bitmapBytes.length);
- _variableSizeDataByteArrayOutputStream.write(bitmapBytes);
+ _variableSizeDataOutputStream.write(bitmapBytes);
}
}
}
diff --git
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/DataTableBuilderV4.java
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/DataTableBuilderV4.java
index f4da8848140..1e4fcf1bed8 100644
---
a/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/DataTableBuilderV4.java
+++
b/pinot-core/src/main/java/org/apache/pinot/core/common/datatable/DataTableBuilderV4.java
@@ -22,6 +22,7 @@ import it.unimi.dsi.fastutil.objects.Object2IntMap;
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
import java.io.IOException;
import java.io.UncheckedIOException;
+import java.nio.ByteOrder;
import org.apache.pinot.common.datatable.DataTable;
import org.apache.pinot.common.datatable.DataTableFactory;
import org.apache.pinot.common.datatable.DataTableImplV4;
@@ -47,17 +48,17 @@ public class DataTableBuilderV4 extends
BaseDataTableBuilder {
public void setColumn(int colId, ByteArray value)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
byte[] bytes = value.getBytes();
_currentRowDataByteBuffer.putInt(bytes.length);
- _variableSizeDataByteArrayOutputStream.write(bytes);
+ _variableSizeDataOutputStream.write(bytes);
}
@Override
public void setColumn(int colId, String[] values)
throws IOException {
_currentRowDataByteBuffer.position(_columnOffsets[colId]);
-
_currentRowDataByteBuffer.putInt(_variableSizeDataByteArrayOutputStream.size());
+ _currentRowDataByteBuffer.putInt(variableSizeDataOffset());
_currentRowDataByteBuffer.putInt(values.length);
for (String value : values) {
int dictId = _dictionary.computeIntIfAbsent(value, k ->
_dictionary.size());
@@ -70,14 +71,17 @@ public class DataTableBuilderV4 extends
BaseDataTableBuilder {
try {
writeNullBitmaps();
} catch (IOException e) {
- // Both buffers are in-memory byte arrays, so this cannot happen.
+ // Both buffers are in-memory heap pages, so this cannot happen.
throw new UncheckedIOException(e);
}
String[] reverseDictionary = new String[_dictionary.size()];
for (Object2IntMap.Entry<String> entry : _dictionary.object2IntEntrySet())
{
reverseDictionary[entry.getIntValue()] = entry.getKey();
}
+ // The table views the written pages directly. Pages are heap memory, so
nothing has to be released. The views
+ // end at the bytes written so far, and the builder only appends, so later
writes do not change this table.
return new DataTableImplV4(_numRows, _dataSchema, reverseDictionary,
- _fixedSizeDataByteArrayOutputStream.toByteArray(),
_variableSizeDataByteArrayOutputStream.toByteArray());
+ _fixedSizeDataOutputStream.asBuffer(ByteOrder.BIG_ENDIAN, false),
+ _variableSizeDataOutputStream.asBuffer(ByteOrder.BIG_ENDIAN, false));
}
}
diff --git
a/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
b/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
index 521d2082916..41b2ed68b15 100644
---
a/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
+++
b/pinot-core/src/test/java/org/apache/pinot/core/common/datatable/DataTableSerDeTest.java
@@ -20,6 +20,7 @@ package org.apache.pinot.core.common.datatable;
import java.io.IOException;
import java.math.BigDecimal;
+import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashMap;
@@ -28,11 +29,13 @@ import java.util.Random;
import java.util.UUID;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.common.CustomObject;
import org.apache.pinot.common.datatable.DataTable;
import org.apache.pinot.common.datatable.DataTable.MetadataKey;
import org.apache.pinot.common.datatable.DataTableFactory;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.utils.ByteArray;
@@ -305,12 +308,99 @@ public class DataTableSerDeTest {
fillDataTableWithRandomData(dataTableBuilder, columnDataTypes, numColumns);
DataTable dataTable = dataTableBuilder.build();
+ // The built table reads the builder's pages, the deserialized one reads a
single array. Both must agree.
+ verifyDataIsSame(dataTable, columnDataTypes, numColumns);
DataTable newDataTable =
DataTableFactory.getDataTable(dataTable.toBytes());
assertEquals(newDataTable.getDataSchema(), dataSchema, ERROR_MESSAGE);
assertEquals(newDataTable.getNumberOfRows(), NUM_ROWS, ERROR_MESSAGE);
verifyDataIsSame(newDataTable, columnDataTypes, numColumns);
}
+ /// The builder writes the fixed and variable size sections in pages. Values
larger than a page, and values that
+ /// straddle a page boundary, must read back the same from the built table
and from the deserialized one.
+ @Test
+ public void testValuesSpanningPages()
+ throws IOException {
+ DataTableBuilderFactory.setDataTableVersion(DataTableFactory.VERSION_4);
+ DataSchema dataSchema = new DataSchema(
+ new String[]{"bytes", "ints", "object", "long", "strings",
"bigDecimals", "nullableInt"},
+ new ColumnDataType[]{ColumnDataType.BYTES, ColumnDataType.INT_ARRAY,
ColumnDataType.OBJECT,
+ ColumnDataType.LONG, ColumnDataType.STRING_ARRAY,
ColumnDataType.BIG_DECIMAL_ARRAY, ColumnDataType.INT});
+ // 52 bytes per row, so the fixed size section spans several pages as well.
+ int numRows = 2_000;
+ byte[][] bytesValues = new byte[numRows][];
+ int[][] intArrayValues = new int[numRows][];
+ byte[][] objectValues = new byte[numRows][];
+ String[][] stringArrayValues = new String[numRows][];
+ BigDecimal[][] bigDecimalArrayValues = new BigDecimal[numRows][];
+ DataTableBuilder dataTableBuilder =
DataTableBuilderFactory.getDataTableBuilder(dataSchema);
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ // Every 97th row carries values larger than a page.
+ boolean large = rowId % 97 == 0;
+ dataTableBuilder.startRow();
+ bytesValues[rowId] = new byte[large ? 20_000 + RANDOM.nextInt(20_000) :
RANDOM.nextInt(300)];
+ RANDOM.nextBytes(bytesValues[rowId]);
+ dataTableBuilder.setColumn(0, new ByteArray(bytesValues[rowId]));
+ intArrayValues[rowId] = RANDOM.ints(large ? 10_000 :
RANDOM.nextInt(100)).toArray();
+ dataTableBuilder.setColumn(1, intArrayValues[rowId]);
+ if (rowId % 10 == 1) {
+ dataTableBuilder.setNull(2);
+ } else {
+ objectValues[rowId] = new byte[large ? 30_000 : 1 +
RANDOM.nextInt(500)];
+ RANDOM.nextBytes(objectValues[rowId]);
+ dataTableBuilder.setColumn(2, new
AggregationFunction.SerializedIntermediateResult(rowId, objectValues[rowId]));
+ }
+ dataTableBuilder.setColumn(3, (long) rowId);
+ stringArrayValues[rowId] = new String[large ? 5_000 :
RANDOM.nextInt(20)];
+ for (int i = 0; i < stringArrayValues[rowId].length; i++) {
+ stringArrayValues[rowId][i] = "s" + RANDOM.nextInt(1_000);
+ }
+ dataTableBuilder.setColumn(4, stringArrayValues[rowId]);
+ bigDecimalArrayValues[rowId] = new BigDecimal[large ? 2_000 :
RANDOM.nextInt(20)];
+ for (int i = 0; i < bigDecimalArrayValues[rowId].length; i++) {
+ bigDecimalArrayValues[rowId][i] =
BigDecimal.valueOf(RANDOM.nextLong(), RANDOM.nextInt(10));
+ }
+ dataTableBuilder.setColumn(5, bigDecimalArrayValues[rowId]);
+ // A third of the rows are null, so the null bitmap section is not
trivial.
+ if (rowId % 3 == 0) {
+ dataTableBuilder.setNull(6);
+ } else {
+ dataTableBuilder.setColumn(6, rowId);
+ }
+ dataTableBuilder.finishRow();
+ }
+
+ DataTable dataTable = dataTableBuilder.build();
+ DataTable newDataTable =
DataTableFactory.getDataTable(dataTable.toBytes());
+ for (DataTable table : new DataTable[]{dataTable,
dataTable.toDataOnlyDataTable(), newDataTable}) {
+ assertEquals(table.getNumberOfRows(), numRows, ERROR_MESSAGE);
+ RoaringBitmap nullRowIds = table.getNullRowIds(6);
+ assertNotNull(nullRowIds, ERROR_MESSAGE);
+ for (int rowId = 0; rowId < numRows; rowId++) {
+ assertEquals(table.getBytes(rowId, 0).getBytes(), bytesValues[rowId],
ERROR_MESSAGE);
+ assertEquals(table.getIntArray(rowId, 1), intArrayValues[rowId],
ERROR_MESSAGE);
+ CustomObject customObject = table.getCustomObject(rowId, 2);
+ if (objectValues[rowId] == null) {
+ assertNull(customObject, ERROR_MESSAGE);
+ } else {
+ assertNotNull(customObject, ERROR_MESSAGE);
+ assertEquals(customObject.getType(), rowId, ERROR_MESSAGE);
+ ByteBuffer buffer = customObject.getBuffer();
+ byte[] objectBytes = new byte[buffer.remaining()];
+ buffer.get(objectBytes);
+ assertEquals(objectBytes, objectValues[rowId], ERROR_MESSAGE);
+ }
+ assertEquals(table.getLong(rowId, 3), rowId, ERROR_MESSAGE);
+ assertEquals(table.getStringArray(rowId, 4), stringArrayValues[rowId],
ERROR_MESSAGE);
+ assertEquals(table.getBigDecimalArray(rowId, 5),
bigDecimalArrayValues[rowId], ERROR_MESSAGE);
+ assertEquals(nullRowIds.contains(rowId), rowId % 3 == 0,
ERROR_MESSAGE);
+ if (rowId % 3 != 0) {
+ assertEquals(table.getInt(rowId, 6), rowId, ERROR_MESSAGE);
+ }
+ }
+ }
+ }
+
@Test(dataProvider = "versionProvider")
public void testAllDataTypes(int dataTableVersion)
throws IOException {
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/CompoundDataBuffer.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/CompoundDataBuffer.java
index ed6d2894676..0942bccfdcd 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/CompoundDataBuffer.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/memory/CompoundDataBuffer.java
@@ -25,6 +25,7 @@ import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.List;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
@@ -112,21 +113,19 @@ public class CompoundDataBuffer implements DataBuffer {
}
private int getBufferIndex(long offset) {
- // this optimistically assumes that lookups are going to be in ascending
order
- // we don't care about concurrency here given that this is only used to
speed up the lookup
+ // Most lookups land in the same buffer as the previous one, so check it
first. Otherwise use a binary search,
+ // which keeps random and backward access at O(log n) instead of a scan
from the first buffer.
+ // We don't care about concurrency here given that this is only used to
speed up the lookup.
int lastBufferIndex = _lastBufferIndex;
- if (_bufferOffsets[lastBufferIndex] > offset) {
- lastBufferIndex = 0;
- }
-
- for (int i = lastBufferIndex; i < _bufferOffsets.length; i++) {
- if (offset < _bufferOffsets[i]) {
- int result = i - 1;
- _lastBufferIndex = result;
- return result;
- }
- }
- int result = _bufferOffsets.length - 1;
+ int nextBufferIndex = lastBufferIndex + 1;
+ if (_bufferOffsets[lastBufferIndex] <= offset
+ && (nextBufferIndex == _bufferOffsets.length || offset <
_bufferOffsets[nextBufferIndex])) {
+ return lastBufferIndex;
+ }
+ // Offsets are strictly ascending because empty buffers are rejected. On a
miss, binarySearch returns
+ // (-insertionPoint - 1), and the buffer that contains the offset is the
one before the insertion point.
+ int index = Arrays.binarySearch(_bufferOffsets, offset);
+ int result = index >= 0 ? index : -index - 2;
_lastBufferIndex = result;
return result;
}
diff --git
a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/CompoundDataBufferTest.java
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/CompoundDataBufferTest.java
index 23a37d6fec0..8bd205adc14 100644
---
a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/CompoundDataBufferTest.java
+++
b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/memory/CompoundDataBufferTest.java
@@ -186,6 +186,32 @@ public class CompoundDataBufferTest {
}
}
+ /// Reads in random order, so that lookups move backward and forward across
buffers of different sizes, and also land
+ /// on the first and last byte of each buffer.
+ @Test
+ void randomAccessAcrossBuffersOfDifferentSizes() {
+ Random r = new Random(42);
+ int[] sizes = {1, 7, 16, 3, 64, 1, 32};
+ int totalSize = Arrays.stream(sizes).sum();
+ byte[] expected = new byte[totalSize];
+ r.nextBytes(expected);
+ DataBuffer[] buffers = new DataBuffer[sizes.length];
+ int start = 0;
+ for (int i = 0; i < sizes.length; i++) {
+ buffers[i] = PinotByteBuffer.wrap(Arrays.copyOfRange(expected, start,
start + sizes[i]));
+ start += sizes[i];
+ }
+ CompoundDataBuffer compoundDataBuffer = new CompoundDataBuffer(buffers,
ByteOrder.BIG_ENDIAN, false);
+
+ for (int i = 0; i < 10_000; i++) {
+ int offset = r.nextInt(totalSize);
+ assertEquals(compoundDataBuffer.getByte(offset), expected[offset],
"Offset " + offset);
+ }
+ for (int offset = totalSize - 1; offset >= 0; offset--) {
+ assertEquals(compoundDataBuffer.getByte(offset), expected[offset],
"Offset " + offset);
+ }
+ }
+
@Test
void readWriteByteBuffer() {
Random r = new Random(42);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
