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 41f48ed1287f1b13f4fb0328209d6d8ed00d5328 Author: Frank Chen <[email protected]> AuthorDate: Tue Sep 1 10:36:30 2026 +0800 perf(sql): batch authorized system table rows --- NATIVE_SYS_SEGMENTS_IMPLEMENTATION_PLAN.md | 70 ++ .../calcite/schema/SysSegmentsSqlBenchmark.java | 424 ++++++++++-- .../druid/query/BatchedInlineDataSource.java | 760 +++++++++++++++++++++ .../druid/query/BatchedInlineDataSourceTest.java | 95 +++ .../apache/druid/guice/SegmentWranglerModule.java | 2 + .../system/handler/SystemTableQueryClient.java | 2 +- .../system/handler/SystemTableQueryHandler.java | 18 +- .../system/table/SystemTableDataProvider.java | 15 + .../system/table/SystemTableQueryRequest.java | 36 + .../server/SpecificSegmentsQuerySegmentWalker.java | 2 + .../server/system/SystemTableQueryHandlerTest.java | 111 ++- .../calcite/schema/SegmentsTableDataProvider.java | 21 + .../schema/SegmentsTableDataProviderTest.java | 38 ++ 13 files changed, 1550 insertions(+), 44 deletions(-) diff --git a/NATIVE_SYS_SEGMENTS_IMPLEMENTATION_PLAN.md b/NATIVE_SYS_SEGMENTS_IMPLEMENTATION_PLAN.md new file mode 100644 index 00000000000..fcead688fa8 --- /dev/null +++ b/NATIVE_SYS_SEGMENTS_IMPLEMENTATION_PLAN.md @@ -0,0 +1,70 @@ +# Native `sys.segments` implementation plan + +## Base and scope + +- Base this work on Apache Druid PR 20183 head `8035122b2a3ab3f60be8aef4e8d164454c1fe266`. +- Add native-query support for `sys.segments`; do not change `sys.server_segments`. +- Keep the existing Bindable path and row contract backward compatible. +- Keep native execution opt-in through `useNativeQueryForSystemTables`. +- Execute native aggregations, expressions, sorting, and residual filters on the Broker. The provider supplies rows through the existing component-local Scan transport. + +## Design constraints + +1. The Broker that receives and plans the SQL query must execute the `sys.segments` provider locally. +2. Do not discover, contact, or fan out to other Broker nodes. +3. Register the provider only on Brokers. +4. Preserve the current `sys.segments` semantics, which combine Coordinator-derived segment status with the executing Broker's segment metadata cache. +5. Preserve datasource `READ` row authorization for the original user. +6. The native row values must match the descriptor's `RowSignature` exactly. Complex columns declared as strings must use the same JSON serialization as the Bindable path: + - `shard_spec` + - `dimensions` + - `metrics` + - `projections` + - `last_compaction_state` + +## Implementation steps + +1. Add a shared `SegmentsTableDescriptor` containing the table name, existing 20-column signature, Broker ownership, local-only routing, and datasource `READ` row authorization. +2. Add the smallest generic local-only routing capability to the native system-table framework. `SystemTableQueryClient` must invoke the raw local `SystemTableQueryHandler` directly for this mode without `SystemTableNodeLocator`, service discovery, HTTP, or retries to another Broker. +3. Add a Broker-local `SegmentsTableDataProvider` using `BrokerSegmentMetadataCache`, `MetadataSegmentView`, and `ObjectMapper`. +4. Preserve all existing row semantics: published/available merging, segment-ID deduplication, row-count precedence, replica and availability calculations, realtime/active/published/overshadowed flags, replication-factor fallback, and JSON serialization. +5. Advertise safe `datasource` equality/IN pushdown. Keep the original native filter in the Broker query as the correctness-preserving residual filter. +6. Share row construction and serialization with the Bindable implementation so native and Bindable output cannot drift. Avoid an unrelated refactor or a speculative abstraction. +7. Bind the provider only in `CliBroker`, and register the descriptor with the existing native system-table framework. +8. Make `SegmentsTable` implement `NativeSystemTable` and expose a `SystemTableDataSource("segments")`-backed `DruidTable`. +9. Update the native system-table documentation to list `sys.segments`, its local Broker source, supported datasource pushdown, and Broker-side execution behavior. +10. Add an optional generic provider capability that converts only framework-authorized rows into a query-local datasource. Implement `sys.segments` with a batched column-oriented cursor for `STRING` and `LONG` projections, retain row-cursor fallback for unsupported query shapes and types, and do not cache user-specific batches. + +## Verification + +1. Unit-test the descriptor signature and datasource authorization. +2. Unit-test local-only routing and prove it bypasses node discovery and remote clients. +3. Unit-test provider rows for published, unpublished, realtime, unavailable, overshadowed, duplicate, and row-count fallback cases, including complex-column JSON serialization and datasource pushdown. +4. Verify the Bindable path remains unchanged when native execution is disabled. +5. Verify COUPLED and DECOUPLED native planning. +6. Add embedded native SQL tests for representative native functionality, including distinct aggregation, grouping or expressions, nested aggregation, filters, and projections. +7. Include a multiple-Broker test or equivalent routing assertion proving a query is executed only by the SQL-receiving Broker and rows are not multiplied. +8. Run focused `server`, `sql`, `services`, and embedded tests with `-Pskip-static-checks -Dweb.console.skip=true -T1C`, followed by relevant static checks. +9. Review the complete diff against PR head `8035122b2a3ab3f60be8aef4e8d164454c1fe266` and ensure every changed line is required by this feature. +10. Benchmark the exact Web Console datasource-tab SQL over 500,000 segments, validate every result row against Bindable before measurement, and compare Bindable, legacy native-row, and provider-backed batched execution separately. + +## Latest 500K benchmark result + +JMH configuration: one fork, two 2-second warmups, three 2-second measurements, JDK 25.0.3. Lower is better. + +| Path | Average | +|---|---:| +| Bindable | 238.689 ms/op | +| Legacy native row | 335.129 ms/op | +| Benchmark-only batched wrapper | 187.089 ms/op | +| Provider-backed authorized batches | 172.031 ms/op | + +The provider-backed path was about 28% faster than Bindable and 49% faster than the legacy native-row path in this run. This is a local microbenchmark, not a cluster-level latency guarantee. + +## Explicitly rejected alternatives + +- Do not use `ALL_NODES` for the Broker role; it duplicates rows and introduces inconsistent cache snapshots. +- Do not move the provider to the Coordinator; that would lose Broker-cache-derived availability, replica, and row-count behavior. +- Do not query all Brokers and deduplicate afterward; it adds network and merge cost and makes conflicting cache values arbitrary. +- Do not cache columnar batches across queries; authorization and metadata snapshots are request-specific. +- Do not build frames per query for this path; measured frame construction cost exceeded both Bindable and row-native execution. 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 8a0022aab28..c244a3ea9f3 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,21 +30,38 @@ 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; import org.apache.druid.java.util.common.parsers.CloseableIterator; +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; +import org.apache.druid.segment.join.JoinableFactoryWrapper; import org.apache.druid.segment.metadata.CentralizedDatasourceSchemaConfig; import org.apache.druid.server.QueryLifecycleFactory; import org.apache.druid.server.QueryScheduler; @@ -55,12 +72,14 @@ import org.apache.druid.server.log.NoopRequestLogger; import org.apache.druid.server.metrics.NoopServiceEmitter; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthTestUtils; +import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.system.handler.SystemTableNodeLocator; import org.apache.druid.server.system.handler.SystemTableQueryClient; import org.apache.druid.server.system.handler.SystemTableQueryHandler; import org.apache.druid.server.system.table.SegmentsTableDescriptor; import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.server.system.table.SystemTableDescriptor; +import org.apache.druid.server.system.table.SystemTablePushdownFilter; import org.apache.druid.sql.SqlStatementFactory; import org.apache.druid.sql.calcite.planner.CalciteRulesManager; import org.apache.druid.sql.calcite.planner.CatalogResolver; @@ -80,6 +99,7 @@ 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; @@ -88,6 +108,7 @@ 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; @@ -98,6 +119,7 @@ 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; @@ -105,7 +127,18 @@ import java.util.concurrent.TimeUnit; /** Compares Bindable and native execution of the Web Console datasource-tab query over 500,000 segments. */ @State(Scope.Benchmark) -@Fork(value = 1, jvmArgsAppend = {"-Xmx12g"}) +@Fork( + value = 1, + jvmArgsAppend = { + "-Xmx12g", + "--add-exports=java.base/sun.nio.ch=ALL-UNNAMED", + "--add-opens=java.base/java.nio=ALL-UNNAMED", + "--add-opens=java.base/sun.nio.ch=ALL-UNNAMED", + "--add-opens=java.base/java.lang=ALL-UNNAMED", + "--add-opens=java.base/java.lang.reflect=ALL-UNNAMED", + "--add-opens=java.base/jdk.internal.ref=ALL-UNNAMED" + } +) @Warmup(iterations = 3, time = 3) @Measurement(iterations = 5, time = 3) @BenchmarkMode(Mode.AverageTime) @@ -114,6 +147,12 @@ public class SysSegmentsSqlBenchmark { private static final int NUM_SEGMENTS = 500_000; private static final int NUM_DATASOURCES = 1_000; + 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" @@ -152,7 +191,49 @@ public class SysSegmentsSqlBenchmark private final Closer closer = Closer.create(); private PlannerFactory plannerFactory; - private SqlEngine engine; + 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}) + private String executionPath; + private PreparedQuery preparedQuery; + + @Setup(Level.Invocation) + public void setup(final SysSegmentsSqlBenchmark benchmark) + { + preparedQuery = benchmark.prepare(executionPath); + } + + @TearDown(Level.Invocation) + public void tearDown() + { + preparedQuery.close(); + } + } + + private static class PreparedQuery implements AutoCloseable + { + private final DruidPlanner planner; + private final PlannerResult plannerResult; + + PreparedQuery(final DruidPlanner planner, final PlannerResult plannerResult) + { + this.planner = planner; + this.plannerResult = plannerResult; + } + + @Override + public void close() + { + planner.close(); + } + } private static class EmptyBrokerSegmentMetadataCache extends BrokerSegmentMetadataCache { @@ -175,6 +256,145 @@ 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 + { + private final SystemTableDataProvider delegate; + + RowOnlySystemTableDataProvider(final SystemTableDataProvider delegate) + { + this.delegate = delegate; + } + + @Override + public List<SystemTablePushdownFilter> getPushdownFilters() + { + return delegate.getPushdownFilters(); + } + + @Override + public Iterable<Object[]> getRows( + final List<DimFilter> filters, + final AuthenticationResult internalAuthenticationResult + ) + { + return delegate.getRows(filters, internalAuthenticationResult); + } + + @Override + public Iterable<Object[]> getRawRows( + final List<DimFilter> filters, + final AuthenticationResult internalAuthenticationResult + ) + { + return delegate.getRawRows(filters, internalAuthenticationResult); + } + + @Override + public Object[] projectRow(final Object[] row, final int[] projects) + { + return delegate.projectRow(row, projects); + } + } + @Setup(Level.Trial) public void setup() { @@ -210,40 +430,57 @@ public class SysSegmentsSqlBenchmark final QueryRunnerFactoryConglomerate conglomerate = QueryStackTests.createQueryRunnerFactoryConglomerate(closer); final SpecificSegmentsQuerySegmentWalker walker = closer.register( - SpecificSegmentsQuerySegmentWalker.createWalker(conglomerate) + SpecificSegmentsQuerySegmentWalker.createWalker( + QueryStackTests.injectorWithLookup(), + conglomerate, + new MapSegmentWrangler( + Map.of( + InlineDataSource.class, + new InlineSegmentWrangler(), + FrameBasedInlineDataSource.class, + new FrameBasedInlineSegmentWrangler(), + BatchedInlineDataSource.class, + new BatchedInlineDataSource.Wrangler() + ) + ), + new JoinableFactoryWrapper(QueryStackTests.makeJoinableFactoryFromDefault(null, null, null)), + QueryStackTests.DEFAULT_NOOP_SCHEDULER + ) ); - final SystemTableQueryHandler localQueryHandler = new SystemTableQueryHandler( + final SystemTableQueryHandler providerQueryHandler = new SystemTableQueryHandler( Map.<String, SystemTableDataProvider>of(descriptor.getTableName(), dataProvider), Map.<String, SystemTableDescriptor>of(descriptor.getTableName(), descriptor), new ScanQueryEngine(), AuthTestUtils.TEST_AUTHORIZER_MAPPER ); - final SystemTableQueryClient queryClient = new SystemTableQueryClient( - EasyMock.mock(SystemTableNodeLocator.class), - EasyMock.mock(DirectDruidClientFactory.class), - EasyMock.mock(QueryScheduler.class), + final SystemTableQueryHandler rowQueryHandler = new SystemTableQueryHandler( + Map.<String, SystemTableDataProvider>of( + descriptor.getTableName(), + new RowOnlySystemTableDataProvider(dataProvider) + ), + Map.<String, SystemTableDescriptor>of(descriptor.getTableName(), descriptor), + new ScanQueryEngine(), + AuthTestUtils.TEST_AUTHORIZER_MAPPER + ); + rowEngine = makeEngine(conglomerate, walker, descriptor, rowQueryHandler); + providerEngine = makeEngine(conglomerate, walker, descriptor, providerQueryHandler); + batchedEngine = makeEngine( + conglomerate, walker, - Map.of(descriptor.getTableName(), descriptor), - AuthTestUtils.TEST_AUTHORIZER_MAPPER, - localQueryHandler, - CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, - CalciteTests.mockCoordinatorNode() + descriptor, + new BatchedSystemTableQueryHandler(rowQueryHandler) ); - final QueryLifecycleFactory queryLifecycleFactory = new QueryLifecycleFactory( + frameBuildEngine = makeEngine( conglomerate, walker, - new DefaultGenericQueryMetricsFactory(), - NoopServiceEmitter.instance(), - NoopRequestLogger.instance(), - new AuthConfig(), - NoopPolicyEnforcer.instance(), - AuthTestUtils.TEST_AUTHORIZER_MAPPER, - new DefaultQueryConfig(Map.of()), - Map.<Class<? extends DataSource>, org.apache.druid.server.DataSourceQueryHandler>of( - SystemTableDataSource.class, - queryClient - ), - null + descriptor, + new FrameSystemTableQueryHandler(rowQueryHandler, false) + ); + frameCachedEngine = makeEngine( + conglomerate, + walker, + descriptor, + new FrameSystemTableQueryHandler(rowQueryHandler, true) ); final PlannerConfig plannerConfig = new PlannerConfig(); @@ -282,7 +519,6 @@ public class SysSegmentsSqlBenchmark plannerConfig ); - engine = new NativeSqlEngine(queryLifecycleFactory, CalciteTests.getJsonMapper(), (SqlStatementFactory) null); plannerFactory = new PlannerFactory( schemaProvider, CalciteTests.createOperatorTable(), @@ -299,13 +535,53 @@ public class SysSegmentsSqlBenchmark new DruidHookDispatcher() ); - final List<Object[]> bindableResults = runQuery(BINDABLE_CONTEXT); - final List<Object[]> nativeResults = runQuery(NATIVE_CONTEXT); - if (bindableResults.size() != NUM_DATASOURCES || !rowsEqual(bindableResults, nativeResults)) { - throw new IllegalStateException("Bindable and native benchmark results do not match"); + final List<Object[]> bindableResults = runQuery(BINDABLE); + for (final String executionPath : + List.of(NATIVE_ROW, NATIVE_PROVIDER, NATIVE_BATCHED, NATIVE_FRAME_BUILD, NATIVE_FRAME_CACHED)) { + 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); + } } } + private static SqlEngine makeEngine( + final QueryRunnerFactoryConglomerate conglomerate, + final SpecificSegmentsQuerySegmentWalker walker, + final SystemTableDescriptor descriptor, + final SystemTableQueryHandler localQueryHandler + ) + { + final SystemTableQueryClient queryClient = new SystemTableQueryClient( + EasyMock.mock(SystemTableNodeLocator.class), + EasyMock.mock(DirectDruidClientFactory.class), + EasyMock.mock(QueryScheduler.class), + walker, + Map.of(descriptor.getTableName(), descriptor), + AuthTestUtils.TEST_AUTHORIZER_MAPPER, + localQueryHandler, + CalciteTests.TEST_AUTHENTICATOR_ESCALATOR, + CalciteTests.mockCoordinatorNode() + ); + final QueryLifecycleFactory queryLifecycleFactory = new QueryLifecycleFactory( + conglomerate, + walker, + new DefaultGenericQueryMetricsFactory(), + NoopServiceEmitter.instance(), + NoopRequestLogger.instance(), + new AuthConfig(), + NoopPolicyEnforcer.instance(), + AuthTestUtils.TEST_AUTHORIZER_MAPPER, + new DefaultQueryConfig(Map.of()), + Map.<Class<? extends DataSource>, org.apache.druid.server.DataSourceQueryHandler>of( + SystemTableDataSource.class, + queryClient + ), + null + ); + return new NativeSqlEngine(queryLifecycleFactory, CalciteTests.getJsonMapper(), (SqlStatementFactory) null); + } + private static List<SegmentStatusInCluster> buildSegments() { final List<SegmentStatusInCluster> segments = new ArrayList<>(NUM_SEGMENTS); @@ -338,25 +614,97 @@ public class SysSegmentsSqlBenchmark return true; } - private List<Object[]> runQuery(final Map<String, Object> context) + private PreparedQuery prepare(final String executionPath) + { + final SqlEngine engine; + final Map<String, Object> context; + switch (executionPath) { + case BINDABLE: + engine = rowEngine; + context = BINDABLE_CONTEXT; + break; + case NATIVE_ROW: + engine = rowEngine; + context = NATIVE_CONTEXT; + break; + case NATIVE_PROVIDER: + 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); + } + + final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, SQL, context); + try { + return new PreparedQuery(planner, planner.plan()); + } + catch (RuntimeException | Error t) { + planner.close(); + throw t; + } + } + + private List<Object[]> runQuery(final String executionPath) { - try (final DruidPlanner planner = plannerFactory.createPlannerForTesting(engine, SQL, context)) { - final PlannerResult plannerResult = planner.plan(); - final Sequence<Object[]> resultSequence = plannerResult.run().getResults(); - return resultSequence.toList(); + try (final PreparedQuery preparedQuery = prepare(executionPath)) { + return preparedQuery.plannerResult.run().getResults().toList(); } } @Benchmark public void queryBindable(final Blackhole blackhole) { - blackhole.consume(runQuery(BINDABLE_CONTEXT)); + blackhole.consume(runQuery(BINDABLE)); } @Benchmark public void queryNative(final Blackhole blackhole) { - blackhole.consume(runQuery(NATIVE_CONTEXT)); + blackhole.consume(runQuery(NATIVE_ROW)); + } + + @Benchmark + public void queryNativeProvider(final Blackhole blackhole) + { + 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) + { + final Sequence<Object[]> resultSequence = state.preparedQuery.plannerResult.run().getResults(); + blackhole.consume(resultSequence.toList()); } @TearDown(Level.Trial) diff --git a/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java b/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java new file mode 100644 index 00000000000..3d792f62041 --- /dev/null +++ b/processing/src/main/java/org/apache/druid/query/BatchedInlineDataSource.java @@ -0,0 +1,760 @@ +/* + * 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.druid.query; + +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.java.util.common.guava.Sequences; +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.segment.Cursor; +import org.apache.druid.segment.CursorBuildSpec; +import org.apache.druid.segment.CursorFactory; +import org.apache.druid.segment.CursorHolder; +import org.apache.druid.segment.IdLookup; +import org.apache.druid.segment.ResidentCursorFactory; +import org.apache.druid.segment.RowAdapter; +import org.apache.druid.segment.RowBasedCursorFactory; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.SegmentWrangler; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.column.ColumnCapabilitiesImpl; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.column.ValueType; +import org.apache.druid.segment.vector.MultiValueDimensionVectorSelector; +import org.apache.druid.segment.vector.NilVectorSelector; +import org.apache.druid.segment.vector.ReadableVectorInspector; +import org.apache.druid.segment.vector.ReadableVectorOffset; +import org.apache.druid.segment.vector.SingleValueDimensionVectorSelector; +import org.apache.druid.segment.vector.VectorColumnSelectorFactory; +import org.apache.druid.segment.vector.VectorCursor; +import org.apache.druid.segment.vector.VectorObjectSelector; +import org.apache.druid.segment.vector.VectorValueSelector; +import org.apache.druid.timeline.SegmentId; +import org.joda.time.Interval; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Internal datasource that transposes authorized rows into reusable, query-local column batches. + * It deliberately has no persistent cache. Unsupported cursor shapes fall back to row-based processing. + */ +public class BatchedInlineDataSource extends LeafDataSource +{ + static final int BATCH_SIZE = 1_024; + + private final Iterable<Object[]> rows; + private final RowSignature signature; + + public BatchedInlineDataSource(final Iterable<Object[]> rows, final RowSignature signature) + { + this.rows = rows; + this.signature = signature; + } + + public Iterable<Object[]> getRows() + { + return rows; + } + + public RowSignature getRowSignature() + { + return signature; + } + + @Override + public Set<String> getTableNames() + { + return Set.of(); + } + + @Override + public boolean isCacheable(final boolean isBroker) + { + return false; + } + + @Override + public boolean isGlobal() + { + return true; + } + + @Override + public boolean isProcessable() + { + return true; + } + + @Override + public byte[] getCacheKey() + { + return null; + } + + public static class Wrangler implements SegmentWrangler + { + @Override + public Iterable<Segment> getSegmentsForIntervals( + final DataSource dataSource, + final Iterable<Interval> intervals + ) + { + return List.of(new BatchedInlineSegment((BatchedInlineDataSource) dataSource)); + } + } + + private static class BatchedInlineSegment implements Segment + { + private final BatchedInlineDataSource dataSource; + + BatchedInlineSegment(final BatchedInlineDataSource dataSource) + { + this.dataSource = dataSource; + } + + @Nullable + @Override + public SegmentId getId() + { + return null; + } + + @Override + public Interval getDataInterval() + { + return Intervals.ETERNITY; + } + + @Nullable + @Override + public <T> T as(@Nonnull final Class<T> clazz) + { + if (CursorFactory.class.equals(clazz)) { + return clazz.cast(new BatchedCursorFactory(dataSource.rows, dataSource.signature)); + } + return null; + } + + @Override + public void close() + { + // Nothing to close. + } + } + + private static class BatchedCursorFactory implements ResidentCursorFactory + { + private final Iterable<Object[]> rows; + private final RowSignature signature; + private final RowBasedCursorFactory<Object[]> rowCursorFactory; + + BatchedCursorFactory(final Iterable<Object[]> rows, final RowSignature signature) + { + this.rows = rows; + this.signature = signature; + final RowAdapter<Object[]> rowAdapter = column -> { + final int columnNumber = signature.indexOf(column); + return row -> columnNumber < 0 ? null : row[columnNumber]; + }; + this.rowCursorFactory = new RowBasedCursorFactory<>(Sequences.simple(rows), rowAdapter, signature); + } + + @Override + public CursorHolder makeCursorHolder(final CursorBuildSpec spec) + { + final CursorHolder rowCursorHolder = rowCursorFactory.makeCursorHolder(spec); + return new CursorHolder() + { + @Nullable + @Override + public Cursor asCursor() + { + return rowCursorHolder.asCursor(); + } + + @Override + public boolean canVectorize() + { + if (spec.getFilter() != null || !supportsVectorTypes(signature)) { + return false; + } + final VirtualColumns virtualColumns = spec.getVirtualColumns(); + final ColumnCapabilitiesInspector inspector = new ColumnCapabilitiesInspector(signature); + if (!virtualColumns.isEmpty() && !virtualColumns.canVectorize(inspector)) { + return false; + } + final List<AggregatorFactory> aggregators = spec.getAggregators(); + if (aggregators != null) { + for (final AggregatorFactory aggregator : aggregators) { + if (!aggregator.canVectorize(virtualColumns.wrapInspector(inspector))) { + return false; + } + } + } + return true; + } + + @Nullable + @Override + public VectorCursor asVectorCursor() + { + if (!canVectorize()) { + throw new ISE("Batched inline cursor cannot vectorize this query"); + } + return new BatchVectorCursor(rows, signature, spec); + } + + @Override + public List<OrderBy> getOrdering() + { + return rowCursorHolder.getOrdering(); + } + + @Override + public void close() + { + rowCursorHolder.close(); + } + }; + } + + @Override + public RowSignature getRowSignature() + { + return signature; + } + + @Nullable + @Override + public ColumnCapabilities getColumnCapabilities(final String column) + { + return capabilities(signature, column); + } + } + + private static class BatchVectorCursor implements VectorCursor + { + private final BatchOffset offset; + private final VectorColumnSelectorFactory selectorFactory; + + BatchVectorCursor( + final Iterable<Object[]> rows, + final RowSignature signature, + final CursorBuildSpec spec + ) + { + this.offset = new BatchOffset(rows, signature); + this.selectorFactory = new BatchVectorColumnSelectorFactory(offset, signature, spec.getVirtualColumns()); + } + + @Override + public VectorColumnSelectorFactory getColumnSelectorFactory() + { + return selectorFactory; + } + + @Override + public void advance() + { + offset.advance(); + BaseQuery.checkInterrupted(); + } + + @Override + public boolean isDone() + { + return offset.isDone(); + } + + @Override + public void reset() + { + offset.reset(); + } + + @Override + public int getMaxVectorSize() + { + return BATCH_SIZE; + } + + @Override + public int getCurrentVectorSize() + { + return offset.getCurrentVectorSize(); + } + } + + private static class BatchOffset implements ReadableVectorOffset + { + private final Iterable<Object[]> rows; + private final RowSignature signature; + private final long[][] longColumns; + private final Object[][] objectColumns; + private final boolean[][] nullColumns; + + private Iterator<Object[]> iterator; + private int currentSize; + private int id; + + BatchOffset(final Iterable<Object[]> rows, final RowSignature signature) + { + this.rows = rows; + this.signature = signature; + 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]; + } else { + objectColumns[i] = new Object[BATCH_SIZE]; + } + } + reset(); + } + + long[] getLongColumn(final int column) + { + return longColumns[column]; + } + + Object[] getObjectColumn(final int column) + { + return objectColumns[column]; + } + + boolean[] getNullColumn(final int column) + { + return nullColumns[column]; + } + + void advance() + { + loadBatch(); + } + + boolean isDone() + { + return currentSize == 0; + } + + void reset() + { + iterator = rows.iterator(); + id = 0; + loadBatch(); + } + + private void loadBatch() + { + int rowNumber = 0; + while (rowNumber < BATCH_SIZE && iterator.hasNext()) { + final Object[] row = iterator.next(); + for (int column = 0; column < signature.size(); column++) { + if (longColumns[column] != null) { + final Object value = row[column]; + nullColumns[column][rowNumber] = value == null; + longColumns[column][rowNumber] = value == null ? 0L : ((Number) value).longValue(); + } else { + objectColumns[column][rowNumber] = row[column]; + } + } + rowNumber++; + } + currentSize = rowNumber; + id++; + } + + @Override + public int getId() + { + return id; + } + + @Override + public boolean isContiguous() + { + return true; + } + + @Override + public int getMaxVectorSize() + { + return BATCH_SIZE; + } + + @Override + public int getCurrentVectorSize() + { + return currentSize; + } + + @Override + public int getStartOffset() + { + return 0; + } + + @Override + public int[] getOffsets() + { + throw new UnsupportedOperationException("contiguous batch"); + } + } + + private static class BatchVectorColumnSelectorFactory implements VectorColumnSelectorFactory + { + private final BatchOffset offset; + private final RowSignature signature; + private final VirtualColumns virtualColumns; + private final Map<DimensionSpec, SingleValueDimensionVectorSelector> dimensionSelectors = new HashMap<>(); + private final Map<String, VectorValueSelector> valueSelectors = new HashMap<>(); + private final Map<String, VectorObjectSelector> objectSelectors = new HashMap<>(); + + BatchVectorColumnSelectorFactory( + final BatchOffset offset, + final RowSignature signature, + final VirtualColumns virtualColumns + ) + { + this.offset = offset; + this.signature = signature; + this.virtualColumns = virtualColumns; + } + + @Override + public ReadableVectorInspector getReadableVectorInspector() + { + return offset; + } + + @Override + public SingleValueDimensionVectorSelector makeSingleValueDimensionSelector(final DimensionSpec dimensionSpec) + { + SingleValueDimensionVectorSelector selector = dimensionSelectors.get(dimensionSpec); + if (selector == null) { + if (virtualColumns.exists(dimensionSpec.getDimension())) { + selector = virtualColumns.makeSingleValueDimensionVectorSelector( + dimensionSpec, + this, + null, + offset + ); + } else { + final int column = signature.indexOf(dimensionSpec.getDimension()); + selector = column < 0 + ? NilVectorSelector.create(offset) + : dimensionSpec.decorate(new BatchStringDimensionSelector(offset, offset.getObjectColumn(column))); + } + dimensionSelectors.put(dimensionSpec, selector); + } + return selector; + } + + @Override + public MultiValueDimensionVectorSelector makeMultiValueDimensionSelector(final DimensionSpec dimensionSpec) + { + throw new UnsupportedOperationException("Batched inline cursors do not support multi-value dimensions"); + } + + @Override + public VectorValueSelector makeValueSelector(final String column) + { + VectorValueSelector selector = valueSelectors.get(column); + if (selector == null) { + if (virtualColumns.exists(column)) { + selector = virtualColumns.makeVectorValueSelector(column, this, null, offset); + } else { + final int columnNumber = signature.indexOf(column); + selector = columnNumber < 0 + ? NilVectorSelector.create(offset) + : new BatchLongVectorValueSelector( + offset, + offset.getLongColumn(columnNumber), + offset.getNullColumn(columnNumber) + ); + } + valueSelectors.put(column, selector); + } + return selector; + } + + @Override + public VectorObjectSelector makeObjectSelector(final String column) + { + VectorObjectSelector selector = objectSelectors.get(column); + if (selector == null) { + if (virtualColumns.exists(column)) { + selector = virtualColumns.makeVectorObjectSelector(column, this, null, offset); + } else { + final int columnNumber = signature.indexOf(column); + selector = columnNumber < 0 + ? NilVectorSelector.create(offset) + : new BatchObjectVectorSelector(offset, offset.getObjectColumn(columnNumber)); + } + objectSelectors.put(column, selector); + } + return selector; + } + + @Nullable + @Override + public ColumnCapabilities getColumnCapabilities(final String column) + { + return virtualColumns.getColumnCapabilitiesWithFallback( + new ColumnCapabilitiesInspector(signature), + column + ); + } + } + + private static class BatchLongVectorValueSelector implements VectorValueSelector + { + 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 int floatId = ReadableVectorInspector.NULL_ID; + private int doubleId = ReadableVectorInspector.NULL_ID; + + BatchLongVectorValueSelector( + final ReadableVectorInspector inspector, + final long[] longs, + final boolean[] nulls + ) + { + this.inspector = inspector; + this.longs = longs; + this.nulls = nulls; + } + + @Override + public long[] getLongVector() + { + return longs; + } + + @Override + public float[] getFloatVector() + { + if (floatId != inspector.getId()) { + for (int i = 0; i < inspector.getCurrentVectorSize(); i++) { + floats[i] = longs[i]; + } + floatId = inspector.getId(); + } + return floats; + } + + @Override + public double[] getDoubleVector() + { + if (doubleId != inspector.getId()) { + for (int i = 0; i < inspector.getCurrentVectorSize(); i++) { + doubles[i] = longs[i]; + } + doubleId = inspector.getId(); + } + return doubles; + } + + @Override + public boolean[] getNullVector() + { + return nulls; + } + + @Override + public int getMaxVectorSize() + { + return inspector.getMaxVectorSize(); + } + + @Override + public int getCurrentVectorSize() + { + return inspector.getCurrentVectorSize(); + } + } + + private static class BatchObjectVectorSelector implements VectorObjectSelector + { + private final ReadableVectorInspector inspector; + private final Object[] objects; + + BatchObjectVectorSelector(final ReadableVectorInspector inspector, final Object[] objects) + { + this.inspector = inspector; + this.objects = objects; + } + + @Override + public Object[] getObjectVector() + { + return objects; + } + + @Override + public int getMaxVectorSize() + { + return inspector.getMaxVectorSize(); + } + + @Override + public int getCurrentVectorSize() + { + return inspector.getCurrentVectorSize(); + } + } + + private static class BatchStringDimensionSelector implements SingleValueDimensionVectorSelector + { + private final ReadableVectorInspector inspector; + private final Object[] objects; + private final int[] ids = new int[BATCH_SIZE]; + private final Map<String, Integer> valueToId = new HashMap<>(); + private final List<String> idToValue = new ArrayList<>(); + private int vectorId = ReadableVectorInspector.NULL_ID; + + BatchStringDimensionSelector(final ReadableVectorInspector inspector, final Object[] objects) + { + this.inspector = inspector; + this.objects = objects; + } + + @Override + public int[] getRowVector() + { + if (vectorId != inspector.getId()) { + for (int i = 0; i < inspector.getCurrentVectorSize(); i++) { + final String value = (String) objects[i]; + ids[i] = valueToId.computeIfAbsent( + value, + ignored -> { + idToValue.add(value); + return idToValue.size() - 1; + } + ); + } + vectorId = inspector.getId(); + } + return ids; + } + + @Override + public int getValueCardinality() + { + return CARDINALITY_UNKNOWN; + } + + @Nullable + @Override + public String lookupName(final int id) + { + return idToValue.get(id); + } + + @Override + public boolean nameLookupPossibleInAdvance() + { + return false; + } + + @Nullable + @Override + public IdLookup idLookup() + { + return null; + } + + @Override + public int getMaxVectorSize() + { + return inspector.getMaxVectorSize(); + } + + @Override + public int getCurrentVectorSize() + { + return inspector.getCurrentVectorSize(); + } + } + + private static class ColumnCapabilitiesInspector implements org.apache.druid.segment.ColumnInspector + { + private final RowSignature signature; + + ColumnCapabilitiesInspector(final RowSignature signature) + { + this.signature = signature; + } + + @Nullable + @Override + public ColumnCapabilities getColumnCapabilities(final String column) + { + return capabilities(signature, column); + } + } + + private static boolean supportsVectorTypes(final RowSignature signature) + { + for (int i = 0; i < signature.size(); i++) { + final ColumnType type = signature.getColumnType(i).orElse(null); + if (!ColumnType.LONG.equals(type) && !ColumnType.STRING.equals(type)) { + return false; + } + } + return true; + } + + @Nullable + private static ColumnCapabilities capabilities(final RowSignature signature, final String column) + { + final ColumnType type = signature.getColumnType(column).orElse(null); + if (type == null) { + return null; + } + if (type.is(ValueType.STRING)) { + return ColumnCapabilitiesImpl.createDefault() + .setType(ColumnType.STRING) + .setDictionaryEncoded(true) + .setDictionaryValuesSorted(false) + .setDictionaryValuesUnique(true) + .setHasMultipleValues(false) + .setHasNulls(true); + } + return ColumnCapabilitiesImpl.createSimpleNumericColumnCapabilities(type).setHasNulls(true); + } +} diff --git a/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java b/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java new file mode 100644 index 00000000000..7394437f392 --- /dev/null +++ b/processing/src/test/java/org/apache/druid/query/BatchedInlineDataSourceTest.java @@ -0,0 +1,95 @@ +/* + * 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.druid.query; + +import org.apache.druid.segment.CursorBuildSpec; +import org.apache.druid.segment.CursorFactory; +import org.apache.druid.segment.CursorHolder; +import org.apache.druid.segment.Segment; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; +import org.apache.druid.segment.vector.VectorCursor; +import org.apache.druid.segment.vector.VectorValueSelector; +import org.joda.time.Interval; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +public class BatchedInlineDataSourceTest +{ + @Test + public void testVectorCursorLoadsMultipleBatches() + { + 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++) { + rows.add(new Object[]{(long) i}); + } + + final CursorFactory cursorFactory = makeCursorFactory(new BatchedInlineDataSource(rows, signature)); + try (final CursorHolder cursorHolder = cursorFactory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) { + 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(0L, selector.getLongVector()[0]); + Assertions.assertEquals( + BatchedInlineDataSource.BATCH_SIZE - 1L, + selector.getLongVector()[BatchedInlineDataSource.BATCH_SIZE - 1] + ); + + cursor.advance(); + + Assertions.assertEquals(1, cursor.getCurrentVectorSize()); + Assertions.assertEquals(BatchedInlineDataSource.BATCH_SIZE, selector.getLongVector()[0]); + cursor.advance(); + Assertions.assertTrue(cursor.isDone()); + } + } + + @Test + public void testUnsupportedTypeUsesRowCursor() + { + final RowSignature signature = RowSignature.builder() + .add("value", ColumnType.DOUBLE) + .build(); + final CursorFactory cursorFactory = makeCursorFactory( + new BatchedInlineDataSource(List.<Object[]>of(new Object[]{1.5D}), signature) + ); + + try (final CursorHolder cursorHolder = cursorFactory.makeCursorHolder(CursorBuildSpec.FULL_SCAN)) { + Assertions.assertFalse(cursorHolder.canVectorize()); + Assertions.assertNotNull(cursorHolder.asCursor()); + } + } + + private static CursorFactory makeCursorFactory(final BatchedInlineDataSource dataSource) + { + final Segment segment = new BatchedInlineDataSource.Wrangler() + .getSegmentsForIntervals(dataSource, List.<Interval>of()) + .iterator() + .next(); + return segment.as(CursorFactory.class); + } +} diff --git a/server/src/main/java/org/apache/druid/guice/SegmentWranglerModule.java b/server/src/main/java/org/apache/druid/guice/SegmentWranglerModule.java index a6cc1ce987b..a13deeee4d4 100644 --- a/server/src/main/java/org/apache/druid/guice/SegmentWranglerModule.java +++ b/server/src/main/java/org/apache/druid/guice/SegmentWranglerModule.java @@ -25,6 +25,7 @@ import com.google.inject.Binder; import com.google.inject.Module; import com.google.inject.Scopes; import com.google.inject.multibindings.MapBinder; +import org.apache.druid.query.BatchedInlineDataSource; import org.apache.druid.query.DataSource; import org.apache.druid.query.FrameBasedInlineDataSource; import org.apache.druid.query.InlineDataSource; @@ -48,6 +49,7 @@ public class SegmentWranglerModule implements Module @VisibleForTesting static final Map<Class<? extends DataSource>, Class<? extends SegmentWrangler>> WRANGLER_MAPPINGS = ImmutableMap.of( + BatchedInlineDataSource.class, BatchedInlineDataSource.Wrangler.class, InlineDataSource.class, InlineSegmentWrangler.class, FrameBasedInlineDataSource.class, FrameBasedInlineSegmentWrangler.class, LookupDataSource.class, LookupSegmentWrangler.class diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java index 44e2a0db525..0bc48b7aa9f 100644 --- a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryClient.java @@ -202,7 +202,7 @@ public class SystemTableQueryClient implements DataSourceQueryHandler return dataSource.withChildren(resolvedChildren); } - private InlineDataSource resolveSystemTableDataSource( + private DataSource resolveSystemTableDataSource( final SystemTableDataSource dataSource, final Query<?> owningQuery, final AuthenticationResult authenticationResult, diff --git a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java index 311621999c3..a1862c479da 100644 --- a/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java +++ b/server/src/main/java/org/apache/druid/server/system/handler/SystemTableQueryHandler.java @@ -26,6 +26,7 @@ import org.apache.druid.java.util.common.IAE; import org.apache.druid.java.util.common.ISE; import org.apache.druid.java.util.common.JodaUtils; import org.apache.druid.java.util.common.guava.Sequence; +import org.apache.druid.query.DataSource; import org.apache.druid.query.Druids; import org.apache.druid.query.InlineDataSource; import org.apache.druid.query.Query; @@ -44,6 +45,7 @@ import org.apache.druid.server.security.AuthorizerMapper; import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.server.system.table.SystemTableDescriptor; import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.apache.druid.server.system.table.SystemTableQueryRequest; import java.util.List; import java.util.Map; @@ -120,7 +122,7 @@ public class SystemTableQueryHandler implements DataSourceQueryHandler } /** Resolves a local system table directly to lazily projected, user-authorized inline rows. */ - public InlineDataSource resolveDataSource( + public DataSource resolveDataSource( final ScanQuery query, final AuthenticationResult requestAuthenticationResult ) @@ -141,15 +143,27 @@ public class SystemTableQueryHandler implements DataSourceQueryHandler signatureBuilder.add(column, descriptor.getRowSignature().getColumnType(columnNumber).orElse(null)); } + final RowSignature projectedSignature = signatureBuilder.build(); final Iterable<Object[]> authorizedRows = getAuthorizedRows( dataSupplier, descriptor, query, requestAuthenticationResult ); + final DataSource authorizedDataSource = dataSupplier.getAuthorizedDataSource( + new SystemTableQueryRequest( + columns, + projectedSignature + ), + authorizedRows + ).orElse(null); + if (authorizedDataSource != null) { + return authorizedDataSource; + } + return InlineDataSource.fromIterable( Iterables.transform(authorizedRows, row -> dataSupplier.projectRow(row, projects)), - signatureBuilder.build() + projectedSignature ); } diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java index a4d6aa18943..cc235f17e52 100644 --- a/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableDataProvider.java @@ -20,12 +20,14 @@ package org.apache.druid.server.system.table; import jakarta.validation.constraints.NotNull; +import org.apache.druid.query.DataSource; import org.apache.druid.query.filter.DimFilter; import org.apache.druid.server.security.AuthenticationResult; import javax.annotation.Nullable; import java.util.Collections; import java.util.List; +import java.util.Optional; /** * Supplies storage-prefiltered rows authorized for the internal caller of one native system table. @@ -39,6 +41,19 @@ public interface SystemTableDataProvider return Collections.emptyList(); } + /** + * Returns a query-local datasource over the framework-authorized raw rows, or empty to use the inline-row fallback. + * Implementations must derive the returned datasource solely from {@code authorizedRows}; authorization and storage + * filter pushdown have already been applied by the framework. + */ + default Optional<DataSource> getAuthorizedDataSource( + final SystemTableQueryRequest request, + final Iterable<Object[]> authorizedRows + ) + { + return Optional.empty(); + } + Iterable<Object[]> getRows( @NotNull List<DimFilter> filters, AuthenticationResult internalAuthenticationResult diff --git a/server/src/main/java/org/apache/druid/server/system/table/SystemTableQueryRequest.java b/server/src/main/java/org/apache/druid/server/system/table/SystemTableQueryRequest.java new file mode 100644 index 00000000000..17a717018cb --- /dev/null +++ b/server/src/main/java/org/apache/druid/server/system/table/SystemTableQueryRequest.java @@ -0,0 +1,36 @@ +/* + * 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.druid.server.system.table; + +import org.apache.druid.segment.column.RowSignature; + +import java.util.List; + +/** Query-local projection requested from a native system-table provider. */ +public record SystemTableQueryRequest( + List<String> columns, + RowSignature rowSignature +) +{ + public SystemTableQueryRequest + { + columns = List.copyOf(columns); + } +} diff --git a/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java b/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java index 9eda84d52a2..f238bc1228f 100644 --- a/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java +++ b/server/src/test/java/org/apache/druid/server/SpecificSegmentsQuerySegmentWalker.java @@ -26,6 +26,7 @@ import com.google.common.collect.Ordering; import com.google.common.io.Closeables; import com.google.inject.Injector; import org.apache.druid.error.DruidException; +import org.apache.druid.query.BatchedInlineDataSource; import org.apache.druid.query.DataSource; import org.apache.druid.query.FrameBasedInlineDataSource; import org.apache.druid.query.InlineDataSource; @@ -170,6 +171,7 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C conglomerate, new MapSegmentWrangler( ImmutableMap.<Class<? extends DataSource>, SegmentWrangler>builder() + .put(BatchedInlineDataSource.class, new BatchedInlineDataSource.Wrangler()) .put(InlineDataSource.class, new InlineSegmentWrangler()) .put(FrameBasedInlineDataSource.class, new FrameBasedInlineSegmentWrangler()) .put( diff --git a/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java b/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java index 866f2c86964..084dc0958fd 100644 --- a/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java +++ b/server/src/test/java/org/apache/druid/server/system/SystemTableQueryHandlerTest.java @@ -20,6 +20,8 @@ package org.apache.druid.server.system; import org.apache.druid.discovery.NodeRole; +import org.apache.druid.query.BatchedInlineDataSource; +import org.apache.druid.query.DataSource; import org.apache.druid.query.Druids; import org.apache.druid.query.InlineDataSource; import org.apache.druid.query.QueryPlus; @@ -39,14 +41,17 @@ import org.apache.druid.server.system.handler.SystemTableQueryHandler; import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.server.system.table.SystemTableDescriptor; import org.apache.druid.server.system.table.SystemTableRowAuthorizer; +import org.apache.druid.server.system.table.SystemTableQueryRequest; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; public class SystemTableQueryHandlerTest { @@ -227,9 +232,12 @@ public class SystemTableQueryHandlerTest .columns(List.of("duration")) .build(); - final InlineDataSource inlineDataSource = handler.resolveDataSource( - query, - new AuthenticationResult("alice", "allow", "external", null) + final InlineDataSource inlineDataSource = Assertions.assertInstanceOf( + InlineDataSource.class, + handler.resolveDataSource( + query, + new AuthenticationResult("alice", "allow", "external", null) + ) ); Assertions.assertEquals(List.of("duration"), inlineDataSource.getRowSignature().getColumnNames()); @@ -242,4 +250,101 @@ public class SystemTableQueryHandlerTest Assertions.assertArrayEquals(new Object[]{20L}, rows.get(0)); } + @Test + public void testProviderDataSourceReceivesOnlyAuthorizedRows() + { + final AtomicReference<List<Object[]>> rowsReceivedByProvider = new AtomicReference<>(); + final AtomicReference<SystemTableQueryRequest> requestReceivedByProvider = new AtomicReference<>(); + final SystemTableDataProvider supplier = new SystemTableDataProvider() + { + @Override + public Iterable<Object[]> getRows( + final List<DimFilter> filters, + final AuthenticationResult authenticationResult + ) + { + return List.of( + new Object[]{"task-a", 10L}, + new Object[]{"task-b", 20L} + ); + } + + @Override + public Optional<DataSource> getAuthorizedDataSource( + final SystemTableQueryRequest request, + final Iterable<Object[]> authorizedRows + ) + { + final List<Object[]> rows = java.util.stream.StreamSupport.stream( + authorizedRows.spliterator(), + false + ).toList(); + requestReceivedByProvider.set(request); + rowsReceivedByProvider.set(rows); + return Optional.of( + new BatchedInlineDataSource( + rows.stream().map(row -> new Object[]{row[1]}).toList(), + request.rowSignature() + ) + ); + } + }; + final SystemTableDescriptor descriptor = new SystemTableDescriptor() + { + @Override + public String getTableName() + { + return "test"; + } + + @Override + public Set<NodeRole> getNodeRoles() + { + return Set.of(); + } + + @Override + public RowSignature getRowSignature() + { + return ROW_SIGNATURE; + } + + @Override + public SystemTableRowAuthorizer getRowAuthorizer() + { + return (rows, authenticationResult, authorizerMapper) -> + () -> java.util.stream.StreamSupport.stream(rows.spliterator(), false) + .filter(row -> "task-b".equals(row[0])) + .iterator(); + } + }; + final SystemTableQueryHandler handler = new SystemTableQueryHandler( + Map.of("test", supplier), + Map.of("test", descriptor), + new ScanQueryEngine(), + new AuthorizerMapper(Map.of()) + ); + final AuthenticationResult authenticationResult = new AuthenticationResult( + "alice", + "allow", + "external", + null + ); + final ScanQuery query = Druids.newScanQueryBuilder() + .dataSource(new SystemTableDataSource("test")) + .eternityInterval() + .columns(List.of("duration")) + .build(); + + final BatchedInlineDataSource dataSource = Assertions.assertInstanceOf( + BatchedInlineDataSource.class, + handler.resolveDataSource(query, authenticationResult) + ); + + Assertions.assertEquals(List.of("duration"), requestReceivedByProvider.get().columns()); + Assertions.assertEquals(List.of("duration"), dataSource.getRowSignature().getColumnNames()); + Assertions.assertEquals(1, rowsReceivedByProvider.get().size()); + Assertions.assertArrayEquals(new Object[]{"task-b", 20L}, rowsReceivedByProvider.get().get(0)); + } + } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProvider.java b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProvider.java index 8aaaa0c6ff4..923be45bf4a 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProvider.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProvider.java @@ -28,6 +28,8 @@ import com.google.inject.Inject; import com.google.inject.Provider; import it.unimi.dsi.fastutil.ints.IntOpenHashSet; import it.unimi.dsi.fastutil.ints.IntSet; +import org.apache.druid.query.BatchedInlineDataSource; +import org.apache.druid.query.DataSource; import org.apache.druid.query.filter.DimFilter; import org.apache.druid.query.filter.EqualityFilter; import org.apache.druid.query.filter.InDimFilter; @@ -40,6 +42,7 @@ import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.system.table.SegmentsTableDescriptor; import org.apache.druid.server.system.table.SystemTableDataProvider; import org.apache.druid.server.system.table.SystemTablePushdownFilter; +import org.apache.druid.server.system.table.SystemTableQueryRequest; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.apache.druid.timeline.SegmentStatusInCluster; @@ -49,6 +52,7 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.stream.IntStream; @@ -100,6 +104,23 @@ public class SegmentsTableDataProvider implements SystemTableDataProvider return PUSHDOWN_FILTERS; } + @Override + public Optional<DataSource> getAuthorizedDataSource( + final SystemTableQueryRequest request, + final Iterable<Object[]> authorizedRows + ) + { + final int[] projects = request.columns() + .stream() + .mapToInt(SegmentsTableDescriptor.ROW_SIGNATURE::indexOf) + .toArray(); + final Iterable<Object[]> projectedRows = Iterables.transform( + authorizedRows, + row -> projectRow(row, projects, jsonMapper) + ); + return Optional.of(new BatchedInlineDataSource(projectedRows, request.rowSignature())); + } + @Override public Iterable<Object[]> getRows( final List<DimFilter> filters, diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProviderTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProviderTest.java index 17c8c792f88..564444cb693 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProviderTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/schema/SegmentsTableDataProviderTest.java @@ -23,11 +23,16 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Provider; import org.apache.druid.jackson.DefaultObjectMapper; import org.apache.druid.java.util.common.Intervals; +import org.apache.druid.query.BatchedInlineDataSource; +import org.apache.druid.query.DataSource; import org.apache.druid.query.filter.DimFilter; import org.apache.druid.query.filter.SelectorDimFilter; +import org.apache.druid.segment.column.ColumnType; +import org.apache.druid.segment.column.RowSignature; import org.apache.druid.server.security.AuthConfig; import org.apache.druid.server.security.AuthenticationResult; import org.apache.druid.server.system.table.SegmentsTableDescriptor; +import org.apache.druid.server.system.table.SystemTableQueryRequest; import org.apache.druid.timeline.DataSegment; import org.apache.druid.timeline.SegmentId; import org.junit.jupiter.api.Assertions; @@ -109,6 +114,39 @@ public class SegmentsTableDataProviderTest Mockito.verify(mapper, Mockito.never()).writeValueAsString(Mockito.any()); } + @Test + public void testAuthorizedDataSourceProjectsRowsIntoBatches() throws Exception + { + final ObjectMapper mapper = Mockito.mock(ObjectMapper.class); + final SegmentsTableDataProvider provider = new SegmentsTableDataProvider( + () -> Mockito.mock(BrokerSegmentMetadataCache.class), + Mockito.mock(MetadataSegmentView.class), + mapper + ); + final Object[] rawRow = new Object[SegmentsTableDescriptor.ROW_SIGNATURE.size()]; + rawRow[SegmentsTableDescriptor.ROW_SIGNATURE.indexOf("datasource")] = "foo"; + rawRow[SegmentsTableDescriptor.ROW_SIGNATURE.indexOf("dimensions")] = List.of("unused"); + final RowSignature projectedSignature = RowSignature.builder() + .add("datasource", ColumnType.STRING) + .build(); + + final DataSource dataSource = provider.getAuthorizedDataSource( + new SystemTableQueryRequest( + List.of("datasource"), + projectedSignature + ), + List.<Object[]>of(rawRow) + ).orElseThrow(); + + final BatchedInlineDataSource batchedDataSource = Assertions.assertInstanceOf( + BatchedInlineDataSource.class, + dataSource + ); + Assertions.assertEquals(projectedSignature, batchedDataSource.getRowSignature()); + Assertions.assertArrayEquals(new Object[]{"foo"}, batchedDataSource.getRows().iterator().next()); + Mockito.verify(mapper, Mockito.never()).writeValueAsString(Mockito.any()); + } + private static List<Object[]> toRows(final Iterable<Object[]> rows) { final List<Object[]> result = new java.util.ArrayList<>(); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
