This is an automated email from the ASF dual-hosted git repository. xiangfu0 pushed a commit to branch xiangfu0/mse-optional-segments in repository https://gitbox.apache.org/repos/asf/pinot.git
commit 1745fee9a18e68ee9c7794cd0fa513585e71eaea Author: Xiang Fu <[email protected]> AuthorDate: Fri Oct 23 13:39:12 2026 -0700 Query optional segments in multi-stage leaf stages The broker routes a new segment it has not seen ONLINE yet as optional. The single-stage engine sends such segments to the servers, which query them if loaded and skip them otherwise, but multi-stage leaf stages dropped them, so right after an upload or a realtime commit a multi-stage query could silently return fewer rows. Carry the optional segments of each leaf worker from the broker to the server leaf, apart from its segments, and set them on the server request. Covers non-partitioned leaf stages of physical and logical tables on both planners. Older servers ignore the new field and keep skipping them. --- pinot-common/src/main/proto/worker.proto | 4 + .../core/routing/MockRoutingManagerFactory.java | 18 ++- .../planner/physical/DispatchablePlanContext.java | 8 ++ .../planner/physical/DispatchablePlanMetadata.java | 15 +++ .../v2/PlanFragmentAndMailboxAssignment.java | 1 + .../planner/physical/v2/TableScanMetadata.java | 13 +++ .../opt/rules/LeafStageWorkerAssignmentRule.java | 70 ++++++++++-- .../pinot/query/routing/QueryPlanSerDeUtils.java | 19 +++- .../apache/pinot/query/routing/WorkerManager.java | 62 ++++++++++- .../apache/pinot/query/routing/WorkerMetadata.java | 31 ++++++ .../rules/LeafStageWorkerAssignmentRuleTest.java | 33 +++++- .../query/routing/QueryPlanSerDeUtilsTest.java | 36 ++++++ .../pinot/query/routing/WorkerManagerTest.java | 76 +++++++++++++ .../plan/server/ServerPlanRequestUtils.java | 47 ++++++-- .../plan/server/ServerPlanRequestUtilsTest.java | 120 ++++++++++++++++++++ .../queries/OptionalSegmentsQueryRunnerTest.java | 122 +++++++++++++++++++++ .../testutils/MockInstanceDataManagerFactory.java | 19 +++- 17 files changed, 662 insertions(+), 32 deletions(-) diff --git a/pinot-common/src/main/proto/worker.proto b/pinot-common/src/main/proto/worker.proto index 2a84f200035..b14d6b4b94f 100644 --- a/pinot-common/src/main/proto/worker.proto +++ b/pinot-common/src/main/proto/worker.proto @@ -101,6 +101,10 @@ message WorkerMetadata { SegmentsMap tableSegmentsMap = 4; // Leaf-stage segments of a logical table, keyed by physical table name (with type suffix). Same encoding rules. SegmentsMap logicalTableSegmentsMap = 5; + // Leaf-stage optional segments (new segments the broker has not yet seen ONLINE), keyed like the segment map above; + // a server queries them if it has them and skips them otherwise. Absent when the worker has none. Same encoding + // rules, with "optionalSegmentsMap" as the custom property. + SegmentsMap optionalSegmentsMap = 6; } // Segment names to scan per table type or physical table name, see WorkerMetadata. diff --git a/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java index c0f30ec9cd4..e29ba80bf4d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java @@ -53,6 +53,7 @@ public class MockRoutingManagerFactory { private final Set<String> _hybridTables; private final Map<String, ServerInstance> _serverInstances; private final Map<String, Map<String, List<ServerInstance>>> _tableSegmentServersMap; + private final Map<String, Map<String, ServerInstance>> _tableOptionalSegmentServerMap; private final Set<String> _disabledTables; public MockRoutingManagerFactory(int... ports) { @@ -61,6 +62,7 @@ public class MockRoutingManagerFactory { _hybridTables = new HashSet<>(); _serverInstances = new HashMap<>(); _tableSegmentServersMap = new HashMap<>(); + _tableOptionalSegmentServerMap = new HashMap<>(); _disabledTables = new HashSet<>(); for (int port : ports) { _serverInstances.put(toHostname(port), getServerInstance(HOST_NAME, port, port, port, port)); @@ -91,6 +93,13 @@ public class MockRoutingManagerFactory { .add(serverInstance); } + /// Registers a segment routed to the given server as optional, like a new segment the broker has not seen ONLINE yet. + /// As in `BaseBrokerRoutingManager`, it is only routed when the server also gets a non-optional segment of the table. + public void registerOptionalSegment(int insertToServerPort, String tableNameWithType, String segmentName) { + _tableOptionalSegmentServerMap.computeIfAbsent(tableNameWithType, k -> new HashMap<>()) + .put(segmentName, _serverInstances.get(toHostname(insertToServerPort))); + } + /// Registers a routing entry for a table that has no routable segments (e.g. an empty table, or one whose segments /// were all pruned). This mirrors `BaseBrokerRoutingManager`, which returns a non-null [RoutingTable] /// with an empty server-to-segments map for such a table (as opposed to `null` when no routing entry exists). @@ -123,10 +132,17 @@ public class MockRoutingManagerFactory { server = servers.get(serverId % numServers); serverId++; } - segmentsToQueryMap.computeIfAbsent(server, k -> new SegmentsToQuery(new ArrayList<>(), null)) + segmentsToQueryMap.computeIfAbsent(server, k -> new SegmentsToQuery(new ArrayList<>(), new ArrayList<>())) .getSegments() .add(segment); } + for (Map.Entry<String, ServerInstance> optionalSegmentServerEntry + : _tableOptionalSegmentServerMap.getOrDefault(tableNameWithType, Map.of()).entrySet()) { + SegmentsToQuery segmentsToQuery = segmentsToQueryMap.get(optionalSegmentServerEntry.getValue()); + if (segmentsToQuery != null) { + segmentsToQuery.getOptionalSegments().add(optionalSegmentServerEntry.getKey()); + } + } routingTableMap.put(tableNameWithType, new RoutingTable(segmentsToQueryMap, List.of(), 0)); tableSegmentsMap.put(tableNameWithType, new ArrayList<>(segmentServersMap.keySet())); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java index 58e595d131e..ef855128b66 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanContext.java @@ -187,6 +187,8 @@ public class DispatchablePlanContext { dispatchablePlanMetadata.getWorkerIdToSegmentsMap(); Map<Integer, Map<String, List<String>>> workerIdToTableNameSegmentsMap = dispatchablePlanMetadata.getWorkerIdToTableSegmentsMap(); + Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap = + dispatchablePlanMetadata.getWorkerIdToOptionalSegmentsMap(); Map<Integer, Map<Integer, MailboxInfos>> workerIdToMailboxesMap = dispatchablePlanMetadata.getWorkerIdToMailboxesMap(); Preconditions.checkArgument(workerIdToSegmentsMap == null || workerIdToTableNameSegmentsMap == null, @@ -214,6 +216,12 @@ public class DispatchablePlanContext { workerId); workerMetadata.setLogicalTableSegmentsMap(tableNameSegmentsMap); } + if (workerIdToOptionalSegmentsMap != null) { + Map<String, List<String>> optionalSegmentsMap = workerIdToOptionalSegmentsMap.get(workerId); + if (optionalSegmentsMap != null) { + workerMetadata.setOptionalSegmentsMap(optionalSegmentsMap); + } + } workerMetadataArray[workerId] = workerMetadata; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java index c47b434110b..b1f8770887d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java @@ -84,6 +84,9 @@ public class DispatchablePlanMetadata implements Serializable { /// Map from workerId -> {physicalTableName -> segments} is required for logical tables. private Map<Integer, Map<String, List<String>>> _workerIdToTableSegmentsMap; + /// Map from workerId -> {tableType or physicalTableName -> optional segments}, keyed like the segment map of the + /// stage; only workers with optional segments have an entry. + private Map<Integer, Map<String, List<String>>> _workerIdToOptionalSegmentsMap; private LogicalTableRouteInfo _logicalTableRouteInfo; public List<String> getScannedTables() { @@ -248,4 +251,16 @@ public class DispatchablePlanMetadata implements Serializable { Map<Integer, Map<String, List<String>>> workerIdToTableSegmentsMap) { _workerIdToTableSegmentsMap = workerIdToTableSegmentsMap; } + + /// Returns the optional segments of each leaf-stage worker (see `WorkerMetadata#getOptionalSegmentsMap()`), or `null` + /// when no worker has any. + @Nullable + public Map<Integer, Map<String, List<String>>> getWorkerIdToOptionalSegmentsMap() { + return _workerIdToOptionalSegmentsMap; + } + + public void setWorkerIdToOptionalSegmentsMap( + @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap) { + _workerIdToOptionalSegmentsMap = workerIdToOptionalSegmentsMap; + } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PlanFragmentAndMailboxAssignment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PlanFragmentAndMailboxAssignment.java index ee419aa11f3..7782114a55f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PlanFragmentAndMailboxAssignment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PlanFragmentAndMailboxAssignment.java @@ -166,6 +166,7 @@ public class PlanFragmentAndMailboxAssignment { // For physical tables, use workerIdToSegmentsMap (tableType -> segments) fragmentMetadata.setWorkerIdToSegmentsMap(tableScanMetadata.getWorkedIdToSegmentsMap()); } + fragmentMetadata.setWorkerIdToOptionalSegmentsMap(tableScanMetadata.getWorkerIdToOptionalSegmentsMap()); NodeHint nodeHint = NodeHint.fromRelHints(tableScan.getHints()); fragmentMetadata.setTableOptions(nodeHint.getHintOptions().get(PinotHintOptions.TABLE_HINT_OPTIONS)); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java index daed25b86f3..47bd6f3122b 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java @@ -60,6 +60,10 @@ public class TableScanMetadata { /// ``` @Nullable private final Map<Integer, Map<String, List<String>>> _workerIdToLogicalTableSegmentsMap; + /// Stores workerId to the optional segments of the worker, keyed like the segment map in use (table type, or + /// physical table name for logical tables). Only workers with optional segments have an entry; `null` when none has. + @Nullable + private final Map<Integer, Map<String, List<String>>> _workerIdToOptionalSegmentsMap; private final Map<String, String> _tableOptions; private final Map<String, Set<String>> _unavailableSegmentsMap; @Nullable @@ -69,11 +73,13 @@ public class TableScanMetadata { /// Constructor for physical table scan metadata. public TableScanMetadata(Set<String> scannedTables, Map<Integer, Map<String, List<String>>> workedIdToSegmentsMap, + @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap, Map<String, String> tableOptions, Map<String, Set<String>> unavailableSegmentsMap, @Nullable TimeBoundaryInfo timeBoundaryInfo) { _scannedTables = scannedTables; _workedIdToSegmentsMap = workedIdToSegmentsMap; _workerIdToLogicalTableSegmentsMap = null; + _workerIdToOptionalSegmentsMap = workerIdToOptionalSegmentsMap; _tableOptions = tableOptions; _unavailableSegmentsMap = unavailableSegmentsMap; _timeBoundaryInfo = timeBoundaryInfo; @@ -83,9 +89,11 @@ public class TableScanMetadata { /// Constructor for logical table scan metadata. public TableScanMetadata(Set<String> scannedTables, Map<Integer, Map<String, List<String>>> workerIdToLogicalTableSegmentsMap, + @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap, Map<String, String> tableOptions, Map<String, Set<String>> unavailableSegmentsMap, @Nullable TimeBoundaryInfo timeBoundaryInfo, boolean isLogicalTable) { _scannedTables = scannedTables; + _workerIdToOptionalSegmentsMap = workerIdToOptionalSegmentsMap; if (isLogicalTable) { _workedIdToSegmentsMap = null; _workerIdToLogicalTableSegmentsMap = workerIdToLogicalTableSegmentsMap; @@ -113,6 +121,11 @@ public class TableScanMetadata { return _workerIdToLogicalTableSegmentsMap; } + @Nullable + public Map<Integer, Map<String, List<String>>> getWorkerIdToOptionalSegmentsMap() { + return _workerIdToOptionalSegmentsMap; + } + public Map<String, String> getTableOptions() { return _tableOptions; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java index 45a07bfb8bf..4740923cb95 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java @@ -172,19 +172,25 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { String tableType = routingEntry.getKey(); RoutingTable routingTable = routingEntry.getValue(); Map<String, List<String>> currentSegmentsMap = new HashMap<>(); + Map<String, List<String>> currentOptionalSegmentsMap = new HashMap<>(); Map<ServerInstance, SegmentsToQuery> tmp = routingTable.getServerInstanceToSegmentsMap(); for (Map.Entry<ServerInstance, SegmentsToQuery> serverEntry : tmp.entrySet()) { - // TODO: Optional segments are not supported yet by the MSE. String instanceId = serverEntry.getKey().getInstanceId(); Preconditions.checkState(currentSegmentsMap.put(instanceId, serverEntry.getValue().getSegments()) == null, "Entry for server %s and table type: %s already exist!", serverEntry.getKey(), tableType); + List<String> optionalSegments = serverEntry.getValue().getOptionalSegments(); + if (CollectionUtils.isNotEmpty(optionalSegments)) { + currentOptionalSegmentsMap.put(instanceId, optionalSegments); + } _physicalPlannerContext.getInstanceIdToQueryServerInstance().computeIfAbsent(instanceId, (ignore) -> new QueryServerInstance(serverEntry.getKey())); } if (tableType.equalsIgnoreCase(TableType.OFFLINE.name())) { instanceIdToSegments._offlineTableSegmentsMap = currentSegmentsMap; + instanceIdToSegments._offlineTableOptionalSegmentsMap = currentOptionalSegmentsMap; } else { instanceIdToSegments._realtimeTableSegmentsMap = currentSegmentsMap; + instanceIdToSegments._realtimeTableOptionalSegmentsMap = currentOptionalSegmentsMap; } if (!routingTable.getUnavailableSegments().isEmpty()) { // Set unavailable segments in context, keyed by PRelNode ID. @@ -202,7 +208,7 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { TableScanWorkerAssignmentResult workerAssignmentResult = assignTableScan(tableName, fieldNames, instanceIdToSegments, tablePartitionInfoMap, inferInvalidPartitionSegment, inferRealtimeSegmentPartition); TableScanMetadata metadata = new TableScanMetadata(Set.of(tableName), workerAssignmentResult._workerIdToSegmentsMap, - tableOptions, segmentUnavailableMap, timeBoundaryInfo); + workerAssignmentResult._workerIdToOptionalSegmentsMap, tableOptions, segmentUnavailableMap, timeBoundaryInfo); return tableScan.with(workerAssignmentResult._pinotDataDistribution, metadata); } @@ -240,6 +246,8 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { tpi.getSegmentsWithInvalidPartition().size()); } if (assignmentResult != null) { + // A partitioned assignment carries no optional segments: the server picked for a new segment need not be + // the one its partition is assigned to, so they are skipped as before. return assignmentResult; } } @@ -258,6 +266,7 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { } // For each server, assign one worker each. Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap = new HashMap<>(); + Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap = new HashMap<>(); List<String> workers = new ArrayList<>(); for (Map.Entry<String, Map<String, List<String>>> entry : instanceIdToTableTypeToSegmentsMap.entrySet()) { String instanceId = entry.getKey(); @@ -265,6 +274,12 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { String tableType = tableTypeAndSegments.getKey(); List<String> segments = tableTypeAndSegments.getValue(); workerIdToSegmentsMap.computeIfAbsent(workers.size(), (x) -> new HashMap<>()).put(tableType, segments); + List<String> optionalSegments = + instanceIdToSegments.getOptionalSegments(TableType.valueOf(tableType), instanceId); + if (CollectionUtils.isNotEmpty(optionalSegments)) { + workerIdToOptionalSegmentsMap.computeIfAbsent(workers.size(), (x) -> new HashMap<>()) + .put(tableType, optionalSegments); + } } workers.add(String.format("%s@%s", workers.size(), instanceId)); } @@ -272,7 +287,8 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { : RelDistribution.Type.RANDOM_DISTRIBUTED; PinotDataDistribution pinotDataDistribution = new PinotDataDistribution(distType, workers, workers.hashCode(), null, null); - return new TableScanWorkerAssignmentResult(pinotDataDistribution, workerIdToSegmentsMap); + return new TableScanWorkerAssignmentResult(pinotDataDistribution, workerIdToSegmentsMap, + workerIdToOptionalSegmentsMap.isEmpty() ? null : workerIdToOptionalSegmentsMap); } /// Tries to assign workers for the table-scan node to generate a partitioned data distribution. If this is not @@ -626,13 +642,15 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { String logicalTableName, LogicalTableRouteInfo logicalTableRouteInfo, Map<String, String> tableOptions) { // Collect server to segments mapping across all physical tables Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap = new HashMap<>(); + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalOptionalSegmentsMap = new HashMap<>(); // Process offline tables if (logicalTableRouteInfo.getOfflineTables() != null) { for (TableRouteInfo physicalTableRoute : logicalTableRouteInfo.getOfflineTables()) { if (physicalTableRoute.getOfflineRoutingTable() != null) { transferToServerInstanceLogicalSegmentsMap(physicalTableRoute.getOfflineTableName(), - physicalTableRoute.getOfflineRoutingTable(), serverInstanceToLogicalSegmentsMap); + physicalTableRoute.getOfflineRoutingTable(), serverInstanceToLogicalSegmentsMap, + serverInstanceToLogicalOptionalSegmentsMap); } } } @@ -642,7 +660,8 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { for (TableRouteInfo physicalTableRoute : logicalTableRouteInfo.getRealtimeTables()) { if (physicalTableRoute.getRealtimeRoutingTable() != null) { transferToServerInstanceLogicalSegmentsMap(physicalTableRoute.getRealtimeTableName(), - physicalTableRoute.getRealtimeRoutingTable(), serverInstanceToLogicalSegmentsMap); + physicalTableRoute.getRealtimeRoutingTable(), serverInstanceToLogicalSegmentsMap, + serverInstanceToLogicalOptionalSegmentsMap); } } } @@ -655,6 +674,7 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { // Assign 1 worker per server int numWorkers = sortedServerInstanceToSegmentsMap.size(); Map<Integer, Map<String, List<String>>> workerIdToLogicalTableSegmentsMap = new HashMap<>(); + Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap = new HashMap<>(); List<String> workers = new ArrayList<>(); List<ServerInstance> serverInstances = new ArrayList<>(); @@ -664,6 +684,11 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { Map<String, List<String>> physicalTableSegments = entry.getValue(); workerIdToLogicalTableSegmentsMap.put(workerId, physicalTableSegments); + Map<String, List<String>> physicalTableOptionalSegments = + serverInstanceToLogicalOptionalSegmentsMap.get(serverInstance); + if (physicalTableOptionalSegments != null) { + workerIdToOptionalSegmentsMap.put(workerId, physicalTableOptionalSegments); + } workers.add(String.format("%s@%s", workerId, serverInstance.getInstanceId())); serverInstances.add(serverInstance); } @@ -689,6 +714,7 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { TableScanMetadata metadata = new TableScanMetadata( Set.of(logicalTableName), workerIdToLogicalTableSegmentsMap, + workerIdToOptionalSegmentsMap.isEmpty() ? null : workerIdToOptionalSegmentsMap, tableOptions, unavailableSegmentsMap, timeBoundaryInfo, @@ -698,28 +724,46 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { return new LogicalTableScanWorkerAssignmentResult(pinotDataDistribution, metadata, serverInstances); } - /// Helper method to transfer routing table entries to the server instance logical segments map. + /// Helper method to transfer routing table entries to the server instance logical segments map, and the optional + /// segments, if any, to the server instance logical optional segments map. @VisibleForTesting static void transferToServerInstanceLogicalSegmentsMap(String physicalTableName, Map<ServerInstance, SegmentsToQuery> routingTable, - Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap) { + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap, + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalOptionalSegmentsMap) { for (Map.Entry<ServerInstance, SegmentsToQuery> entry : routingTable.entrySet()) { ServerInstance serverInstance = entry.getKey(); List<String> segments = entry.getValue().getSegments(); serverInstanceToLogicalSegmentsMap .computeIfAbsent(serverInstance, k -> new HashMap<>()) .put(physicalTableName, segments); + List<String> optionalSegments = entry.getValue().getOptionalSegments(); + if (CollectionUtils.isNotEmpty(optionalSegments)) { + serverInstanceToLogicalOptionalSegmentsMap + .computeIfAbsent(serverInstance, k -> new HashMap<>()) + .put(physicalTableName, optionalSegments); + } } } static class TableScanWorkerAssignmentResult { final PinotDataDistribution _pinotDataDistribution; final Map<Integer, Map<String, List<String>>> _workerIdToSegmentsMap; + /// Optional segments per worker, keyed by table type, or `null` when no worker has any. + @Nullable + final Map<Integer, Map<String, List<String>>> _workerIdToOptionalSegmentsMap; TableScanWorkerAssignmentResult(PinotDataDistribution pinotDataDistribution, Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap) { + this(pinotDataDistribution, workerIdToSegmentsMap, null); + } + + TableScanWorkerAssignmentResult(PinotDataDistribution pinotDataDistribution, + Map<Integer, Map<String, List<String>>> workerIdToSegmentsMap, + @Nullable Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap) { _pinotDataDistribution = pinotDataDistribution; _workerIdToSegmentsMap = workerIdToSegmentsMap; + _workerIdToOptionalSegmentsMap = workerIdToOptionalSegmentsMap; } } @@ -743,12 +787,24 @@ public class LeafStageWorkerAssignmentRule extends PRelOptRule { Map<String, List<String>> _offlineTableSegmentsMap; @Nullable Map<String, List<String>> _realtimeTableSegmentsMap; + /// Optional segments by instance id, holding only the instances that have any. + @Nullable + Map<String, List<String>> _offlineTableOptionalSegmentsMap; + @Nullable + Map<String, List<String>> _realtimeTableOptionalSegmentsMap; @Nullable Map<String, List<String>> getSegmentsMap(TableType tableType) { return tableType == TableType.OFFLINE ? _offlineTableSegmentsMap : _realtimeTableSegmentsMap; } + @Nullable + List<String> getOptionalSegments(TableType tableType, String instanceId) { + Map<String, List<String>> optionalSegmentsMap = + tableType == TableType.OFFLINE ? _offlineTableOptionalSegmentsMap : _realtimeTableOptionalSegmentsMap; + return optionalSegmentsMap != null ? optionalSegmentsMap.get(instanceId) : null; + } + Set<String> getActiveTableTypes() { Set<String> tableTypes = new HashSet<>(); if (MapUtils.isNotEmpty(_offlineTableSegmentsMap)) { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java index 5ad7b559d85..4e10f7f89c0 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/QueryPlanSerDeUtils.java @@ -86,10 +86,12 @@ public class QueryPlanSerDeUtils { Map<String, String> customProperties = protoWorkerMetadata.getCustomPropertyMap(); String tableSegmentsJson = customProperties.get(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY); String logicalTableSegmentsJson = customProperties.get(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY); - if (tableSegmentsJson != null || logicalTableSegmentsJson != null) { + String optionalSegmentsJson = customProperties.get(WorkerMetadata.OPTIONAL_SEGMENTS_MAP_KEY); + if (tableSegmentsJson != null || logicalTableSegmentsJson != null || optionalSegmentsJson != null) { Map<String, String> strippedProperties = new HashMap<>(customProperties); strippedProperties.remove(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY); strippedProperties.remove(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY); + strippedProperties.remove(WorkerMetadata.OPTIONAL_SEGMENTS_MAP_KEY); customProperties = Collections.unmodifiableMap(strippedProperties); } WorkerMetadata workerMetadata = @@ -105,6 +107,11 @@ public class QueryPlanSerDeUtils { } else if (logicalTableSegmentsJson != null) { workerMetadata.setLogicalTableSegmentsMapJson(logicalTableSegmentsJson); } + if (protoWorkerMetadata.hasOptionalSegmentsMap()) { + workerMetadata.setOptionalSegmentsMap(fromProtoSegmentsMap(protoWorkerMetadata.getOptionalSegmentsMap())); + } else if (optionalSegmentsJson != null) { + workerMetadata.setOptionalSegmentsMapJson(optionalSegmentsJson); + } return workerMetadata; } @@ -170,6 +177,16 @@ public class QueryPlanSerDeUtils { encodeSegmentsMapJson(logicalTableSegmentsMap)); } } + // A server that predates optional segments ignores them under either encoding (an unknown proto field, or an + // unknown custom property), and so keeps skipping them as before. + Map<String, List<String>> optionalSegmentsMap = workerMetadata.getOptionalSegmentsMap(); + if (optionalSegmentsMap != null) { + if (protoSegmentList) { + builder.setOptionalSegmentsMap(toProtoSegmentsMap(optionalSegmentsMap)); + } else { + builder.putCustomProperty(WorkerMetadata.OPTIONAL_SEGMENTS_MAP_KEY, encodeSegmentsMapJson(optionalSegmentsMap)); + } + } return builder.build(); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java index 3539df3720c..eedcedd30a4 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java @@ -963,6 +963,7 @@ public class WorkerManager { // extract all the instances associated to each table type Map<ServerInstance, Map<String, List<String>>> serverInstanceToSegmentsMap = new HashMap<>(); + Map<ServerInstance, Map<String, List<String>>> serverInstanceToOptionalSegmentsMap = new HashMap<>(); for (Map.Entry<String, RoutingTable> routingEntry : routingTableMap.entrySet()) { String tableType = routingEntry.getKey(); RoutingTable routingTable = routingEntry.getValue(); @@ -971,9 +972,10 @@ public class WorkerManager { for (Map.Entry<ServerInstance, SegmentsToQuery> serverEntry : segmentsMap.entrySet()) { Map<String, List<String>> tableTypeToSegmentListMap = serverInstanceToSegmentsMap.computeIfAbsent(serverEntry.getKey(), k -> new HashMap<>()); - // TODO: support optional segments for multi-stage engine. Preconditions.checkState(tableTypeToSegmentListMap.put(tableType, serverEntry.getValue().getSegments()) == null, "Entry for server {} and table type: {} already exist!", serverEntry.getKey(), tableType); + addOptionalSegments(serverEntry.getKey(), tableType, serverEntry.getValue(), + serverInstanceToOptionalSegmentsMap); } // attach unavailable segments to metadata @@ -1008,6 +1010,42 @@ public class WorkerManager { metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setWorkerIdToSegmentsMap(workerIdToSegmentsMap); + metadata.setWorkerIdToOptionalSegmentsMap( + toWorkerIdToOptionalSegmentsMap(sortedServerInstanceToSegmentsMap, serverInstanceToOptionalSegmentsMap)); + } + + /// Records the optional segments the broker routed to a server, if any. These are new segments the broker has not + /// seen ONLINE yet: the server queries them if it already has them and skips them otherwise, the same as on the + /// single-stage engine. They must not be merged into the segment list, where a not yet loaded one would fail the + /// query as missing. + private static void addOptionalSegments(ServerInstance serverInstance, String key, SegmentsToQuery segmentsToQuery, + Map<ServerInstance, Map<String, List<String>>> serverInstanceToOptionalSegmentsMap) { + List<String> optionalSegments = segmentsToQuery.getOptionalSegments(); + if (CollectionUtils.isNotEmpty(optionalSegments)) { + serverInstanceToOptionalSegmentsMap.computeIfAbsent(serverInstance, k -> new HashMap<>()) + .put(key, optionalSegments); + } + } + + /// Maps the optional segments collected by [#addOptionalSegments] to the worker id of their server, given the + /// servers in worker id order, or returns `null` when no server has any. + @Nullable + private static Map<Integer, Map<String, List<String>>> toWorkerIdToOptionalSegmentsMap( + List<Map.Entry<ServerInstance, Map<String, List<String>>>> sortedServerInstanceToSegmentsMap, + Map<ServerInstance, Map<String, List<String>>> serverInstanceToOptionalSegmentsMap) { + if (serverInstanceToOptionalSegmentsMap.isEmpty()) { + return null; + } + Map<Integer, Map<String, List<String>>> workerIdToOptionalSegmentsMap = new HashMap<>(); + int numWorkers = sortedServerInstanceToSegmentsMap.size(); + for (int workerId = 0; workerId < numWorkers; workerId++) { + Map<String, List<String>> optionalSegmentsMap = + serverInstanceToOptionalSegmentsMap.get(sortedServerInstanceToSegmentsMap.get(workerId).getKey()); + if (optionalSegmentsMap != null) { + workerIdToOptionalSegmentsMap.put(workerId, optionalSegmentsMap); + } + } + return workerIdToOptionalSegmentsMap; } /// Acquire routing table for items listed in [org.apache.pinot.query.planner.plannode.TableScanNode]. @@ -1231,6 +1269,10 @@ public class WorkerManager { /// segment lists (or the whole map) on the metadata, rather than by editing them in place: part of what the /// assignment is built from is the broker's published partition metadata, shared across queries and read concurrently /// by other planning threads. What is handed over is nevertheless kept safe to edit in place. + /// + /// The optional segments of each worker are held apart from its segment lists, in + /// [DispatchablePlanMetadata#getWorkerIdToOptionalSegmentsMap()], under the same worker ids; an override that filters + /// segments has to filter those too, or they are queried unfiltered. protected void filterLeafStageSegments(DispatchablePlanContext context, DispatchablePlanMetadata metadata) { } @@ -1348,17 +1390,20 @@ public class WorkerManager { return brokerRequest; } - private static void assignTableSegmentsToWorkers(LogicalTableRouteInfo logicalTableRouteInfo, + @VisibleForTesting + static void assignTableSegmentsToWorkers(LogicalTableRouteInfo logicalTableRouteInfo, DispatchablePlanMetadata metadata) { Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap = new HashMap<>(); + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalOptionalSegmentsMap = new HashMap<>(); if (logicalTableRouteInfo.getOfflineTables() != null) { for (TableRouteInfo physicalTableRoute : logicalTableRouteInfo.getOfflineTables()) { // Routing table maybe null if no routing table is found OR there are no segments. if (physicalTableRoute.getOfflineRoutingTable() != null) { transferToServerInstanceLogicalSegmentsMap(physicalTableRoute.getOfflineTableName(), - physicalTableRoute.getOfflineRoutingTable(), serverInstanceToLogicalSegmentsMap); + physicalTableRoute.getOfflineRoutingTable(), serverInstanceToLogicalSegmentsMap, + serverInstanceToLogicalOptionalSegmentsMap); } } } @@ -1368,7 +1413,8 @@ public class WorkerManager { // Routing table maybe null if no routing table is found OR there are no segments. if (physicalTableRoute.getRealtimeRoutingTable() != null) { transferToServerInstanceLogicalSegmentsMap(physicalTableRoute.getRealtimeTableName(), - physicalTableRoute.getRealtimeRoutingTable(), serverInstanceToLogicalSegmentsMap); + physicalTableRoute.getRealtimeRoutingTable(), serverInstanceToLogicalSegmentsMap, + serverInstanceToLogicalOptionalSegmentsMap); } } } @@ -1398,18 +1444,22 @@ public class WorkerManager { metadata.setWorkerIdToServerInstanceMap(workerIdToServerInstanceMap); metadata.setWorkerIdToTableSegmentsMap(workerIdToLogicalTableSegmentsMap); + metadata.setWorkerIdToOptionalSegmentsMap( + toWorkerIdToOptionalSegmentsMap(sortedServerInstanceToSegmentsMap, serverInstanceToLogicalOptionalSegmentsMap)); } private static void transferToServerInstanceLogicalSegmentsMap(String physicalTableName, Map<ServerInstance, SegmentsToQuery> segmentsMap, - Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap) { + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalSegmentsMap, + Map<ServerInstance, Map<String, List<String>>> serverInstanceToLogicalOptionalSegmentsMap) { for (Map.Entry<ServerInstance, SegmentsToQuery> serverEntry : segmentsMap.entrySet()) { Map<String, List<String>> tableNameToSegmentsMap = serverInstanceToLogicalSegmentsMap.computeIfAbsent(serverEntry.getKey(), k -> new HashMap<>()); - // TODO: support optional segments for multi-stage engine. Preconditions.checkState( tableNameToSegmentsMap.put(physicalTableName, serverEntry.getValue().getSegments()) == null, "Entry for server {} and physical table: {} already exist!", serverEntry.getKey(), physicalTableName); + addOptionalSegments(serverEntry.getKey(), physicalTableName, serverEntry.getValue(), + serverInstanceToLogicalOptionalSegmentsMap); } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java index f5781e780d5..18a1d198564 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerMetadata.java @@ -54,6 +54,7 @@ public class WorkerMetadata { /// broker/server versions keep working; never present in [#getCustomProperties()] of a decoded instance. public static final String TABLE_SEGMENTS_MAP_KEY = "tableSegmentsMap"; public static final String LOGICAL_TABLE_SEGMENTS_MAP_KEY = "logicalTableSegmentsMap"; + public static final String OPTIONAL_SEGMENTS_MAP_KEY = "optionalSegmentsMap"; private static final TypeReference<Map<String, List<String>>> SEGMENTS_MAP_TYPE = new TypeReference<>() { }; @@ -65,12 +66,17 @@ public class WorkerMetadata { private volatile Map<String, List<String>> _tableSegmentsMap; @Nullable private volatile Map<String, List<String>> _logicalTableSegmentsMap; + @Nullable + private volatile Map<String, List<String>> _optionalSegmentsMap; /// The legacy JSON encoding of [#_tableSegmentsMap] as received from the broker, parsed on first access. @Nullable private String _tableSegmentsMapJson; /// The legacy JSON encoding of [#_logicalTableSegmentsMap] as received from the broker, parsed on first access. @Nullable private String _logicalTableSegmentsMapJson; + /// The legacy JSON encoding of [#_optionalSegmentsMap] as received from the broker, parsed on first access. + @Nullable + private String _optionalSegmentsMapJson; public WorkerMetadata(int workerId, Map<Integer, MailboxInfos> mailboxInfosMap) { this(workerId, mailboxInfosMap, new HashMap<>()); @@ -141,6 +147,31 @@ public class WorkerMetadata { _logicalTableSegmentsMapJson = logicalTableSegmentsMapJson; } + /// Optional segments to scan if the server has them, keyed like the worker's segment map (table type, or physical + /// table name for a logical table), or `null` when the worker has none. These are the new segments the broker routes + /// as optional because it has not seen them ONLINE yet; unlike the segments of the segment map, a missing one is not + /// reported as an error. + @Nullable + public Map<String, List<String>> getOptionalSegmentsMap() { + Map<String, List<String>> optionalSegmentsMap = _optionalSegmentsMap; + if (optionalSegmentsMap == null && _optionalSegmentsMapJson != null) { + optionalSegmentsMap = decodeSegmentsMapJson(_optionalSegmentsMapJson); + _optionalSegmentsMap = optionalSegmentsMap; + } + return optionalSegmentsMap; + } + + /// Stores `optionalSegmentsMap` by reference, with the same no-mutation contract as [#setTableSegmentsMap]. + public void setOptionalSegmentsMap(Map<String, List<String>> optionalSegmentsMap) { + _optionalSegmentsMap = optionalSegmentsMap; + } + + /// Stores the legacy JSON encoding of the optional segments map, to be parsed by [#getOptionalSegmentsMap] on first + /// access. + void setOptionalSegmentsMapJson(String optionalSegmentsMapJson) { + _optionalSegmentsMapJson = optionalSegmentsMapJson; + } + /// A leaf-stage worker carries a (possibly empty) segment map, parsed or not; an intermediate-stage worker carries /// none. public boolean isLeafStageWorker() { diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRuleTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRuleTest.java index 14b4d9b5992..70053c84205 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRuleTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRuleTest.java @@ -692,13 +692,38 @@ public class LeafStageWorkerAssignmentRuleTest { Map<ServerInstance, Map<String, List<String>>> result = new HashMap<>(); LeafStageWorkerAssignmentRule.transferToServerInstanceLogicalSegmentsMap( - "orders_OFFLINE", routingTable, result); + "orders_OFFLINE", routingTable, result, new HashMap<>()); assertEquals(result.size(), 2); assertEquals(result.get(server1).get("orders_OFFLINE"), List.of("seg1", "seg2")); assertEquals(result.get(server2).get("orders_OFFLINE"), List.of("seg3")); } + /// Optional segments of a logical table are carried per worker keyed by physical table name, like its segments. + @Test + public void testBuildLogicalTableWorkerAssignmentCarriesOptionalSegments() { + ServerInstance server1 = mock(ServerInstance.class); + doReturn("server-1").when(server1).getInstanceId(); + ServerInstance server2 = mock(ServerInstance.class); + doReturn("server-2").when(server2).getInstanceId(); + + TableRouteInfo offlineTableRoute = mock(TableRouteInfo.class); + doReturn("orders_OFFLINE").when(offlineTableRoute).getOfflineTableName(); + doReturn(Map.of(server1, new SegmentsToQuery(List.of("orders_seg1"), List.of()), + server2, new SegmentsToQuery(List.of("orders_seg2"), List.of("orders_newSeg")))).when(offlineTableRoute) + .getOfflineRoutingTable(); + LogicalTableRouteInfo logicalTableRouteInfo = mock(LogicalTableRouteInfo.class); + doReturn(List.of(offlineTableRoute)).when(logicalTableRouteInfo).getOfflineTables(); + + TableScanMetadata metadata = LeafStageWorkerAssignmentRule.buildLogicalTableWorkerAssignment( + "myLogicalTable", logicalTableRouteInfo, Map.of())._tableScanMetadata; + + assertEquals(metadata.getWorkerIdToLogicalTableSegmentsMap().get(1), + Map.of("orders_OFFLINE", List.of("orders_seg2"))); + assertEquals(metadata.getWorkerIdToOptionalSegmentsMap(), + Map.of(1, Map.of("orders_OFFLINE", List.of("orders_newSeg")))); + } + @Test public void testTransferToServerInstanceLogicalSegmentsMapMultipleTables() { ServerInstance server1 = mock(ServerInstance.class); @@ -711,14 +736,16 @@ public class LeafStageWorkerAssignmentRuleTest { SegmentsToQuery seg1 = mock(SegmentsToQuery.class); doReturn(List.of("offline_seg1")).when(seg1).getSegments(); routing1.put(server1, seg1); - LeafStageWorkerAssignmentRule.transferToServerInstanceLogicalSegmentsMap("table_OFFLINE", routing1, result); + LeafStageWorkerAssignmentRule.transferToServerInstanceLogicalSegmentsMap("table_OFFLINE", routing1, result, + new HashMap<>()); // Add second table Map<ServerInstance, SegmentsToQuery> routing2 = new HashMap<>(); SegmentsToQuery seg2 = mock(SegmentsToQuery.class); doReturn(List.of("realtime_seg1")).when(seg2).getSegments(); routing2.put(server1, seg2); - LeafStageWorkerAssignmentRule.transferToServerInstanceLogicalSegmentsMap("table_REALTIME", routing2, result); + LeafStageWorkerAssignmentRule.transferToServerInstanceLogicalSegmentsMap("table_REALTIME", routing2, result, + new HashMap<>()); // Verify single server has both tables assertEquals(result.size(), 1); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java index 4640a709099..d28842fa3f8 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/QueryPlanSerDeUtilsTest.java @@ -44,6 +44,8 @@ public class QueryPlanSerDeUtilsTest { Map.of("OFFLINE", List.of("seg_0", "seg_1"), "REALTIME", List.of("seg__0__0__20240101T0000Z")); private static final Map<String, List<String>> LOGICAL_TABLE_SEGMENTS_MAP = Map.of("t1_OFFLINE", List.of("t1_seg_0"), "t2_REALTIME", List.of("t2_seg_0", "t2_seg_1")); + private static final Map<String, List<String>> OPTIONAL_SEGMENTS_MAP = + Map.of("REALTIME", List.of("seg__0__1__20240101T0100Z")); @DataProvider public static Object[][] encodings() { @@ -60,12 +62,15 @@ public class QueryPlanSerDeUtilsTest { assertFalse(proto.hasLogicalTableSegmentsMap()); assertEquals(proto.getCustomPropertyMap().containsKey(WorkerMetadata.TABLE_SEGMENTS_MAP_KEY), !protoSegmentList); assertFalse(proto.getCustomPropertyMap().containsKey(WorkerMetadata.LOGICAL_TABLE_SEGMENTS_MAP_KEY)); + assertFalse(proto.hasOptionalSegmentsMap()); + assertFalse(proto.getCustomPropertyMap().containsKey(WorkerMetadata.OPTIONAL_SEGMENTS_MAP_KEY)); assertEquals(proto.getCustomPropertyMap().get("foo"), "bar"); WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); assertEquals(decoded.getWorkerId(), 3); assertEquals(decoded.getTableSegmentsMap(), TABLE_SEGMENTS_MAP); assertNull(decoded.getLogicalTableSegmentsMap()); + assertNull(decoded.getOptionalSegmentsMap()); assertTrue(decoded.isLeafStageWorker()); // The JSON is never surfaced as a custom property of the decoded metadata, whichever encoding was used. assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); @@ -93,6 +98,37 @@ public class QueryPlanSerDeUtilsTest { assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); } + @Test(dataProvider = "encodings") + public void testOptionalSegmentsRoundTrip(boolean protoSegmentList) + throws Exception { + WorkerMetadata workerMetadata = leafWorker(TABLE_SEGMENTS_MAP, null); + workerMetadata.setOptionalSegmentsMap(OPTIONAL_SEGMENTS_MAP); + + Worker.WorkerMetadata proto = toProto(workerMetadata, protoSegmentList); + assertEquals(proto.hasOptionalSegmentsMap(), protoSegmentList); + assertEquals(proto.getCustomPropertyMap().containsKey(WorkerMetadata.OPTIONAL_SEGMENTS_MAP_KEY), + !protoSegmentList); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(proto); + assertEquals(decoded.getTableSegmentsMap(), TABLE_SEGMENTS_MAP); + assertEquals(decoded.getOptionalSegmentsMap(), OPTIONAL_SEGMENTS_MAP); + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + } + + /// Optional segments of a logical table are keyed by physical table name, like its segment map. + @Test(dataProvider = "encodings") + public void testLogicalTableOptionalSegmentsRoundTrip(boolean protoSegmentList) + throws Exception { + Map<String, List<String>> optionalSegmentsMap = Map.of("t2_REALTIME", List.of("t2_seg_2")); + WorkerMetadata workerMetadata = leafWorker(null, LOGICAL_TABLE_SEGMENTS_MAP); + workerMetadata.setOptionalSegmentsMap(optionalSegmentsMap); + + WorkerMetadata decoded = QueryPlanSerDeUtils.fromProtoWorkerMetadata(toProto(workerMetadata, protoSegmentList)); + assertEquals(decoded.getLogicalTableSegmentsMap(), LOGICAL_TABLE_SEGMENTS_MAP); + assertEquals(decoded.getOptionalSegmentsMap(), optionalSegmentsMap); + assertEquals(decoded.getCustomProperties(), CUSTOM_PROPERTIES); + } + @Test(dataProvider = "encodings") public void testEmptySegmentListStillMarksLeafWorker(boolean protoSegmentList) throws Exception { diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java index dfc526e6bfd..3dd775a5054 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/WorkerManagerTest.java @@ -33,6 +33,8 @@ import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.Expression; import org.apache.pinot.common.request.Function; import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.core.routing.ImplicitHybridTableRouteInfo; +import org.apache.pinot.core.routing.LogicalTableRouteInfo; import org.apache.pinot.core.routing.MockRoutingManagerFactory; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; @@ -44,6 +46,7 @@ import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.planner.physical.DispatchablePlanContext; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; +import org.apache.pinot.query.planner.physical.DispatchablePlanMetadata; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; @@ -55,6 +58,7 @@ import org.apache.pinot.spi.exception.QueryException; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.sql.FilterKind; import org.apache.pinot.sql.parsers.CalciteSqlCompiler; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.mockito.ArgumentMatchers.anyString; @@ -534,6 +538,78 @@ public class WorkerManagerTest { return new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, tableCache, workerManager); } + /// The broker routes a new segment it has not seen ONLINE yet as optional. A leaf worker must carry the optional + /// segments of its server apart from its segments (so that the server queries them if it has them, and skips them + /// rather than failing otherwise), on both planners. + @Test(dataProvider = "physicalOptimizer") + public void testLeafWorkersCarryOptionalSegments(boolean usePhysicalOptimizer) { + ServerInstance server1 = getServerInstance("localhost", 1); + ServerInstance server2 = getServerInstance("localhost", 2); + Map<String, ServerInstance> serverInstanceMap = + Map.of(server1.getInstanceId(), server1, server2.getInstanceId(), server2); + RoutingTable routingTable = new RoutingTable( + Map.of(server1, new SegmentsToQuery(List.of("segment1"), List.of("newSegment1")), + server2, new SegmentsToQuery(List.of("segment2"), List.of())), List.of(), 0); + CapturingRoutingManager routingManager = + new CapturingRoutingManager(serverInstanceMap, Map.of("testTable_OFFLINE", routingTable)); + QueryEnvironment queryEnvironment = newQueryEnvironment(getSchemaBuilder("testTable").build(), routingManager); + + String query = "SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; SELECT col1 FROM testTable"; + try (QueryEnvironment.CompiledQuery compiledQuery = queryEnvironment.compile(query)) { + DispatchableSubPlan dispatchableSubPlan = compiledQuery.planQuery(0).getQueryPlan(); + DispatchablePlanFragment leafFragment = dispatchableSubPlan.getQueryStagesWithoutRoot().stream() + .filter(fragment -> fragment.getTableName() != null).findFirst().orElseThrow(); + Map<String, WorkerMetadata> workerByServer = new HashMap<>(); + for (Map.Entry<QueryServerInstance, List<Integer>> entry + : leafFragment.getServerInstanceToWorkerIdMap().entrySet()) { + assertEquals(entry.getValue().size(), 1); + workerByServer.put(entry.getKey().getInstanceId(), + leafFragment.getWorkerMetadataList().get(entry.getValue().get(0))); + } + assertEquals(workerByServer.size(), 2); + WorkerMetadata worker1 = workerByServer.get(server1.getInstanceId()); + assertEquals(worker1.getTableSegmentsMap(), Map.of("OFFLINE", List.of("segment1"))); + assertEquals(worker1.getOptionalSegmentsMap(), Map.of("OFFLINE", List.of("newSegment1"))); + WorkerMetadata worker2 = workerByServer.get(server2.getInstanceId()); + assertEquals(worker2.getTableSegmentsMap(), Map.of("OFFLINE", List.of("segment2"))); + assertNull(worker2.getOptionalSegmentsMap()); + } + } + + @DataProvider + public static Object[][] physicalOptimizer() { + return new Object[][]{{false}, {true}}; + } + + /// Optional segments of a logical table are carried per worker keyed by physical table name, like its segments. + @Test + public void testLogicalTableLeafWorkersCarryOptionalSegments() { + ServerInstance server1 = getServerInstance("localhost", 1); + ServerInstance server2 = getServerInstance("localhost", 2); + ImplicitHybridTableRouteInfo offlineTableRoute = mock(ImplicitHybridTableRouteInfo.class); + when(offlineTableRoute.getOfflineTableName()).thenReturn("t1_OFFLINE"); + when(offlineTableRoute.getOfflineRoutingTable()).thenReturn( + Map.of(server1, new SegmentsToQuery(List.of("t1_seg1"), List.of()), + server2, new SegmentsToQuery(List.of("t1_seg2"), List.of()))); + ImplicitHybridTableRouteInfo realtimeTableRoute = mock(ImplicitHybridTableRouteInfo.class); + when(realtimeTableRoute.getRealtimeTableName()).thenReturn("t2_REALTIME"); + when(realtimeTableRoute.getRealtimeRoutingTable()).thenReturn( + Map.of(server1, new SegmentsToQuery(List.of("t2_seg1"), List.of("t2_newSeg1")))); + LogicalTableRouteInfo logicalTableRouteInfo = mock(LogicalTableRouteInfo.class); + when(logicalTableRouteInfo.getOfflineTables()).thenReturn(List.of(offlineTableRoute)); + when(logicalTableRouteInfo.getRealtimeTables()).thenReturn(List.of(realtimeTableRoute)); + + DispatchablePlanMetadata metadata = new DispatchablePlanMetadata(); + WorkerManager.assignTableSegmentsToWorkers(logicalTableRouteInfo, metadata); + + // Workers are assigned in instance id order. + assertEquals(metadata.getWorkerIdToServerInstanceMap().get(0).getInstanceId(), server1.getInstanceId()); + assertEquals(metadata.getWorkerIdToTableSegmentsMap().get(0), + Map.of("t1_OFFLINE", List.of("t1_seg1"), "t2_REALTIME", List.of("t2_seg1"))); + assertEquals(metadata.getWorkerIdToOptionalSegmentsMap(), + Map.of(0, Map.of("t2_REALTIME", List.of("t2_newSeg1")))); + } + @Test public void testBrokerPruningPreservesQueryOptionsOnRoutingRequest() { Schema schema = getSchemaBuilder("testTable").build(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java index 8c35aa5442f..0aad04a2505 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.query.runtime.plan.server; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.math.BigDecimal; import java.util.ArrayList; @@ -28,6 +29,7 @@ import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.function.BiConsumer; import javax.annotation.Nullable; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.request.BrokerRequest; @@ -157,6 +159,7 @@ public class ServerPlanRequestUtils { String rawTableName = TableNameBuilder.extractRawTableName(stageMetadata.getTableName()); Map<String, List<String>> tableSegmentsMap = executionContext.getWorkerMetadata().getTableSegmentsMap(); assert tableSegmentsMap != null; + Map<String, List<String>> optionalSegmentsMap = executionContext.getWorkerMetadata().getOptionalSegmentsMap(); TimeBoundaryInfo timeBoundary = stageMetadata.getTimeBoundary(); int numRequests = tableSegmentsMap.size(); if (numRequests == 1) { @@ -171,7 +174,7 @@ public class ServerPlanRequestUtils { Pair<TableConfig, Schema> tableConfigAndSchema = tableDataManager.getCachedTableConfigAndSchema(); return List.of(compileInstanceRequest(executionContext, pinotQuery, timeBoundary, TableType.OFFLINE, tableDataManager.getTableName(), tableConfigAndSchema.getLeft(), tableConfigAndSchema.getRight(), segments, - null)); + copyOptionalSegments(optionalSegmentsMap, tableType), null)); } else { assert tableType.equals(TableType.REALTIME.name()); String realtimeTableName = TableNameBuilder.forType(TableType.REALTIME).tableNameWithType(rawTableName); @@ -181,7 +184,7 @@ public class ServerPlanRequestUtils { Pair<TableConfig, Schema> tableConfigAndSchema = tableDataManager.getCachedTableConfigAndSchema(); return List.of(compileInstanceRequest(executionContext, pinotQuery, timeBoundary, TableType.REALTIME, tableDataManager.getTableName(), tableConfigAndSchema.getLeft(), tableConfigAndSchema.getRight(), segments, - null)); + copyOptionalSegments(optionalSegmentsMap, tableType), null)); } } else { assert numRequests == 2; @@ -203,18 +206,33 @@ public class ServerPlanRequestUtils { return List.of( compileInstanceRequest(executionContext, new PinotQuery(pinotQuery), timeBoundary, TableType.OFFLINE, offlineTableDataManager.getTableName(), offlineTableConfigAndSchema.getLeft(), - offlineTableConfigAndSchema.getRight(), offlineSegments, null), + offlineTableConfigAndSchema.getRight(), offlineSegments, + copyOptionalSegments(optionalSegmentsMap, TableType.OFFLINE.name()), null), compileInstanceRequest(executionContext, pinotQuery, timeBoundary, TableType.REALTIME, realtimeTableDataManager.getTableName(), realtimeTableConfigAndSchema.getLeft(), - realtimeTableConfigAndSchema.getRight(), realtimeSegments, null)); + realtimeTableConfigAndSchema.getRight(), realtimeSegments, + copyOptionalSegments(optionalSegmentsMap, TableType.REALTIME.name()), null)); } } + /// Returns a copy of the optional segments under the given key, or `null` when there are none. A copy because the + /// list is handed over to [InstanceRequest], whose consumers may append to it (e.g. the newly added segments of an + /// upsert table), while the original belongs to the worker metadata. + @Nullable + private static List<String> copyOptionalSegments(@Nullable Map<String, List<String>> optionalSegmentsMap, + String key) { + if (optionalSegmentsMap == null) { + return null; + } + List<String> optionalSegments = optionalSegmentsMap.get(key); + return CollectionUtils.isNotEmpty(optionalSegments) ? new ArrayList<>(optionalSegments) : null; + } + /// Convert [PinotQuery] into an [InstanceRequest]. private static InstanceRequest compileInstanceRequest(OpChainExecutionContext executionContext, PinotQuery pinotQuery, @Nullable TimeBoundaryInfo timeBoundaryInfo, TableType tableType, String tableNameWithType, TableConfig tableConfig, Schema schema, @Nullable List<String> segmentList, - @Nullable List<TableSegmentsInfo> tableRouteInfoList) { + @Nullable List<String> optionalSegmentList, @Nullable List<TableSegmentsInfo> tableRouteInfoList) { Preconditions.checkArgument(segmentList == null || tableRouteInfoList == null, "Either segmentList OR tableRouteInfoList should be set"); @@ -256,6 +274,9 @@ public class ServerPlanRequestUtils { */ if (segmentList != null) { instanceRequest.setSearchSegments(segmentList); + if (optionalSegmentList != null) { + instanceRequest.setOptionalSegments(optionalSegmentList); + } } else { instanceRequest.setTableSegmentsInfoList(tableRouteInfoList); } @@ -408,7 +429,8 @@ public class ServerPlanRequestUtils { return expressions; } - private static List<InstanceRequest> constructLogicalTableServerQueryRequests( + @VisibleForTesting + static List<InstanceRequest> constructLogicalTableServerQueryRequests( OpChainExecutionContext executionContext, PinotQuery pinotQuery, InstanceDataManager instanceDataManager) { StageMetadata stageMetadata = executionContext.getStageMetadata(); String logicalTableName = stageMetadata.getTableName(); @@ -419,6 +441,7 @@ public class ServerPlanRequestUtils { Map<String, List<String>> logicalTableSegmentsMap = executionContext.getWorkerMetadata().getLogicalTableSegmentsMap(); + Map<String, List<String>> optionalSegmentsMap = executionContext.getWorkerMetadata().getOptionalSegmentsMap(); List<TableSegmentsInfo> offlineTableRouteInfoList = new ArrayList<>(); List<TableSegmentsInfo> realtimeTableRouteInfoList = new ArrayList<>(); @@ -429,6 +452,10 @@ public class ServerPlanRequestUtils { TableSegmentsInfo tableSegmentsInfo = new TableSegmentsInfo(); tableSegmentsInfo.setTableName(physicalTableName); tableSegmentsInfo.setSegments(entry.getValue()); + List<String> optionalSegments = copyOptionalSegments(optionalSegmentsMap, physicalTableName); + if (optionalSegments != null) { + tableSegmentsInfo.setOptionalSegments(optionalSegments); + } if (tableType == TableType.REALTIME) { realtimeTableRouteInfoList.add(tableSegmentsInfo); } else { @@ -447,7 +474,7 @@ public class ServerPlanRequestUtils { String offlineTableName = TableNameBuilder.forType(TableType.OFFLINE).tableNameWithType(logicalTableName); return List.of( compileInstanceRequest(executionContext, pinotQuery, timeBoundaryInfo, TableType.OFFLINE, offlineTableName, - logicalTableContext.getRefOfflineTableConfig(), logicalTableContext.getLogicalTableSchema(), null, + logicalTableContext.getRefOfflineTableConfig(), logicalTableContext.getLogicalTableSchema(), null, null, routeInfoList)); } else { Preconditions.checkNotNull(logicalTableContext.getRefRealtimeTableConfig()); @@ -455,7 +482,7 @@ public class ServerPlanRequestUtils { return List.of( compileInstanceRequest(executionContext, pinotQuery, timeBoundaryInfo, TableType.REALTIME, realtimeTableName, logicalTableContext.getRefRealtimeTableConfig(), - logicalTableContext.getLogicalTableSchema(), null, routeInfoList)); + logicalTableContext.getLogicalTableSchema(), null, null, routeInfoList)); } } else { Preconditions.checkNotNull(logicalTableContext.getRefOfflineTableConfig()); @@ -467,10 +494,10 @@ public class ServerPlanRequestUtils { return List.of( compileInstanceRequest(executionContext, offlinePinotQuery, timeBoundaryInfo, TableType.OFFLINE, offlineTableName, logicalTableContext.getRefOfflineTableConfig(), - logicalTableContext.getLogicalTableSchema(), null, offlineTableRouteInfoList), + logicalTableContext.getLogicalTableSchema(), null, null, offlineTableRouteInfoList), compileInstanceRequest(executionContext, realtimePinotQuery, timeBoundaryInfo, TableType.REALTIME, realtimeTableName, logicalTableContext.getRefRealtimeTableConfig(), - logicalTableContext.getLogicalTableSchema(), null, realtimeTableRouteInfoList)); + logicalTableContext.getLogicalTableSchema(), null, null, realtimeTableRouteInfoList)); } } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java new file mode 100644 index 00000000000..c4933d2d553 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java @@ -0,0 +1,120 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.plan.server; + +import java.util.List; +import java.util.Map; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.request.InstanceRequest; +import org.apache.pinot.core.data.manager.InstanceDataManager; +import org.apache.pinot.core.data.manager.LogicalTableContext; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; +import org.apache.pinot.query.routing.StageMetadata; +import org.apache.pinot.query.routing.WorkerMetadata; +import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.segment.local.data.manager.TableDataManager; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.testng.annotations.Test; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotSame; + + +/// Tests that the leaf stage hands the optional segments of its worker to the requests it builds for the server. +public class ServerPlanRequestUtilsTest { + private static final Schema SCHEMA = new Schema.SchemaBuilder().setSchemaName("tbl") + .addSingleValueDimension("col1", FieldSpec.DataType.STRING, "") + .build(); + + @Test + public void testHybridTableRequestsCarryOptionalSegmentsPerTableType() { + WorkerMetadata workerMetadata = new WorkerMetadata(0, Map.of()); + workerMetadata.setTableSegmentsMap(Map.of("OFFLINE", List.of("offlineSeg"), "REALTIME", List.of("realtimeSeg"))); + workerMetadata.setOptionalSegmentsMap(Map.of("REALTIME", List.of("newRealtimeSeg"))); + InstanceDataManager instanceDataManager = mock(InstanceDataManager.class); + for (TableType tableType : TableType.values()) { + String tableNameWithType = "tbl_" + tableType; + TableDataManager tableDataManager = mock(TableDataManager.class); + when(tableDataManager.getTableName()).thenReturn(tableNameWithType); + when(tableDataManager.getCachedTableConfigAndSchema()).thenReturn(Pair.of(tableConfig(tableType), SCHEMA)); + when(instanceDataManager.getTableDataManager(tableNameWithType)).thenReturn(tableDataManager); + } + + List<InstanceRequest> instanceRequests = ServerPlanRequestUtils.constructServerQueryRequests( + executionContext("tbl", workerMetadata), CalciteSqlParser.compileToPinotQuery("SELECT col1 FROM tbl"), + instanceDataManager); + + assertEquals(instanceRequests.size(), 2); + InstanceRequest offlineRequest = instanceRequests.get(0); + assertEquals(offlineRequest.getSearchSegments(), List.of("offlineSeg")); + assertFalse(offlineRequest.isSetOptionalSegments()); + InstanceRequest realtimeRequest = instanceRequests.get(1); + assertEquals(realtimeRequest.getSearchSegments(), List.of("realtimeSeg")); + assertEquals(realtimeRequest.getOptionalSegments(), List.of("newRealtimeSeg")); + // The server may append to the list (e.g. the newly added segments of an upsert table), so it must be a copy. + assertNotSame(realtimeRequest.getOptionalSegments(), workerMetadata.getOptionalSegmentsMap().get("REALTIME")); + } + + @Test + public void testLogicalTableRequestsCarryOptionalSegmentsPerPhysicalTable() { + WorkerMetadata workerMetadata = new WorkerMetadata(0, Map.of()); + workerMetadata.setLogicalTableSegmentsMap( + Map.of("t1_OFFLINE", List.of("t1Seg"), "t2_REALTIME", List.of("t2Seg"))); + workerMetadata.setOptionalSegmentsMap(Map.of("t2_REALTIME", List.of("t2NewSeg"))); + LogicalTableContext logicalTableContext = mock(LogicalTableContext.class); + when(logicalTableContext.getRefOfflineTableConfig()).thenReturn(tableConfig(TableType.OFFLINE)); + when(logicalTableContext.getRefRealtimeTableConfig()).thenReturn(tableConfig(TableType.REALTIME)); + when(logicalTableContext.getLogicalTableSchema()).thenReturn(SCHEMA); + InstanceDataManager instanceDataManager = mock(InstanceDataManager.class); + when(instanceDataManager.getLogicalTableContext("tbl")).thenReturn(logicalTableContext); + + List<InstanceRequest> instanceRequests = ServerPlanRequestUtils.constructLogicalTableServerQueryRequests( + executionContext("tbl", workerMetadata), CalciteSqlParser.compileToPinotQuery("SELECT col1 FROM tbl"), + instanceDataManager); + + assertEquals(instanceRequests.size(), 2); + InstanceRequest offlineRequest = instanceRequests.get(0); + assertEquals(offlineRequest.getTableSegmentsInfoList().get(0).getSegments(), List.of("t1Seg")); + assertFalse(offlineRequest.getTableSegmentsInfoList().get(0).isSetOptionalSegments()); + InstanceRequest realtimeRequest = instanceRequests.get(1); + assertEquals(realtimeRequest.getTableSegmentsInfoList().get(0).getSegments(), List.of("t2Seg")); + assertEquals(realtimeRequest.getTableSegmentsInfoList().get(0).getOptionalSegments(), List.of("t2NewSeg")); + } + + private static TableConfig tableConfig(TableType tableType) { + return new TableConfigBuilder(tableType).setTableName("tbl").build(); + } + + private static OpChainExecutionContext executionContext(String tableName, WorkerMetadata workerMetadata) { + StageMetadata stageMetadata = new StageMetadata(1, List.of(workerMetadata), + Map.of(DispatchablePlanFragment.TABLE_NAME_KEY, tableName)); + long activeDeadlineMs = System.currentTimeMillis() + 10_000; + return new OpChainExecutionContext(mock(MailboxService.class), 1, "cid", activeDeadlineMs, + activeDeadlineMs + 1_000, "broker", Map.of(), stageMetadata, workerMetadata, null, false, false); + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/OptionalSegmentsQueryRunnerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/OptionalSegmentsQueryRunnerTest.java new file mode 100644 index 00000000000..a3ff4815bae --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/OptionalSegmentsQueryRunnerTest.java @@ -0,0 +1,122 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.queries; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.core.routing.MockRoutingManagerFactory; +import org.apache.pinot.query.QueryEnvironment; +import org.apache.pinot.query.QueryServerEnclosure; +import org.apache.pinot.query.mailbox.MailboxService; +import org.apache.pinot.query.routing.QueryServerInstance; +import org.apache.pinot.query.routing.WorkerManager; +import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory; +import org.apache.pinot.query.testutils.QueryTestUtils; +import org.apache.pinot.spi.config.instance.InstanceType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/// Runs queries over a table routed with optional segments, i.e. new segments the broker has not seen ONLINE yet. As on +/// the single-stage engine, the leaf stage must query the optional segments the server has, and skip without error the +/// ones it does not have yet. +public class OptionalSegmentsQueryRunnerTest extends QueryRunnerTestBase { + private static final String TABLE_NAME = "tbl"; + private static final String OFFLINE_TABLE_NAME = TABLE_NAME + "_OFFLINE"; + private static final int NUM_ROWS_PER_SEGMENT = 5; + + @BeforeClass + public void setUp() + throws Exception { + Schema schema = new Schema.SchemaBuilder().setSchemaName(TABLE_NAME) + .addSingleValueDimension("col1", FieldSpec.DataType.STRING, "") + .addDateTime("ts", FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:HOURS") + .build(); + MockInstanceDataManagerFactory instanceDataManagerFactory = new MockInstanceDataManagerFactory("server"); + instanceDataManagerFactory.registerTable(schema, OFFLINE_TABLE_NAME); + String segment = instanceDataManagerFactory.addSegment(OFFLINE_TABLE_NAME, buildRows()).getSegmentName(); + String loadedNewSegment = instanceDataManagerFactory.addSegment(OFFLINE_TABLE_NAME, buildRows()).getSegmentName(); + + _reducerHostname = "localhost"; + _reducerPort = QueryTestUtils.getAvailablePort(); + _mailboxService = new MailboxService(_reducerHostname, _reducerPort, InstanceType.BROKER, new PinotConfiguration( + Map.of(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME, _reducerHostname, + MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT, _reducerPort))); + _mailboxService.start(); + + QueryServerEnclosure server = new QueryServerEnclosure(instanceDataManagerFactory); + server.start(); + int port = server.getPort(); + _servers.put(new QueryServerInstance("Server_localhost_" + port, "localhost", port, port), server); + + MockRoutingManagerFactory routingManagerFactory = new MockRoutingManagerFactory(port); + routingManagerFactory.registerTable(schema, OFFLINE_TABLE_NAME); + routingManagerFactory.registerSegment(port, OFFLINE_TABLE_NAME, segment); + // Two new segments: one the server has already loaded, and one it has not. + routingManagerFactory.registerOptionalSegment(port, OFFLINE_TABLE_NAME, loadedNewSegment); + routingManagerFactory.registerOptionalSegment(port, OFFLINE_TABLE_NAME, OFFLINE_TABLE_NAME + "_notLoadedYet"); + _queryEnvironment = new QueryEnvironment(CommonConstants.DEFAULT_DATABASE, routingManagerFactory.buildTableCache(), + new WorkerManager("Broker_localhost", "localhost", _reducerPort, + routingManagerFactory.buildRoutingManager(null))); + } + + private static List<GenericRow> buildRows() { + List<GenericRow> rows = new ArrayList<>(NUM_ROWS_PER_SEGMENT); + for (int i = 0; i < NUM_ROWS_PER_SEGMENT; i++) { + GenericRow row = new GenericRow(); + row.putValue("col1", "value" + i); + row.putValue("ts", System.currentTimeMillis()); + rows.add(row); + } + return rows; + } + + @AfterClass + public void tearDown() { + for (QueryServerEnclosure server : _servers.values()) { + server.shutDown(); + } + _mailboxService.shutdown(); + } + + @DataProvider + public static Object[][] physicalOptimizer() { + return new Object[][]{{false}, {true}}; + } + + @Test(dataProvider = "physicalOptimizer") + public void testQueriesLoadedOptionalSegmentAndSkipsMissingOne(boolean usePhysicalOptimizer) { + ResultTable resultTable = queryRunner( + "SET usePhysicalOptimizer=" + usePhysicalOptimizer + "; SELECT COUNT(*) FROM " + TABLE_NAME, false) + .getResultTable(); + assertEquals(resultTable.getRows().get(0)[0], 2L * NUM_ROWS_PER_SEGMENT); + } +} diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java index f4df35b81b6..103b6ae8031 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/testutils/MockInstanceDataManagerFactory.java @@ -47,9 +47,9 @@ import org.apache.pinot.spi.utils.ReadMode; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyList; import static org.mockito.ArgumentMatchers.anyMap; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -158,10 +158,21 @@ public class MockInstanceDataManagerFactory { Map<String, SegmentDataManager> segmentDataManagerMap = segmentList.stream().collect(Collectors.toMap(IndexSegment::getSegmentName, ImmutableSegmentDataManager::new)); - // TODO: support optional segments for multi-stage engine, but for now, it's always null. - when(tableDataManager.acquireSegments(anyList(), eq(null), anyList())).thenAnswer(invocation -> { + // As in BaseTableDataManager, an optional segment the server does not have is skipped. + when(tableDataManager.acquireSegments(anyList(), any(), anyList())).thenAnswer(invocation -> { List<String> segments = invocation.getArgument(0); - return segments.stream().map(segmentDataManagerMap::get).collect(Collectors.toList()); + List<String> optionalSegments = invocation.getArgument(1); + List<SegmentDataManager> segmentDataManagers = + segments.stream().map(segmentDataManagerMap::get).collect(Collectors.toCollection(ArrayList::new)); + if (optionalSegments != null) { + for (String optionalSegment : optionalSegments) { + SegmentDataManager segmentDataManager = segmentDataManagerMap.get(optionalSegment); + if (segmentDataManager != null) { + segmentDataManagers.add(segmentDataManager); + } + } + } + return segmentDataManagers; }); when(tableDataManager.getSegmentContexts(anyList(), anyMap())).thenAnswer(invocation -> { List<IndexSegment> segments = invocation.getArgument(0); --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
