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 <[email protected]> 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: [email protected] For additional commands, e-mail: [email protected]
