You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@druid.apache.org by fr...@apache.org on 2026/09/11 06:29:34 UTC
(druid) 05/09: perf(sql): use query vector size for system table batches
This is an automated email from the ASF dual-hosted git repository.
FrankChen021 pushed a commit to branch codex/native-sys-segments
in repository https://gitbox.apache.org/repos/asf/druid.git
commit 628e678c2f242b6ec87cfd96d087b658dddc8b27
Author: Frank Chen <fr...@outlook.com>
AuthorDate: Wed Sep 2 16:32:25 2026 +0800
perf(sql): use query vector size for system table batches
---
.../calcite/schema/SysSegmentsSqlBenchmark.java | 172 +--------------------
.../druid/query/BatchedInlineDataSource.java | 29 ++--
.../druid/query/BatchedInlineDataSourceTest.java | 45 +++++-
3 files changed, 58 insertions(+), 188 deletions(-)
diff --git a/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsSqlBenchmark.java b/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsSqlBenchmark.java
index c244a3ea9f3..05dc8dba786 100644
--- a/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsSqlBenchmark.java
+++ b/benchmarks/src/test/java/org/apache/druid/sql/calcite/schema/SysSegmentsSqlBenchmark.java
@@ -30,13 +30,8 @@ import org.apache.druid.client.InternalQueryConfig;
import org.apache.druid.client.TimelineServerView;
import org.apache.druid.client.coordinator.CoordinatorClient;
import org.apache.druid.client.coordinator.NoopCoordinatorClient;
-import org.apache.druid.frame.allocation.ArenaMemoryAllocatorFactory;
-import org.apache.druid.frame.segment.FrameCursorUtils;
-import org.apache.druid.frame.write.FrameWriterFactory;
-import org.apache.druid.frame.write.FrameWriters;
import org.apache.druid.java.util.common.CloseableIterators;
import org.apache.druid.java.util.common.Intervals;
-import org.apache.druid.java.util.common.Pair;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.io.Closer;
@@ -45,19 +40,13 @@ import org.apache.druid.query.BatchedInlineDataSource;
import org.apache.druid.query.DataSource;
import org.apache.druid.query.DefaultGenericQueryMetricsFactory;
import org.apache.druid.query.DefaultQueryConfig;
-import org.apache.druid.query.FrameBasedInlineDataSource;
-import org.apache.druid.query.FrameSignaturePair;
import org.apache.druid.query.InlineDataSource;
-import org.apache.druid.query.IterableRowsCursorHelper;
import org.apache.druid.query.QueryRunnerFactoryConglomerate;
import org.apache.druid.query.SystemTableDataSource;
import org.apache.druid.query.filter.DimFilter;
import org.apache.druid.query.policy.NoopPolicyEnforcer;
-import org.apache.druid.query.scan.ScanQuery;
import org.apache.druid.query.scan.ScanQueryEngine;
import org.apache.druid.rpc.indexing.NoopOverlordClient;
-import org.apache.druid.segment.Cursor;
-import org.apache.druid.segment.FrameBasedInlineSegmentWrangler;
import org.apache.druid.segment.InlineSegmentWrangler;
import org.apache.druid.segment.MapSegmentWrangler;
import org.apache.druid.segment.join.JoinableFactory;
@@ -99,7 +88,6 @@ import org.apache.druid.timeline.DataSegment;
import org.apache.druid.timeline.SegmentId;
import org.apache.druid.timeline.SegmentStatusInCluster;
import org.apache.druid.timeline.partition.LinearShardSpec;
-import org.apache.druid.utils.CloseableUtils;
import org.easymock.EasyMock;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -119,7 +107,6 @@ import org.openjdk.jmh.infra.Blackhole;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -150,9 +137,6 @@ public class SysSegmentsSqlBenchmark
private static final String BINDABLE = "bindable";
private static final String NATIVE_ROW = "nativeRow";
private static final String NATIVE_PROVIDER = "nativeProvider";
- private static final String NATIVE_BATCHED = "nativeBatched";
- private static final String NATIVE_FRAME_BUILD = "nativeFrameBuild";
- private static final String NATIVE_FRAME_CACHED = "nativeFrameCached";
private static final String SQL = "SELECT\n"
+ "datasource,\n"
+ "COUNT(*) FILTER (WHERE is_active = 1) AS num_segments,\n"
@@ -193,14 +177,11 @@ public class SysSegmentsSqlBenchmark
private PlannerFactory plannerFactory;
private SqlEngine rowEngine;
private SqlEngine providerEngine;
- private SqlEngine batchedEngine;
- private SqlEngine frameBuildEngine;
- private SqlEngine frameCachedEngine;
@State(Scope.Thread)
public static class ExecutionState
{
- @Param({BINDABLE, NATIVE_ROW, NATIVE_PROVIDER, NATIVE_BATCHED, NATIVE_FRAME_BUILD, NATIVE_FRAME_CACHED})
+ @Param({BINDABLE, NATIVE_ROW, NATIVE_PROVIDER})
private String executionPath;
private PreparedQuery preparedQuery;
@@ -256,104 +237,6 @@ public class SysSegmentsSqlBenchmark
}
}
- private static class FrameSystemTableQueryHandler extends SystemTableQueryHandler
- {
- private final SystemTableQueryHandler rowHandler;
- private final Map<FrameCacheKey, FrameBasedInlineDataSource> frameCache;
-
- /**
- * Benchmark-only cache key. The benchmark uses immutable segment metadata and a fixed authorization setup;
- * production use would additionally require a versioned metadata snapshot and bounded cache lifecycle.
- */
- private record FrameCacheKey(
- List<String> columns,
- DimFilter filter,
- String identity,
- String authorizerName
- )
- {
- }
-
- FrameSystemTableQueryHandler(final SystemTableQueryHandler rowHandler, final boolean cacheFrames)
- {
- super(Map.of(), Map.of(), new ScanQueryEngine(), AuthTestUtils.TEST_AUTHORIZER_MAPPER);
- this.rowHandler = rowHandler;
- this.frameCache = cacheFrames ? new HashMap<>() : null;
- }
-
- @Override
- public DataSource resolveDataSource(
- final ScanQuery query,
- final AuthenticationResult requestAuthenticationResult
- )
- {
- if (frameCache == null) {
- return makeFrameDataSource(query, requestAuthenticationResult);
- }
- final FrameCacheKey cacheKey = new FrameCacheKey(
- List.copyOf(query.getColumns()),
- query.getFilter(),
- requestAuthenticationResult.getIdentity(),
- requestAuthenticationResult.getAuthorizerName()
- );
- return frameCache.computeIfAbsent(cacheKey, ignored -> makeFrameDataSource(query, requestAuthenticationResult));
- }
-
- private FrameBasedInlineDataSource makeFrameDataSource(
- final ScanQuery query,
- final AuthenticationResult requestAuthenticationResult
- )
- {
- final InlineDataSource inlineDataSource = (InlineDataSource) rowHandler.resolveDataSource(
- query,
- requestAuthenticationResult
- );
- final Pair<Cursor, java.io.Closeable> cursorAndCloseable = IterableRowsCursorHelper.getCursorFromIterable(
- inlineDataSource.getRows(),
- inlineDataSource.getRowSignature()
- );
- final FrameWriterFactory frameWriterFactory = FrameWriters.makeColumnBasedFrameWriterFactory(
- ArenaMemoryAllocatorFactory.makeDefault(),
- inlineDataSource.getRowSignature(),
- List.of()
- );
- try {
- final List<FrameSignaturePair> frames = FrameCursorUtils.cursorToFramesSequence(
- cursorAndCloseable.lhs,
- frameWriterFactory
- ).map(frame -> new FrameSignaturePair(frame, inlineDataSource.getRowSignature())).toList();
- return new FrameBasedInlineDataSource(frames, inlineDataSource.getRowSignature());
- }
- finally {
- CloseableUtils.closeAndWrapExceptions(cursorAndCloseable.rhs);
- }
- }
- }
-
- private static class BatchedSystemTableQueryHandler extends SystemTableQueryHandler
- {
- private final SystemTableQueryHandler rowHandler;
-
- BatchedSystemTableQueryHandler(final SystemTableQueryHandler rowHandler)
- {
- super(Map.of(), Map.of(), new ScanQueryEngine(), AuthTestUtils.TEST_AUTHORIZER_MAPPER);
- this.rowHandler = rowHandler;
- }
-
- @Override
- public DataSource resolveDataSource(
- final ScanQuery query,
- final AuthenticationResult requestAuthenticationResult
- )
- {
- final InlineDataSource inlineDataSource = (InlineDataSource) rowHandler.resolveDataSource(
- query,
- requestAuthenticationResult
- );
- return new BatchedInlineDataSource(inlineDataSource.getRows(), inlineDataSource.getRowSignature());
- }
- }
-
/** Keeps the former row-only native path available as a stable benchmark baseline. */
private static class RowOnlySystemTableDataProvider implements SystemTableDataProvider
{
@@ -437,8 +320,6 @@ public class SysSegmentsSqlBenchmark
Map.of(
InlineDataSource.class,
new InlineSegmentWrangler(),
- FrameBasedInlineDataSource.class,
- new FrameBasedInlineSegmentWrangler(),
BatchedInlineDataSource.class,
new BatchedInlineDataSource.Wrangler()
)
@@ -464,24 +345,6 @@ public class SysSegmentsSqlBenchmark
);
rowEngine = makeEngine(conglomerate, walker, descriptor, rowQueryHandler);
providerEngine = makeEngine(conglomerate, walker, descriptor, providerQueryHandler);
- batchedEngine = makeEngine(
- conglomerate,
- walker,
- descriptor,
- new BatchedSystemTableQueryHandler(rowQueryHandler)
- );
- frameBuildEngine = makeEngine(
- conglomerate,
- walker,
- descriptor,
- new FrameSystemTableQueryHandler(rowQueryHandler, false)
- );
- frameCachedEngine = makeEngine(
- conglomerate,
- walker,
- descriptor,
- new FrameSystemTableQueryHandler(rowQueryHandler, true)
- );
final PlannerConfig plannerConfig = new PlannerConfig();
final TimelineServerView timelineServerView = new TestTimelineServerView(Collections.emptyList());
@@ -536,8 +399,7 @@ public class SysSegmentsSqlBenchmark
);
final List<Object[]> bindableResults = runQuery(BINDABLE);
- for (final String executionPath :
- List.of(NATIVE_ROW, NATIVE_PROVIDER, NATIVE_BATCHED, NATIVE_FRAME_BUILD, NATIVE_FRAME_CACHED)) {
+ for (final String executionPath : List.of(NATIVE_ROW, NATIVE_PROVIDER)) {
final List<Object[]> nativeResults = runQuery(executionPath);
if (bindableResults.size() != NUM_DATASOURCES || !rowsEqual(bindableResults, nativeResults)) {
throw new IllegalStateException("Bindable and native benchmark results do not match for " + executionPath);
@@ -631,18 +493,6 @@ public class SysSegmentsSqlBenchmark
engine = providerEngine;
context = NATIVE_CONTEXT;
break;
- case NATIVE_BATCHED:
- engine = batchedEngine;
- context = NATIVE_CONTEXT;
- break;
- case NATIVE_FRAME_BUILD:
- engine = frameBuildEngine;
- context = NATIVE_CONTEXT;
- break;
- case NATIVE_FRAME_CACHED:
- engine = frameCachedEngine;
- context = NATIVE_CONTEXT;
- break;
default:
throw new IllegalArgumentException("Unknown execution path " + executionPath);
}
@@ -682,24 +532,6 @@ public class SysSegmentsSqlBenchmark
blackhole.consume(runQuery(NATIVE_PROVIDER));
}
- @Benchmark
- public void queryNativeBatched(final Blackhole blackhole)
- {
- blackhole.consume(runQuery(NATIVE_BATCHED));
- }
-
- @Benchmark
- public void queryNativeFrameBuild(final Blackhole blackhole)
- {
- blackhole.consume(runQuery(NATIVE_FRAME_BUILD));
- }
-
- @Benchmark
- public void queryNativeFrameCached(final Blackhole blackhole)
- {
- blackhole.consume(runQuery(NATIVE_FRAME_CACHED));
- }
-
@Benchmark
public void queryExecutionOnly(final ExecutionState state, final Blackhole blackhole)
{
diff --git a/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java b/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java
index 3d792f62041..a355df24b04 100644
--- a/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java
+++ b/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java
@@ -67,8 +67,6 @@ import java.util.Set;
*/
public class BatchedInlineDataSource extends LeafDataSource
{
- static final int BATCH_SIZE = 1_024;
-
private final Iterable<Object[]> rows;
private final RowSignature signature;
@@ -270,7 +268,7 @@ public class BatchedInlineDataSource extends LeafDataSource
final CursorBuildSpec spec
)
{
- this.offset = new BatchOffset(rows, signature);
+ this.offset = new BatchOffset(rows, signature, spec.getQueryContext().getVectorSize());
this.selectorFactory = new BatchVectorColumnSelectorFactory(offset, signature, spec.getVirtualColumns());
}
@@ -302,7 +300,7 @@ public class BatchedInlineDataSource extends LeafDataSource
@Override
public int getMaxVectorSize()
{
- return BATCH_SIZE;
+ return offset.getMaxVectorSize();
}
@Override
@@ -316,6 +314,7 @@ public class BatchedInlineDataSource extends LeafDataSource
{
private final Iterable<Object[]> rows;
private final RowSignature signature;
+ private final int batchSize;
private final long[][] longColumns;
private final Object[][] objectColumns;
private final boolean[][] nullColumns;
@@ -324,20 +323,21 @@ public class BatchedInlineDataSource extends LeafDataSource
private int currentSize;
private int id;
- BatchOffset(final Iterable<Object[]> rows, final RowSignature signature)
+ BatchOffset(final Iterable<Object[]> rows, final RowSignature signature, final int batchSize)
{
this.rows = rows;
this.signature = signature;
+ this.batchSize = batchSize;
this.longColumns = new long[signature.size()][];
this.objectColumns = new Object[signature.size()][];
this.nullColumns = new boolean[signature.size()][];
for (int i = 0; i < signature.size(); i++) {
final ColumnType type = signature.getColumnType(i).orElse(null);
if (ColumnType.LONG.equals(type)) {
- longColumns[i] = new long[BATCH_SIZE];
- nullColumns[i] = new boolean[BATCH_SIZE];
+ longColumns[i] = new long[batchSize];
+ nullColumns[i] = new boolean[batchSize];
} else {
- objectColumns[i] = new Object[BATCH_SIZE];
+ objectColumns[i] = new Object[batchSize];
}
}
reset();
@@ -378,7 +378,7 @@ public class BatchedInlineDataSource extends LeafDataSource
private void loadBatch()
{
int rowNumber = 0;
- while (rowNumber < BATCH_SIZE && iterator.hasNext()) {
+ while (rowNumber < batchSize && iterator.hasNext()) {
final Object[] row = iterator.next();
for (int column = 0; column < signature.size(); column++) {
if (longColumns[column] != null) {
@@ -410,7 +410,7 @@ public class BatchedInlineDataSource extends LeafDataSource
@Override
public int getMaxVectorSize()
{
- return BATCH_SIZE;
+ return batchSize;
}
@Override
@@ -543,8 +543,8 @@ public class BatchedInlineDataSource extends LeafDataSource
private final ReadableVectorInspector inspector;
private final long[] longs;
private final boolean[] nulls;
- private final float[] floats = new float[BATCH_SIZE];
- private final double[] doubles = new double[BATCH_SIZE];
+ private final float[] floats;
+ private final double[] doubles;
private int floatId = ReadableVectorInspector.NULL_ID;
private int doubleId = ReadableVectorInspector.NULL_ID;
@@ -557,6 +557,8 @@ public class BatchedInlineDataSource extends LeafDataSource
this.inspector = inspector;
this.longs = longs;
this.nulls = nulls;
+ this.floats = new float[inspector.getMaxVectorSize()];
+ this.doubles = new double[inspector.getMaxVectorSize()];
}
@Override
@@ -642,7 +644,7 @@ public class BatchedInlineDataSource extends LeafDataSource
{
private final ReadableVectorInspector inspector;
private final Object[] objects;
- private final int[] ids = new int[BATCH_SIZE];
+ private final int[] ids;
private final Map<String, Integer> valueToId = new HashMap<>();
private final List<String> idToValue = new ArrayList<>();
private int vectorId = ReadableVectorInspector.NULL_ID;
@@ -651,6 +653,7 @@ public class BatchedInlineDataSource extends LeafDataSource
{
this.inspector = inspector;
this.objects = objects;
+ this.ids = new int[inspector.getMaxVectorSize()];
}
@Override
diff --git a/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java b/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java
index 7394437f392..ca48530198c 100644
--- a/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java
+++ b/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java
@@ -33,17 +33,19 @@ import org.junit.jupiter.api.Test;
import java.util.ArrayList;
import java.util.List;
+import java.util.Map;
public class BatchedInlineDataSourceTest
{
@Test
public void testVectorCursorLoadsMultipleBatches()
{
+ final int batchSize = QueryContexts.DEFAULT_VECTOR_SIZE;
final RowSignature signature = RowSignature.builder()
.add("value", ColumnType.LONG)
.build();
final List<Object[]> rows = new ArrayList<>();
- for (int i = 0; i < BatchedInlineDataSource.BATCH_SIZE + 1; i++) {
+ for (int i = 0; i < batchSize + 1; i++) {
rows.add(new Object[]{(long) i});
}
@@ -52,17 +54,17 @@ public class BatchedInlineDataSourceTest
Assertions.assertTrue(cursorHolder.canVectorize());
final VectorCursor cursor = cursorHolder.asVectorCursor();
final VectorValueSelector selector = cursor.getColumnSelectorFactory().makeValueSelector("value");
- Assertions.assertEquals(BatchedInlineDataSource.BATCH_SIZE, cursor.getCurrentVectorSize());
+ Assertions.assertEquals(batchSize, cursor.getCurrentVectorSize());
Assertions.assertEquals(0L, selector.getLongVector()[0]);
Assertions.assertEquals(
- BatchedInlineDataSource.BATCH_SIZE - 1L,
- selector.getLongVector()[BatchedInlineDataSource.BATCH_SIZE - 1]
+ batchSize - 1L,
+ selector.getLongVector()[batchSize - 1]
);
cursor.advance();
Assertions.assertEquals(1, cursor.getCurrentVectorSize());
- Assertions.assertEquals(BatchedInlineDataSource.BATCH_SIZE, selector.getLongVector()[0]);
+ Assertions.assertEquals(batchSize, selector.getLongVector()[0]);
cursor.advance();
Assertions.assertTrue(cursor.isDone());
}
@@ -84,6 +86,39 @@ public class BatchedInlineDataSourceTest
}
}
+ @Test
+ public void testUsesConfiguredVectorSize()
+ {
+ final RowSignature signature = RowSignature.builder()
+ .add("value", ColumnType.LONG)
+ .build();
+ final CursorFactory cursorFactory = makeCursorFactory(
+ new BatchedInlineDataSource(
+ List.of(
+ new Object[]{1L},
+ new Object[]{2L},
+ new Object[]{3L},
+ new Object[]{4L},
+ new Object[]{5L}
+ ),
+ signature
+ )
+ );
+ final CursorBuildSpec spec = CursorBuildSpec.builder()
+ .setQueryContext(
+ QueryContext.of(Map.of(QueryContexts.VECTOR_SIZE_KEY, 4))
+ )
+ .build();
+
+ try (final CursorHolder cursorHolder = cursorFactory.makeCursorHolder(spec)) {
+ final VectorCursor cursor = cursorHolder.asVectorCursor();
+ Assertions.assertEquals(4, cursor.getMaxVectorSize());
+ Assertions.assertEquals(4, cursor.getCurrentVectorSize());
+ cursor.advance();
+ Assertions.assertEquals(1, cursor.getCurrentVectorSize());
+ }
+ }
+
private static CursorFactory makeCursorFactory(final BatchedInlineDataSource dataSource)
{
final Segment segment = new BatchedInlineDataSource.Wrangler()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org