Copilot commented on code in PR #17532: URL: https://github.com/apache/pinot/pull/17532#discussion_r2708769538
########## pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/NPerDaySegmentsTableSampler.java: ########## @@ -0,0 +1,145 @@ +/** + * 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.broker.routing.tablesampler; + +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZonedDateTime; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.collections4.MapUtils; +import org.apache.helix.store.zk.ZkHelixPropertyStore; +import org.apache.helix.zookeeper.datamodel.ZNRecord; +import org.apache.pinot.common.metadata.ZKMetadataProvider; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig; +import org.apache.pinot.spi.utils.CommonConstants.Segment; + + +/** + * Selects up to N segments per day based on segment ZK metadata start time. + * + * <p>Config: + * <ul> + * <li>{@code properties.numSegmentsPerDay}: positive integer</li> + * <li>{@code properties.timezone}: optional timezone ID, defaults to {@code UTC}</li> + * </ul> + * + * <p>Notes: + * <ul> + * <li>Day bucketing is derived from the segment start time (in ms) converted to the configured timezone.</li> + * <li>Segments without valid start time metadata are included (not sampled) to avoid unexpectedly dropping data.</li> + * </ul> + */ +public class NPerDaySegmentsTableSampler implements TableSampler { + public static final String TYPE = "nPerDay"; + public static final String PROP_NUM_SEGMENTS_PER_DAY = "numSegmentsPerDay"; + public static final String PROP_TIMEZONE = "timezone"; + + private String _tableNameWithType; + private ZkHelixPropertyStore<ZNRecord> _propertyStore; + private int _numSegmentsPerDay; + private ZoneId _zoneId; + + @Override + public void init(String tableNameWithType, TableConfig tableConfig, TableSamplerConfig samplerConfig, + ZkHelixPropertyStore<ZNRecord> propertyStore) { + _tableNameWithType = tableNameWithType; + _propertyStore = propertyStore; + + Map<String, String> props = samplerConfig.getProperties(); + if (MapUtils.isEmpty(props) || !props.containsKey(PROP_NUM_SEGMENTS_PER_DAY)) { + throw new IllegalArgumentException( + "Missing required property '" + PROP_NUM_SEGMENTS_PER_DAY + "' for table sampler type '" + TYPE + "'"); + } + _numSegmentsPerDay = Integer.parseInt(props.get(PROP_NUM_SEGMENTS_PER_DAY)); + if (_numSegmentsPerDay <= 0) { + throw new IllegalArgumentException("'" + PROP_NUM_SEGMENTS_PER_DAY + "' must be positive"); + } + + String timezone = props.getOrDefault(PROP_TIMEZONE, "UTC"); + _zoneId = ZoneId.of(timezone); + } + + @Override + public Set<String> selectSegments(Set<String> onlineSegments) { + if (onlineSegments.isEmpty()) { + return Collections.emptySet(); + } + + Map<Long, List<String>> epochDayToSegments = new HashMap<>(); + List<String> unknownTimeSegments = new ArrayList<>(); + + for (String segmentName : onlineSegments) { + SegmentZKMetadata zkMetadata = + ZKMetadataProvider.getSegmentZKMetadata(_propertyStore, _tableNameWithType, segmentName); + if (zkMetadata == null) { + unknownTimeSegments.add(segmentName); + continue; + } + long startTimeMs = getStartTimeMs(zkMetadata); + if (startTimeMs < 0) { + unknownTimeSegments.add(segmentName); + continue; + } + ZonedDateTime zdt = Instant.ofEpochMilli(startTimeMs).atZone(_zoneId); + long epochDay = zdt.toLocalDate().toEpochDay(); + epochDayToSegments.computeIfAbsent(epochDay, k -> new ArrayList<>()).add(segmentName); + } + + Set<String> selected = new HashSet<>(); + for (List<String> segmentsForDay : epochDayToSegments.values()) { + Collections.sort(segmentsForDay); + int limit = Math.min(_numSegmentsPerDay, segmentsForDay.size()); + selected.addAll(segmentsForDay.subList(0, limit)); + } + + // Include segments without valid time metadata to avoid dropping them unexpectedly. + selected.addAll(unknownTimeSegments); + return selected; + } + + /** + * {@link SegmentZKMetadata#getStartTimeMs()} treats start time {@code 0} as invalid (returns {@code -1}). + * For table sampling, allow epoch-start segments when the raw start time is explicitly set to {@code 0} with a valid + * time unit. + */ + private static long getStartTimeMs(SegmentZKMetadata zkMetadata) { + long startTimeMs = zkMetadata.getStartTimeMs(); + if (startTimeMs > 0) { + return startTimeMs; + } + // Handle epoch-start segment (start time explicitly set to 0). + Map<String, String> simpleFields = zkMetadata.getSimpleFields(); + String startTimeString = simpleFields.get(Segment.START_TIME); + if ("0".equals(startTimeString)) { + String timeUnit = simpleFields.get(Segment.TIME_UNIT); + if (timeUnit != null) { + return 0L; + } + } Review Comment: The logic returns `0L` if `startTimeString` is \"0\" and `timeUnit` is non-null, but it does not validate that the time unit is actually valid or convertible to milliseconds. If the time unit is invalid, this may cause inconsistent behavior elsewhere. Consider validating the time unit value before returning 0L, or documenting that the caller must handle this. ########## pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java: ########## @@ -1087,10 +1222,31 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri @Override public List<String> getSegments(BrokerRequest brokerRequest) { String tableNameWithType = brokerRequest.getQuerySource().getTableName(); - RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType); + RoutingEntry routingEntry = getRoutingEntry(brokerRequest, tableNameWithType); return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null; } + @Nullable + private RoutingEntry getRoutingEntry(BrokerRequest brokerRequest, String tableNameWithType) { + String samplerName = null; + if (brokerRequest != null && brokerRequest.isSetPinotQuery()) { + org.apache.pinot.common.request.PinotQuery pinotQuery = brokerRequest.getPinotQuery(); + if (pinotQuery != null && pinotQuery.isSetQueryOptions()) { + samplerName = pinotQuery.getQueryOptions().get(CommonConstants.Broker.Request.QueryOptionKey.TABLE_SAMPLER); + } + } Review Comment: The method performs null checks on `brokerRequest`, `pinotQuery`, and query options, then extracts `samplerName`. This logic is duplicated in multiple places in WorkerManager.java. Consider extracting this into a utility method (e.g., `extractSamplerName(BrokerRequest)`) to reduce duplication and improve maintainability. ```suggestion private String extractSamplerName(@Nullable BrokerRequest brokerRequest) { if (brokerRequest == null || !brokerRequest.isSetPinotQuery()) { return null; } org.apache.pinot.common.request.PinotQuery pinotQuery = brokerRequest.getPinotQuery(); if (pinotQuery == null || !pinotQuery.isSetQueryOptions()) { return null; } return pinotQuery.getQueryOptions().get(CommonConstants.Broker.Request.QueryOptionKey.TABLE_SAMPLER); } @Nullable private RoutingEntry getRoutingEntry(BrokerRequest brokerRequest, String tableNameWithType) { String samplerName = extractSamplerName(brokerRequest); ``` ########## pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java: ########## @@ -820,6 +878,70 @@ private void buildRoutingInternal(String tableNameWithType) { LOGGER.info("Rebuilt routing for table: {}", tableNameWithType); } + // Build routing entries for configured table samplers (if any). + // These entries use the same underlying routing components, but operate on a deterministically sampled set of + // segments selected during routing build. + List<TableSamplerConfig> tableSamplerConfigs = tableConfig.getTableSamplers(); + if (tableSamplerConfigs != null && !tableSamplerConfigs.isEmpty()) { + Map<String, RoutingEntry> samplerRoutingEntries = new HashMap<>(); + for (TableSamplerConfig samplerConfig : tableSamplerConfigs) { + if (samplerConfig == null) { + continue; + } Review Comment: The loop continues if `samplerConfig` is null, but `tableSamplerConfigs` is a non-null list that may contain null elements. This suggests the list construction allows nulls. Consider validating and filtering null entries at the configuration level (in TableConfig or TableConfigSerDeUtils) to avoid defensive checks at runtime. ########## pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java: ########## @@ -539,40 +541,52 @@ private void assignWorkersToNonPartitionedLeafFragment(DispatchablePlanMetadata * @return keyed-map from table type(s) to routing table(s). */ private Map<String, RoutingTable> getRoutingTable(String tableName, long requestId) { + return getRoutingTable(tableName, requestId, Map.of()); + } + + private Map<String, RoutingTable> getRoutingTable(String tableName, long requestId, + Map<String, String> queryOptions) { TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableName); if (tableType == null) { // Raw table name Map<String, RoutingTable> routingTableMap = new HashMap<>(4); RoutingTable offlineRoutingTable = - getRoutingTableHelper(TableNameBuilder.OFFLINE.tableNameWithType(tableName), requestId); + getRoutingTableHelper(TableNameBuilder.OFFLINE.tableNameWithType(tableName), requestId, queryOptions); if (offlineRoutingTable != null) { routingTableMap.put(TableType.OFFLINE.name(), offlineRoutingTable); } RoutingTable realtimeRoutingTable = - getRoutingTableHelper(TableNameBuilder.REALTIME.tableNameWithType(tableName), requestId); + getRoutingTableHelper(TableNameBuilder.REALTIME.tableNameWithType(tableName), requestId, queryOptions); if (realtimeRoutingTable != null) { routingTableMap.put(TableType.REALTIME.name(), realtimeRoutingTable); } return routingTableMap; } else { // Table name with type - RoutingTable routingTable = getRoutingTableHelper(tableName, requestId); + RoutingTable routingTable = getRoutingTableHelper(tableName, requestId, queryOptions); return routingTable != null ? Map.of(tableType.name(), routingTable) : Map.of(); } } @Nullable - private RoutingTable getRoutingTableHelper(String tableNameWithType, long requestId) { - return _routingManager.getRoutingTable( - CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" + tableNameWithType + "\""), requestId); + private RoutingTable getRoutingTableHelper(String tableNameWithType, long requestId, + Map<String, String> queryOptions) { + BrokerRequest brokerRequest = + CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM \"" + tableNameWithType + "\""); + if (MapUtils.isNotEmpty(queryOptions) && brokerRequest.isSetPinotQuery()) { + // Ensure query options (e.g. tableSampler) are visible to routing selection. + brokerRequest.getPinotQuery().setQueryOptions(new HashMap<>(queryOptions)); + } Review Comment: Creating a new `HashMap` from `queryOptions` on every call may be unnecessary if `queryOptions` is already a mutable map. Consider passing `queryOptions` directly or documenting why a defensive copy is required here. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
