You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by "xiangfu0 (via GitHub)" <gi...@apache.org> on 2026/01/19 14:23:16 UTC
[PR] Add table sampler routing entries (precomputed segment subsets) with nPerDay sampler and tableSampler query option [pinot]
xiangfu0 opened a new pull request, #17532:
URL: https://github.com/apache/pinot/pull/17532
### Motivation
Large-table sampling needs to be deterministic and avoid query-time segment selection overhead. This adds a pluggable “table sampler” definition in table config and precomputes sampler-specific routing entries at the broker.
### Key changes
- Config: add `tableSamplers` to `TableConfig` (+ ZK SerDe support) and query option `tableSampler=<name>`.
- Broker routing:
- Build and cache sampler-specific routing entries per table sampler name.
- Select routing entry at query time by tableSampler option (fallback to default when absent/unknown).
- Refresh sampler routing entries on Helix assignment changes (IdealState/ExternalView updates).
- Built-in samplers:
- `firstN`: select first N segments (lexicographic)
- `nPerDay`: select up to N segments per day using segment ZK start time metadata
- MSQ support: propagate query options into MSQ leaf routing requests so `tableSampler` works with multi-stage engine.
- Tests:
- Unit test for `nPerDay`
- Integration test (shared cluster) validating 10 segments/day × 7 days → sampler returns 1 segment/day and group-by results reflect that
- Quickstart: add sample `tableSamplers` config to batch airlineStats table config.
### How to use
- Table config example:
- `tableSamplers: [{ name: "perDay", type: "nPerDay", properties: { numSegmentsPerDay: "1", timezone: "UTC" } }]`
- Query:
- `SET tableSampler=perDay; ...` or request `queryOptions: "tableSampler=perDay"`
### Compatibility
- Fully backward compatible: if no sampler is configured or selected, routing behavior is unchanged.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 merged PR #17532:
URL: https://github.com/apache/pinot/pull/17532
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add table sampler routing entries (precomputed segment subsets) with nPerDay sampler and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2713824794
##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -254,6 +263,30 @@ public void setCustomConfig(TableCustomConfig customConfig) {
_customConfig = customConfig;
}
+ @JsonProperty(TABLE_SAMPLERS_KEY)
+ @Nullable
+ public List<TableSamplerConfig> getTableSamplers() {
+ return _tableSamplers;
+ }
+
+ public void setTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ _tableSamplers = sanitizeTableSamplers(tableSamplers);
+ }
+
+ @Nullable
+ private static List<TableSamplerConfig> sanitizeTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ if (tableSamplers == null || tableSamplers.isEmpty()) {
+ return tableSamplers;
+ }
+ List<TableSamplerConfig> sanitized = new ArrayList<>(tableSamplers.size());
+ for (TableSamplerConfig config : tableSamplers) {
+ if (config != null) {
+ sanitized.add(config);
+ }
+ }
+ return sanitized;
Review Comment:
The `sanitizeTableSamplers` method removes null entries but returns the original list when empty, which could preserve an empty list instead of returning null consistently. Consider returning `null` when the sanitized list is empty to maintain consistency with the null-handling pattern, or document why empty lists should be preserved differently from null inputs.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1221,35 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
+ }
+ return _routingEntryMap.get(tableNameWithType);
+ }
+
+ @Nullable
+ private static 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()) {
Review Comment:
The method checks `brokerRequest.isSetPinotQuery()` but then also checks if `pinotQuery == null` on line 1247. If `isSetPinotQuery()` returns true, the subsequent null check is redundant. Consider removing the null check on line 1247 or clarifying if the Thrift `isSet` method can return true with a null value in this context.
```suggestion
if (!pinotQuery.isSetQueryOptions()) {
```
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2773559266
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TableSamplerIntegrationTest.java:
##########
@@ -0,0 +1,140 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.SchemaBuilder;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class TableSamplerIntegrationTest extends CustomDataQueryClusterIntegrationTest {
+ private static final int DAYS = 7;
+ private static final int SEGMENTS_PER_DAY = 10;
+ private static final int RECORDS_PER_SEGMENT = 1;
+ private static final int BASE_DAY = 20000;
+
+ private static final String DAYS_SINCE_EPOCH_COL = "DaysSinceEpoch";
+
+ @Override
+ public String getTableName() {
+ return "TableSamplerIntegrationTest";
+ }
+
+ @Override
+ protected long getCountStarResult() {
+ return (long) DAYS * SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT;
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder().setSchemaName(getTableName())
+ .addDateTime(DAYS_SINCE_EPOCH_COL, FieldSpec.DataType.INT, "1:DAYS:EPOCH", "1:DAYS")
+ .build();
+ }
+
+ @Override
+ public TableConfig createOfflineTableConfig() {
+ TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName())
+ .setTimeColumnName(DAYS_SINCE_EPOCH_COL)
+ .setTimeType("DAYS")
+ .build();
+ tableConfig.setTableSamplers(List.of(
+ new TableSamplerConfig("firstOnly", "firstN", Map.of("numSegments", "1")),
+ new TableSamplerConfig("firstTwo", "firstN", Map.of("numSegments", "2"))));
+ return tableConfig;
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws Exception {
+ var fieldAssembler = SchemaBuilder.record("myRecord").fields();
+ fieldAssembler.name(DAYS_SINCE_EPOCH_COL).type().intType().noDefault();
+ var avroSchema = fieldAssembler.endRecord();
+
+ List<File> files = new ArrayList<>();
+ for (int day = 0; day < DAYS; day++) {
+ int dayValue = BASE_DAY + day;
+ for (int seg = 0; seg < SEGMENTS_PER_DAY; seg++) {
+ File avroFile = new File(_tempDir, "data_day_" + day + "_seg_" + seg + ".avro");
+ try (DataFileWriter<GenericData.Record> fileWriter =
+ new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) {
+ fileWriter.create(avroSchema, avroFile);
+ for (int docId = 0; docId < RECORDS_PER_SEGMENT; docId++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ record.put(DAYS_SINCE_EPOCH_COL, dayValue);
+ fileWriter.append(record);
+ }
+ }
+ files.add(avroFile);
+ }
+ }
+ return files;
+ }
+
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFirstNSamplerForGroupByDay(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+
+ JsonNode full = postQuery("SELECT DaysSinceEpoch, COUNT(*) AS cnt FROM " + getTableName()
+ + " GROUP BY DaysSinceEpoch ORDER BY DaysSinceEpoch");
+ JsonNode fullRows = full.path("resultTable").path("rows");
+ Assert.assertEquals(fullRows.size(), DAYS);
+ for (int i = 0; i < DAYS; i++) {
+ Assert.assertEquals(fullRows.get(i).get(0).asInt(), BASE_DAY + i);
+ Assert.assertEquals(fullRows.get(i).get(1).asLong(), (long) SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT);
+ }
+ Assert.assertEquals(Integer.parseInt(full.path("numSegmentsQueried").asText()), DAYS * SEGMENTS_PER_DAY);
Review Comment:
Parsing the numSegmentsQueried field as text and then converting to int is fragile. Use asInt() directly if the field is numeric, or handle the case where the field might not be present or numeric.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1309,59 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
+ if (MissingSamplerWarningTracker.shouldLogWarn(tableNameWithType, samplerName)) {
+ LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
+ samplerName, tableNameWithType);
+ } else if (LOGGER.isDebugEnabled()) {
+ LOGGER.debug("Requested sampler '{}' not found for table '{}'; falling back to default routing entry "
+ + "(subsequent occurrences logged at DEBUG level only)", samplerName, tableNameWithType);
Review Comment:
The debug log message mentions 'subsequent occurrences' but this is actually the first occurrence that didn't log at WARN level. The message should clarify that the WARN was already logged once and this is a subsequent occurrence.
```suggestion
+ "(WARN already logged once for this sampler/table; subsequent occurrences logged at DEBUG level only)",
samplerName, tableNameWithType);
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSamplerFactory.java:
##########
@@ -0,0 +1,147 @@
+/**
+ * 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 com.google.common.annotations.VisibleForTesting;
+import java.lang.reflect.Modifier;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.annotations.tablesampler.TableSamplerProvider;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.plugin.PluginManager;
+import org.apache.pinot.spi.utils.PinotReflectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class TableSamplerFactory {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TableSamplerFactory.class);
+ private static final String ANNOTATION_PACKAGES_KEY = "annotation.packages";
+ // Keep this list in sync with built-in TableSampler locations; additional packages can be configured via broker
+ // config, but defaults should always include built-ins.
+ private static final List<String> DEFAULT_ANNOTATION_PACKAGES =
+ List.of("org.apache.pinot.broker.routing.tablesampler");
+ private static final Map<String, String> REGISTERED_TABLE_SAMPLER_CLASS_MAP = new ConcurrentHashMap<>();
+
+ private TableSamplerFactory() {
+ }
+
+ public static void init(PinotConfiguration tableSamplerConfig) {
+ if (tableSamplerConfig == null) {
+ return;
+ }
+ registerAnnotatedTableSamplers(tableSamplerConfig);
+ }
+
+ public static void register(String alias, String className) {
+ if (StringUtils.isBlank(alias)) {
+ LOGGER.warn("Skipping table sampler registration because alias is blank");
+ return;
+ }
+ if (StringUtils.isBlank(className)) {
+ LOGGER.warn("Skipping table sampler registration for alias '{}' because class name is blank", alias);
+ return;
+ }
+ String normalizedAlias = normalizeType(alias);
+ String trimmedClassName = className.trim();
+ String previousClassName = REGISTERED_TABLE_SAMPLER_CLASS_MAP.put(normalizedAlias, trimmedClassName);
+ if (previousClassName == null) {
+ LOGGER.info("Registered table sampler alias '{}' -> '{}'", alias, trimmedClassName);
+ } else if (!previousClassName.equals(trimmedClassName)) {
+ LOGGER.warn("Overriding table sampler alias '{}' from '{}' to '{}'", alias, previousClassName, trimmedClassName);
+ }
+ }
+
+ public static TableSampler create(String type) {
+ String resolvedClassName = resolveClassName(type);
+ String classNameToLoad = resolvedClassName != null ? resolvedClassName : type;
+ try {
+ return PluginManager.get().createInstance(classNameToLoad);
+ } catch (Exception e) {
+ String errorMessage = resolvedClassName != null
+ ? String.format("Failed to create TableSampler for alias '%s' mapped to class '%s'", type,
+ resolvedClassName)
+ : "Failed to create TableSampler for type: " + type;
+ throw new RuntimeException(errorMessage, e);
+ }
+ }
+
+ @VisibleForTesting
+ static void clearRegistry() {
+ REGISTERED_TABLE_SAMPLER_CLASS_MAP.clear();
+ }
+
+ private static void registerAnnotatedTableSamplers(PinotConfiguration tableSamplerConfig) {
+ List<String> configuredPackages = tableSamplerConfig.getProperty(ANNOTATION_PACKAGES_KEY, List.of());
Review Comment:
The getProperty method is being called with a List.of() default, but PinotConfiguration.getProperty returns a String or collection based on the actual stored type. If the property is not set, this could return an empty list, but if it's set as a comma-separated string, it won't be properly parsed into a list. Use the appropriate method for retrieving list properties or handle string-to-list conversion.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add table sampler routing entries (precomputed segment subsets) with nPerDay sampler and tableSampler query option [pinot]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #17532:
URL: https://github.com/apache/pinot/pull/17532#issuecomment-3768822591
## [Codecov](https://app.codecov.io/gh/apache/pinot/pull/17532?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `54.34783%` with `21 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 55.51%. Comparing base ([`c899956`](https://app.codecov.io/gh/apache/pinot/commit/c899956b2d084dcad7aa4e7cd5515edb69f0ac9d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`19f856b`](https://app.codecov.io/gh/apache/pinot/commit/19f856b6c6322faece1c1dd1127f3d2fcc583e61?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
| [Files with missing lines](https://app.codecov.io/gh/apache/pinot/pull/17532?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...t/spi/config/table/sampler/TableSamplerConfig.java](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&filepath=pinot-spi%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fspi%2Fconfig%2Ftable%2Fsampler%2FTableSamplerConfig.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvY29uZmlnL3RhYmxlL3NhbXBsZXIvVGFibGVTYW1wbGVyQ29uZmlnLmphdmE=) | 0.00% | [8 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../org/apache/pinot/query/routing/WorkerManager.java](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&filepath=pinot-query-planner%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fquery%2Frouting%2FWorkerManager.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtcXVlcnktcGxhbm5lci9zcmMvbWFpbi9qYXZhL29yZy9hcGFjaGUvcGlub3QvcXVlcnkvcm91dGluZy9Xb3JrZXJNYW5hZ2VyLmphdmE=) | 78.26% | [3 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...not/common/utils/config/TableConfigSerDeUtils.java](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&filepath=pinot-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fcommon%2Futils%2Fconfig%2FTableConfigSerDeUtils.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3QtY29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9jb21tb24vdXRpbHMvY29uZmlnL1RhYmxlQ29uZmlnU2VyRGVVdGlscy5qYXZh) | 42.85% | [2 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...org/apache/pinot/spi/config/table/TableConfig.java](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&filepath=pinot-spi%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fspi%2Fconfig%2Ftable%2FTableConfig.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvY29uZmlnL3RhYmxlL1RhYmxlQ29uZmlnLmphdmE=) | 66.66% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...he/pinot/spi/utils/builder/TableConfigBuilder.java](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&filepath=pinot-spi%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fpinot%2Fspi%2Futils%2Fbuilder%2FTableConfigBuilder.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-cGlub3Qtc3BpL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9waW5vdC9zcGkvdXRpbHMvYnVpbGRlci9UYWJsZUNvbmZpZ0J1aWxkZXIuamF2YQ==) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/pinot/pull/17532?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
> :exclamation: There is a different number of reports uploaded between BASE (c899956) and HEAD (19f856b). Click for more details.
>
> <details><summary>HEAD has 32 uploads less than BASE</summary>
>
>| Flag | BASE (c899956) | HEAD (19f856b) |
>|------|------|------|
>|java-21|4|1|
>|unittests1|2|1|
>|unittests|3|1|
>|temurin|9|1|
>|java-11|5|0|
>|integration|6|0|
>|integration2|2|0|
>|integration1|2|0|
>|unittests2|1|0|
>|custom-integration1|2|0|
></details>
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #17532 +/- ##
============================================
- Coverage 63.21% 55.51% -7.70%
+ Complexity 1476 702 -774
============================================
Files 3170 2474 -696
Lines 189508 139856 -49652
Branches 28997 22295 -6702
============================================
- Hits 119789 77642 -42147
+ Misses 60417 55657 -4760
+ Partials 9302 6557 -2745
```
| [Flag](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [custom-integration1](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration1](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [integration2](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-11](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
| [java-21](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `55.51% <54.34%> (-0.01%)` | :arrow_down: |
| [temurin](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `55.51% <54.34%> (-7.70%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `55.51% <54.34%> (-7.70%)` | :arrow_down: |
| [unittests1](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `55.51% <54.34%> (-0.04%)` | :arrow_down: |
| [unittests2](https://app.codecov.io/gh/apache/pinot/pull/17532/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `?` | |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
</details>
[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/pinot/pull/17532?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2796334155
##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java:
##########
@@ -68,6 +68,8 @@ private QueryOptionsUtils() {
}
}
}
+ // Backward-compatible alias used by SQL SET syntax: SET sampler='x'
+ configResolver.put("sampler", QueryOptionKey.TABLE_SAMPLER);
Review Comment:
SQL is using SAMPLE at run time to pick rows, our concept is bit different, making query option `sampler` to match it.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2776716275
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -820,6 +971,69 @@ 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()) {
Review Comment:
(minor) Use `CollectionUtils.isNotEmpty()`
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSampler.java:
##########
@@ -0,0 +1,46 @@
+/**
+ * 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.util.Set;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+
+
+/**
+ * A {@code TableSampler} deterministically selects a subset of segments from the set of online segments for a table.
+ *
+ * <p>Selection is performed during routing table build/update so there is no additional per-query overhead beyond
+ * selecting the pre-built routing entry.
+ */
+public interface TableSampler {
+
+ /**
+ * Initializes the sampler for a specific table and sampler config.
+ */
+ void init(String tableNameWithType, TableConfig tableConfig, TableSamplerConfig samplerConfig,
Review Comment:
No need to pass in `tableNameWithType`. Use `TableConfig.getTableNameI()`
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -143,6 +152,9 @@ public abstract class BaseBrokerRoutingManager implements RoutingManager, Cluste
// and OFFLINE tables. LogicalTables will also use the raw table name of the table list underneath
private final Map<String, Object> _routingTableBuildLocks = new ConcurrentHashMap<>();
+ // Per-table locks for sampler routing entry updates.
+ private final Map<String, Object> _routingTableSamplerBuildLocks = new ConcurrentHashMap<>();
Review Comment:
We shouldn't need this extra lock. When updating the routing for a given table, we should grab the routing table build lock for the table, update both the routing entry, and also the sampler routing entries if exist
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -820,6 +971,69 @@ 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) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableNameWithType, tableConfig, samplerConfig, _propertyStore);
+
+ SegmentPreSelector samplerSegmentPreSelector =
Review Comment:
The sampler shouldn't be modeled as a pre-selector. It should be applied after pre-selection, where lineage is handled, or it can give wrong result.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -820,6 +971,69 @@ 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) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableNameWithType, tableConfig, samplerConfig, _propertyStore);
+
+ SegmentPreSelector samplerSegmentPreSelector =
+ new TableSamplerSegmentPreSelector(segmentPreSelector, sampler);
+ Set<String> samplerPreSelectedOnlineSegments =
+ // SegmentPreSelector contract allows mutation of the input set. Create a defensive copy so the
+ // sampler cannot mutate the shared onlineSegments set, which is reused by other samplers.
+ samplerSegmentPreSelector.preSelect(new HashSet<>(onlineSegments));
Review Comment:
This should be applied to `preSelectedOnlineSegments`.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -261,25 +355,33 @@ protected void processSegmentAssignmentChangeInternal() {
if (idealStateStat != null && externalViewStat != null) {
String tableNameWithType = cachedTableNames.get(i);
futures.add(_executorService.submit(() -> {
+ AssignmentChangeResult assignmentChangeResult;
+ RoutingEntry routingEntry;
Object tableLock = getRoutingTableBuildLock(tableNameWithType);
synchronized (tableLock) {
// The routingEntry may have been removed from the _routingEntryMap by the time we get here in case
// one of the other functions such as 'removeRouting' was called since taking the snapshot. Check for
// existence before proceeding. Also note that if new entries were added since the snapshot was taken, we
// will miss processing them in this call. The buildRouting() method tries to handle that by checking for
// changes in the IS / EV version after adding a new entry
- RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
+ routingEntry = _routingEntryMap.get(tableNameWithType);
if (routingEntry == null) {
LOGGER.info("Table {} was removed while processing segment assignment change, skipping",
tableNameWithType);
return;
}
- boolean hasISOrEVVersionChanged = processAssignmentChangeForTable(idealStateStat.getVersion(),
- externalViewStat.getVersion(), routingEntry);
- if (hasISOrEVVersionChanged) {
+ assignmentChangeResult =
+ processAssignmentChangeForTable(idealStateStat.getVersion(), externalViewStat.getVersion(),
+ routingEntry);
+ if (assignmentChangeResult.hasVersionChanged()) {
tablesUpdated.add(tableNameWithType);
}
}
+ // Avoid reusing the per-table lock for sampler routing update so sampler updates can proceed in parallel.
Review Comment:
I don't follow this. Wouldn't this cause race condition? Routing entry is not protected
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2796354576
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,9 +1146,30 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
- return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
+ if (routingEntry == null) {
+ return null;
+ }
+ return routingEntry.getSegments(brokerRequest, samplerName);
+ }
+
+ private static String normalizeSamplerName(String samplerName) {
+ return samplerName.trim().toLowerCase(Locale.ROOT);
+ }
+
+ @Nullable
+ static String extractSamplerName(BrokerRequest brokerRequest) {
+ if (!brokerRequest.isSetPinotQuery() || !brokerRequest.getPinotQuery().isSetQueryOptions()) {
+ return null;
+ }
Review Comment:
Do you still remember which tests are failing? All `BrokerRequest` should have `PinotQuery`
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2762026991
##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -254,6 +263,33 @@ public void setCustomConfig(TableCustomConfig customConfig) {
_customConfig = customConfig;
}
+ @JsonProperty(TABLE_SAMPLERS_KEY)
+ @Nullable
+ public List<TableSamplerConfig> getTableSamplers() {
+ return _tableSamplers;
+ }
+
+ public void setTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ _tableSamplers = sanitizeTableSamplers(tableSamplers);
+ }
+
+ @Nullable
+ private static List<TableSamplerConfig> sanitizeTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ if (tableSamplers == null || tableSamplers.isEmpty()) {
+ return null;
+ }
+ List<TableSamplerConfig> sanitized = new ArrayList<>(tableSamplers.size());
+ for (TableSamplerConfig config : tableSamplers) {
+ if (config != null) {
+ sanitized.add(config);
+ }
+ }
+ if (sanitized.isEmpty()) {
+ return null;
+ }
+ return sanitized;
+ }
Review Comment:
The sanitizeTableSamplers method silently filters out null entries but doesn't validate for duplicate sampler names. If multiple samplers with the same name are configured, the last one will effectively override earlier ones during routing entry creation (line 991 in BaseBrokerRoutingManager uses a HashMap keyed by sampler name). Consider adding validation to detect and reject duplicate sampler names during sanitization.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -993,9 +1178,9 @@ public void refreshSegment(String tableNameWithType, String segment) {
}
private void refreshSegmentInternal(String tableNameWithType, String segment) {
+ LOGGER.info("Refreshing segment: {} for table: {}", segment, tableNameWithType);
Object tableLock = getRoutingTableBuildLock(tableNameWithType);
synchronized (tableLock) {
Review Comment:
The log statement at line 1181 was moved outside the synchronized block. This ordering change means the log will execute before acquiring the lock, which could lead to slightly misleading timestamps in logs if there's lock contention. Consider whether the log should remain inside the synchronized block for accurate timing, or add a note explaining the intentional ordering.
```suggestion
Object tableLock = getRoutingTableBuildLock(tableNameWithType);
synchronized (tableLock) {
LOGGER.info("Refreshing segment: {} for table: {}", segment, tableNameWithType);
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TimeBucketSegmentsTableSampler.java:
##########
@@ -0,0 +1,337 @@
+/**
+ * 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.LocalDateTime;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeParseException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.TimeUnit;
+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.SegmentPartitionMetadata;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.utils.LLCSegmentName;
+import org.apache.pinot.common.utils.UploadedRealtimeSegmentName;
+import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
+import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
+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;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+
+/**
+ * Selects up to N segments per time bucket, where the bucket size is configurable in days or hours.
+ *
+ * <p>Config:
+ * <ul>
+ * <li>{@code properties.numSegmentsPerDay}: positive integer (required unless numSegmentsPerHour is set). Note:
+ * interpreted as "per bucket".</li>
+ * <li>{@code properties.numSegmentsPerHour}: positive integer (optional, interpreted as "per bucket").</li>
+ * <li>{@code properties.bucketDays}: positive integer (optional, default 1)</li>
+ * <li>{@code properties.bucketHours}: positive integer (optional, mutually exclusive with bucketDays)</li>
+ * </ul>
+ *
+ * <p>Notes:
+ * <ul>
+ * <li>Bucket boundaries are computed in UTC.</li>
+ * <li>If the table has a single partition column configured, selection is done per partition within each time
+ * bucket (one partition id per segment required).</li>
+ * <li>For OFFLINE tables, bucketing uses segment end time from ZK metadata.</li>
+ * <li>For REALTIME tables, bucketing tries to derive a timestamp from the segment name (LLC / uploaded realtime)
+ * to avoid ZK reads.</li>
+ * <li>Segments with unparsable/missing timestamps are skipped.</li>
+ * <li>Selection is deterministic: within each bucket, choose lexicographically first N segment names.</li>
+ * </ul>
+ */
+public class TimeBucketSegmentsTableSampler implements TableSampler {
+ public static final String TYPE = "timeBucket";
+
+ public static final String PROP_NUM_SEGMENTS = "numSegmentsPerDay";
Review Comment:
The property name 'numSegmentsPerDay' is misleading since it's actually used as 'numSegments per bucket' where buckets can be configured as days or hours. The javadoc correctly notes this at line 57, but the constant name creates confusion. Consider renaming to PROP_NUM_SEGMENTS_PER_BUCKET or updating the documentation to make this dual usage clearer.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1277,35 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
Review Comment:
When a user specifies a non-existent sampler name in the query option, the code silently falls back to the default routing entry (line 1294). This could lead to confusion where queries appear to succeed but aren't actually using the intended sampler. Consider logging a warning when a requested sampler name is not found to aid debugging.
```suggestion
}
// Log a warning when a requested sampler is not found and we fall back to the default routing entry.
LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
samplerName, tableNameWithType);
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TimeBucketSegmentsTableSampler.java:
##########
@@ -0,0 +1,337 @@
+/**
+ * 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.LocalDateTime;
+import java.time.ZoneOffset;
+import java.time.ZonedDateTime;
+import java.time.format.DateTimeFormatter;
+import java.time.format.DateTimeParseException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NavigableSet;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.concurrent.TimeUnit;
+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.SegmentPartitionMetadata;
+import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
+import org.apache.pinot.common.utils.LLCSegmentName;
+import org.apache.pinot.common.utils.UploadedRealtimeSegmentName;
+import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
+import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
+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;
+import org.apache.pinot.spi.utils.builder.TableNameBuilder;
+
+/**
+ * Selects up to N segments per time bucket, where the bucket size is configurable in days or hours.
+ *
+ * <p>Config:
+ * <ul>
+ * <li>{@code properties.numSegmentsPerDay}: positive integer (required unless numSegmentsPerHour is set). Note:
+ * interpreted as "per bucket".</li>
+ * <li>{@code properties.numSegmentsPerHour}: positive integer (optional, interpreted as "per bucket").</li>
+ * <li>{@code properties.bucketDays}: positive integer (optional, default 1)</li>
+ * <li>{@code properties.bucketHours}: positive integer (optional, mutually exclusive with bucketDays)</li>
+ * </ul>
+ *
+ * <p>Notes:
+ * <ul>
+ * <li>Bucket boundaries are computed in UTC.</li>
+ * <li>If the table has a single partition column configured, selection is done per partition within each time
+ * bucket (one partition id per segment required).</li>
+ * <li>For OFFLINE tables, bucketing uses segment end time from ZK metadata.</li>
+ * <li>For REALTIME tables, bucketing tries to derive a timestamp from the segment name (LLC / uploaded realtime)
+ * to avoid ZK reads.</li>
+ * <li>Segments with unparsable/missing timestamps are skipped.</li>
+ * <li>Selection is deterministic: within each bucket, choose lexicographically first N segment names.</li>
+ * </ul>
+ */
+public class TimeBucketSegmentsTableSampler implements TableSampler {
+ public static final String TYPE = "timeBucket";
+
+ public static final String PROP_NUM_SEGMENTS = "numSegmentsPerDay";
+ public static final String PROP_NUM_SEGMENTS_PER_HOUR = "numSegmentsPerHour";
+ public static final String PROP_BUCKET_DAYS = "bucketDays";
+ public static final String PROP_BUCKET_HOURS = "bucketHours";
+
+ private static final DateTimeFormatter REALTIME_SEGMENT_NAME_TIME_FORMATTER =
+ DateTimeFormatter.ofPattern("yyyyMMdd'T'HHmm'Z'");
+
+ private String _tableNameWithType;
+ private ZkHelixPropertyStore<ZNRecord> _propertyStore;
+ private int _numSegmentsPerBucket;
+ private int _bucketDays;
+ private int _bucketHours;
+ private boolean _useHourBuckets;
+ private boolean _usePartitionAware;
+ private String _partitionColumn;
+ private int _numPartitions;
+ private boolean _isRealtimeTable;
+
+ // Cache bucket->segments to avoid re-reading ZK metadata for unchanged segments.
+ private static final int DEFAULT_PARTITION_KEY = 0;
+ private final Map<Long, Map<Integer, NavigableSet<String>>> _bucketIdToPartitionSegments = new HashMap<>();
+ private final Map<String, Long> _segmentToBucketId = new HashMap<>();
+ private final Map<String, Integer> _segmentToPartitionId = new HashMap<>();
+
+ @Override
+ public void init(String tableNameWithType, TableConfig tableConfig, TableSamplerConfig samplerConfig,
+ ZkHelixPropertyStore<ZNRecord> propertyStore) {
+ _tableNameWithType = tableNameWithType;
+ _propertyStore = propertyStore;
+ _isRealtimeTable = TableNameBuilder.isRealtimeTableResource(tableNameWithType);
+
+ Map<String, String> props = samplerConfig.getProperties();
+ if (MapUtils.isEmpty(props)
+ || (!props.containsKey(PROP_NUM_SEGMENTS) && !props.containsKey(PROP_NUM_SEGMENTS_PER_HOUR))) {
+ throw new IllegalArgumentException("Missing required property '" + PROP_NUM_SEGMENTS + "' or '"
+ + PROP_NUM_SEGMENTS_PER_HOUR + "' for table sampler type '" + TYPE + "'");
+ }
+
+ boolean hasNumSegmentsPerDay = props.containsKey(PROP_NUM_SEGMENTS);
+ boolean hasNumSegmentsPerHour = props.containsKey(PROP_NUM_SEGMENTS_PER_HOUR);
+ if (hasNumSegmentsPerDay && hasNumSegmentsPerHour) {
+ throw new IllegalArgumentException("'" + PROP_NUM_SEGMENTS + "' and '" + PROP_NUM_SEGMENTS_PER_HOUR
+ + "' are mutually exclusive");
+ }
+ String numSegmentsStr =
+ hasNumSegmentsPerHour ? props.get(PROP_NUM_SEGMENTS_PER_HOUR) : props.get(PROP_NUM_SEGMENTS);
+ _numSegmentsPerBucket = Integer.parseInt(numSegmentsStr);
+ if (_numSegmentsPerBucket <= 0) {
+ throw new IllegalArgumentException("numSegments per bucket must be positive");
+ }
+
+ boolean hasBucketDays = props.containsKey(PROP_BUCKET_DAYS);
+ boolean hasBucketHours = props.containsKey(PROP_BUCKET_HOURS);
+ if (hasBucketDays && hasBucketHours) {
+ throw new IllegalArgumentException("'" + PROP_BUCKET_DAYS + "' and '" + PROP_BUCKET_HOURS
+ + "' are mutually exclusive");
+ }
+
+ if (hasNumSegmentsPerHour || hasBucketHours) {
+ _bucketHours = Integer.parseInt(props.getOrDefault(PROP_BUCKET_HOURS, "1"));
+ if (_bucketHours <= 0) {
+ throw new IllegalArgumentException("'" + PROP_BUCKET_HOURS + "' must be positive");
+ }
+ _useHourBuckets = true;
+ } else {
+ String bucketDaysStr = props.getOrDefault(PROP_BUCKET_DAYS, "1");
+ _bucketDays = Integer.parseInt(bucketDaysStr);
+ if (_bucketDays <= 0) {
+ throw new IllegalArgumentException("'" + PROP_BUCKET_DAYS + "' must be positive");
+ }
+ }
+
+ SegmentPartitionConfig segmentPartitionConfig = tableConfig.getIndexingConfig() != null
+ ? tableConfig.getIndexingConfig().getSegmentPartitionConfig() : null;
+ if (segmentPartitionConfig != null && segmentPartitionConfig.getColumnPartitionMap() != null
+ && !segmentPartitionConfig.getColumnPartitionMap().isEmpty()) {
+ if (segmentPartitionConfig.getColumnPartitionMap().size() != 1) {
+ throw new IllegalArgumentException("Time bucket sampler supports only a single partition column");
+ }
+ Map.Entry<String, ColumnPartitionConfig> entry =
+ segmentPartitionConfig.getColumnPartitionMap().entrySet().iterator().next();
+ _partitionColumn = entry.getKey();
+ _numPartitions = entry.getValue().getNumPartitions();
+ _usePartitionAware = true;
+ } else {
+ _usePartitionAware = false;
+ }
+ }
+
+ @Override
+ public Set<String> selectSegments(Set<String> onlineSegments) {
+ if (onlineSegments.isEmpty()) {
+ return Collections.emptySet();
+ }
+
+ // Remove segments that are no longer online.
+ Iterator<Map.Entry<String, Long>> iterator = _segmentToBucketId.entrySet().iterator();
+ while (iterator.hasNext()) {
+ Map.Entry<String, Long> entry = iterator.next();
+ String segmentName = entry.getKey();
+ if (!onlineSegments.contains(segmentName)) {
+ Long bucketId = entry.getValue();
+ Integer partitionId = _segmentToPartitionId.get(segmentName);
+ if (partitionId != null) {
+ Map<Integer, NavigableSet<String>> partitionsForBucket = _bucketIdToPartitionSegments.get(bucketId);
+ if (partitionsForBucket != null) {
+ NavigableSet<String> segmentsForPartition = partitionsForBucket.get(partitionId);
+ if (segmentsForPartition != null) {
+ segmentsForPartition.remove(segmentName);
+ if (segmentsForPartition.isEmpty()) {
+ partitionsForBucket.remove(partitionId);
+ }
+ }
+ if (partitionsForBucket.isEmpty()) {
+ _bucketIdToPartitionSegments.remove(bucketId);
+ }
+ }
+ }
+ _segmentToPartitionId.remove(segmentName);
+ iterator.remove();
+ }
+ }
+
+ // Add new segments (only these may require metadata lookups).
+ for (String segmentName : onlineSegments) {
+ if (_segmentToBucketId.containsKey(segmentName)) {
+ continue;
+ }
+ long bucketId = getBucketIdForSegment(segmentName);
+ if (bucketId < 0) {
+ // Skip segments without parsable bucket id.
+ continue;
+ }
+ int partitionId = getPartitionIdForSegment(segmentName);
+ if (partitionId < 0) {
+ // Skip segments without valid partition id for partition-aware sampling.
+ continue;
+ }
+ _segmentToBucketId.put(segmentName, bucketId);
+ _segmentToPartitionId.put(segmentName, partitionId);
+ _bucketIdToPartitionSegments
+ .computeIfAbsent(bucketId, k -> new HashMap<>())
+ .computeIfAbsent(partitionId, k -> new TreeSet<>())
+ .add(segmentName);
+ }
+
+ Set<String> selected = new HashSet<>();
+ for (Map<Integer, NavigableSet<String>> partitionsForBucket : _bucketIdToPartitionSegments.values()) {
+ for (NavigableSet<String> segmentsForPartition : partitionsForBucket.values()) {
+ int count = 0;
+ for (String segmentName : segmentsForPartition) {
+ selected.add(segmentName);
+ if (++count >= _numSegmentsPerBucket) {
+ break;
+ }
+ }
+ }
+ }
+ return selected;
+ }
+
+ private long getBucketId(long timeMs) {
+ if (_useHourBuckets) {
+ long epochHour = TimeUnit.MILLISECONDS.toHours(timeMs);
+ return epochHour / _bucketHours;
+ }
+ ZonedDateTime zdt = Instant.ofEpochMilli(timeMs).atZone(ZoneOffset.UTC);
+ long epochDay = zdt.toLocalDate().toEpochDay();
+ return epochDay / _bucketDays;
+ }
+
+ private long getSegmentTimeMs(String segmentName) {
+ if (_isRealtimeTable) {
+ Long timeMs = getTimeMsFromRealtimeSegmentName(segmentName);
+ return timeMs != null ? timeMs : -1L;
+ }
+
+ SegmentZKMetadata zkMetadata = ZKMetadataProvider.getSegmentZKMetadata(_propertyStore, _tableNameWithType,
+ segmentName);
+ if (zkMetadata == null) {
+ return -1L;
+ }
+ return getEndTimeMs(zkMetadata);
+ }
+
+ private long getBucketIdForSegment(String segmentName) {
+ long timeMs = getSegmentTimeMs(segmentName);
+ return timeMs < 0 ? -1L : getBucketId(timeMs);
+ }
+
+ private int getPartitionIdForSegment(String segmentName) {
+ if (!_usePartitionAware) {
+ return DEFAULT_PARTITION_KEY;
+ }
+ SegmentZKMetadata zkMetadata = ZKMetadataProvider.getSegmentZKMetadata(_propertyStore, _tableNameWithType,
+ segmentName);
+ if (zkMetadata == null) {
+ return -1;
+ }
+ SegmentPartitionMetadata partitionMetadata = zkMetadata.getPartitionMetadata();
+ if (partitionMetadata == null) {
+ return -1;
+ }
+ ColumnPartitionMetadata columnPartitionMetadata =
+ partitionMetadata.getColumnPartitionMap().get(_partitionColumn);
+ if (columnPartitionMetadata == null || columnPartitionMetadata.getNumPartitions() != _numPartitions) {
+ return -1;
+ }
+ Set<Integer> partitions = columnPartitionMetadata.getPartitions();
+ if (partitions == null || partitions.size() != 1) {
+ return -1;
+ }
+ return partitions.iterator().next();
+ }
+
+ private static Long getTimeMsFromRealtimeSegmentName(String segmentName) {
+ LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
+ if (llcSegmentName != null) {
+ return llcSegmentName.getCreationTimeMs();
+ }
+ UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName);
+ if (uploadedRealtimeSegmentName != null) {
+ return parseRealtimeSegmentNameTime(uploadedRealtimeSegmentName.getCreationTime());
+ }
+ return null;
+ }
+
+ private static long parseRealtimeSegmentNameTime(String timeString) {
+ try {
+ return LocalDateTime.parse(timeString, REALTIME_SEGMENT_NAME_TIME_FORMATTER).toInstant(ZoneOffset.UTC)
+ .toEpochMilli();
+ } catch (DateTimeParseException e) {
+ return -1L;
+ }
+ }
+
+ private static long getEndTimeMs(SegmentZKMetadata zkMetadata) {
+ long endTimeMs = zkMetadata.getEndTimeMs();
+ if (endTimeMs > 0) {
+ return endTimeMs;
+ }
+ // Handle explicit end time 0 (treated as invalid by SegmentZKMetadata#getEndTimeMs()).
+ Map<String, String> simpleFields = zkMetadata.getSimpleFields();
+ String endTimeString = simpleFields.get(Segment.END_TIME);
+ if ("0".equals(endTimeString)) {
+ String timeUnit = simpleFields.get(Segment.TIME_UNIT);
+ if (timeUnit != null) {
+ try {
+ TimeUnit.valueOf(timeUnit);
+ } catch (IllegalArgumentException e) {
+ return -1L;
+ }
+ return 0L;
+ }
+ }
+ return -1L;
+ }
Review Comment:
The TimeUnit.valueOf call at line 328 is only used to validate the time unit format, but its result is discarded. The validation ensures the time unit string is a valid TimeUnit enum value before returning 0L for endTime=0, but the specific time unit doesn't affect the return value. Consider adding a comment explaining that this validation is intentional and the actual time unit value doesn't matter when endTime is 0.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2767192770
##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/sampler/TableSamplerConfig.java:
##########
@@ -0,0 +1,61 @@
+/**
+ * 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.spi.config.table.sampler;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.config.BaseJsonConfig;
+
+
+/**
+ * Configuration for a table sampler.
+ *
+ * Samplers are defined in {@link org.apache.pinot.spi.config.table.TableConfig} and can be selected at query time
+ * via a query option. The sampler type can be one of the built-in types, a fully qualified class name, or an alias
+ * discovered via {@code pinot.broker.table.sampler.annotation.packages}.
+ */
+public class TableSamplerConfig extends BaseJsonConfig {
+ private final String _name;
+ private final String _type;
+ private final Map<String, String> _properties;
+
+ @JsonCreator
+ public TableSamplerConfig(@JsonProperty(value = "name", required = true) String name,
+ @JsonProperty(value = "type", required = true) String type,
+ @JsonProperty("properties") @Nullable Map<String, String> properties) {
Review Comment:
The Javadoc comment for this class references a package path in the NOTE section that may become outdated if the TableSampler interface is moved. Consider using a less specific reference or adding a check to ensure the documentation stays synchronized with the actual implementation location.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSamplerFactory.java:
##########
@@ -0,0 +1,137 @@
+/**
+ * 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 com.google.common.annotations.VisibleForTesting;
+import java.lang.reflect.Modifier;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.annotations.tablesampler.TableSamplerProvider;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.plugin.PluginManager;
+import org.apache.pinot.spi.utils.PinotReflectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class TableSamplerFactory {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TableSamplerFactory.class);
+ private static final String ANNOTATION_PACKAGES_KEY = "annotation.packages";
Review Comment:
The DEFAULT_ANNOTATION_PACKAGES is hardcoded to a single package. If built-in samplers are added outside this package, they won't be discovered by default. Consider documenting this limitation or making the default list configurable via external configuration.
```suggestion
private static final String ANNOTATION_PACKAGES_KEY = "annotation.packages";
// NOTE: By default we only scan this package for @TableSamplerProvider annotations.
// If built-in samplers are added in other packages, callers must configure additional
// packages via the `annotation.packages` key; they will not be discovered otherwise.
```
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TableSamplerIntegrationTest.java:
##########
@@ -0,0 +1,130 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class TableSamplerIntegrationTest extends CustomDataQueryClusterIntegrationTest {
+ private static final int DAYS = 7;
+ private static final int SEGMENTS_PER_DAY = 10;
+ private static final int RECORDS_PER_SEGMENT = 1;
+ private static final int BASE_DAY = 20000;
+
+ private static final String DAYS_SINCE_EPOCH_COL = "DaysSinceEpoch";
+
+ @Override
+ public String getTableName() {
+ return "TableSamplerIntegrationTest";
+ }
+
+ @Override
+ protected long getCountStarResult() {
+ return (long) DAYS * SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT;
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder().setSchemaName(getTableName())
+ .addDateTime(DAYS_SINCE_EPOCH_COL, FieldSpec.DataType.INT, "1:DAYS:EPOCH", "1:DAYS")
+ .build();
+ }
+
+ @Override
+ public TableConfig createOfflineTableConfig() {
+ TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName())
+ .setTimeColumnName(DAYS_SINCE_EPOCH_COL)
+ .setTimeType("DAYS")
+ .build();
+ tableConfig.setTableSamplers(List.of(
+ new TableSamplerConfig("perDay", "timeBucket", Map.of("numSegmentsPerDay", "1"))));
+ return tableConfig;
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws Exception {
+ org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
+ avroSchema.setFields(List.of(
+ new org.apache.avro.Schema.Field(DAYS_SINCE_EPOCH_COL,
+ org.apache.avro.Schema.create(org.apache.avro.Schema.Type.INT), null, null)));
+
+ List<File> files = new ArrayList<>();
+ for (int day = 0; day < DAYS; day++) {
+ int dayValue = BASE_DAY + day;
+ for (int seg = 0; seg < SEGMENTS_PER_DAY; seg++) {
+ File avroFile = new File(_tempDir, "data_day_" + day + "_seg_" + seg + ".avro");
+ try (DataFileWriter<GenericData.Record> fileWriter =
+ new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) {
+ fileWriter.create(avroSchema, avroFile);
+ for (int docId = 0; docId < RECORDS_PER_SEGMENT; docId++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ record.put(DAYS_SINCE_EPOCH_COL, dayValue);
+ fileWriter.append(record);
+ }
+ }
+ files.add(avroFile);
+ }
+ }
+ return files;
+ }
+
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testTimeBucketSamplerForGroupByDay(boolean useMultiStageQueryEngine)
Review Comment:
The integration test only validates the `timeBucket` sampler with `numSegmentsPerDay=1`. Consider adding test cases for the `firstN` sampler and edge cases such as empty segment sets or invalid sampler configurations to ensure comprehensive coverage.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1277,37 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
+ LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
+ samplerName, tableNameWithType);
Review Comment:
This warning will be logged on every query when a non-existent sampler is requested, potentially generating excessive log noise for misconfigured clients. Consider rate-limiting this warning or logging it at debug level after the first occurrence.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2784428415
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -448,14 +464,22 @@ private void processInstanceConfigChangeInternal() {
// Update routing entry for all tables
for (RoutingEntry routingEntry : _routingEntryMap.values()) {
+ String tableNameWithType = routingEntry.getTableNameWithType();
try {
- Object tableLock = getRoutingTableBuildLock(routingEntry.getTableNameWithType());
+ Object tableLock = getRoutingTableBuildLock(tableNameWithType);
synchronized (tableLock) {
routingEntry.onInstancesChange(_routableServers, changedServers);
+
+ // Handle sampler routing entries
+ if (_samplerRoutingEntryMap.containsKey(tableNameWithType)) {
+ for (RoutingEntry samplerRoutingEntry : _samplerRoutingEntryMap.get(tableNameWithType).values()) {
Review Comment:
`containsKey()` followed by `get()` on a `ConcurrentHashMap` is not atomic; if the entry is removed between calls, `_samplerRoutingEntryMap.get(tableNameWithType)` can be null and `values()` will throw NPE. Safer pattern: fetch once into a local variable (e.g., `Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);`) and null-check before iterating. (This same pattern repeats in exclude/include paths as well.)
```suggestion
Map<String, RoutingEntry> samplerRoutingEntries = _samplerRoutingEntryMap.get(tableNameWithType);
if (samplerRoutingEntries != null) {
for (RoutingEntry samplerRoutingEntry : samplerRoutingEntries.values()) {
```
##########
pinot-spi/src/test/java/org/apache/pinot/spi/config/table/TableConfigTest.java:
##########
@@ -104,4 +106,15 @@ public void testCopyConstructor() {
assertEquals(config, copy);
assertEquals(config.toJsonString(), copy.toJsonString());
}
+
+ @Test
+ public void testDuplicateTableSamplerNamesRejected() {
+ TableConfig config = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+ List<TableSamplerConfig> duplicateSamplers = List.of(
+ new TableSamplerConfig("sampler1", "firstN", Map.of("numSegments", "10")),
+ new TableSamplerConfig("sampler1", "firstN", Map.of("numSegments", "1")));
+ IllegalArgumentException e = Assert.expectThrows(IllegalArgumentException.class,
+ () -> config.setTableSamplers(duplicateSamplers));
+ assertTrue(e.getMessage().contains("Duplicate table sampler name: sampler1"));
+ }
Review Comment:
The duplicate-name test only covers exact string matches. If sampler names are intended to be user-facing and stable (especially given the query option is free-form text), add coverage for normalization behaviors you support (e.g., leading/trailing whitespace and case-insensitive duplicates like `sampler1` vs ` Sampler1 `). This helps prevent regressions when making sampler lookup consistent between table config and query options.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,11 +1188,43 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
- RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
+ RoutingEntry routingEntry = getRoutingEntry(tableNameWithType, samplerName);
return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
}
+ @Nullable
+ private RoutingEntry getRoutingEntry(String tableNameWithType, String samplerName) {
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
+ LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
Review Comment:
This `WARN` is in the per-query routing path and can become very noisy (and expensive) if clients send an unknown `tableSampler` option or have mismatched casing/whitespace. Consider downgrading to `DEBUG`, adding rate-limiting, or logging once per (table, samplerName) with a cache to avoid log amplification on hot paths.
```suggestion
LOGGER.debug("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
```
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2791042010
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -123,6 +130,8 @@ public abstract class BaseBrokerRoutingManager implements RoutingManager, Cluste
protected final BrokerMetrics _brokerMetrics;
protected final Map<String, RoutingEntry> _routingEntryMap = new ConcurrentHashMap<>();
+ // Additional sampler-specific routing components per table, keyed by normalized sampler name.
+ protected final Map<String, Map<String, SamplerInfo>> _samplerRoutingEntryMap = new ConcurrentHashMap<>();
Review Comment:
You can put `Map<String, SamplerInfo>` as one member variable within the `RoutingEntry`, then all the methods within `RoutingEntry` can directly handle the samplers
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2796246624
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1248,6 +1343,45 @@ boolean isDisabled() {
return _disabled;
}
+ private void updateSamplerInfos(IdealState idealState, ExternalView externalView,
+ Set<String> preSelectedOnlineSegments) {
+ if (_samplerInfos.isEmpty()) {
+ return;
+ }
+ for (Map.Entry<String, SamplerInfo> entry : _samplerInfos.entrySet()) {
+ String samplerName = entry.getKey();
+ SamplerInfo samplerInfo = entry.getValue();
+ try {
+ Set<String> samplerPreSelectedOnlineSegments =
+ samplerInfo._tableSampler.sampleSegments(preSelectedOnlineSegments);
+ if (samplerPreSelectedOnlineSegments == null) {
Review Comment:
(minor) Same here, check is redundant
##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java:
##########
@@ -68,6 +68,8 @@ private QueryOptionsUtils() {
}
}
}
+ // Backward-compatible alias used by SQL SET syntax: SET sampler='x'
+ configResolver.put("sampler", QueryOptionKey.TABLE_SAMPLER);
Review Comment:
Is `sampler` a standard concept in SQL? If so, should we just use it as the option name?
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -810,10 +820,58 @@ private void buildRoutingInternal(String tableNameWithType) {
}
segmentZkMetadataFetcher.init(idealState, externalView, preSelectedOnlineSegments);
+ // Build table sampler contexts keyed by normalized sampler name.
+ Map<String, SamplerInfo> samplerInfos = Map.of();
+ List<TableSamplerConfig> tableSamplerConfigs = tableConfig.getTableSamplers();
+ if (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, SamplerInfo> configuredSamplerInfos =
+ new HashMap<>(HashUtil.getHashMapCapacity(tableSamplerConfigs.size()));
Review Comment:
(minor)
```suggestion
Maps.newHashMapWithExpectedSize(tableSamplerConfigs.size());
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -810,10 +820,58 @@ private void buildRoutingInternal(String tableNameWithType) {
}
segmentZkMetadataFetcher.init(idealState, externalView, preSelectedOnlineSegments);
+ // Build table sampler contexts keyed by normalized sampler name.
+ Map<String, SamplerInfo> samplerInfos = Map.of();
+ List<TableSamplerConfig> tableSamplerConfigs = tableConfig.getTableSamplers();
+ if (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, SamplerInfo> configuredSamplerInfos =
+ new HashMap<>(HashUtil.getHashMapCapacity(tableSamplerConfigs.size()));
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ String normalizedSamplerName = normalizeSamplerName(samplerName);
+ if (configuredSamplerInfos.containsKey(normalizedSamplerName)) {
+ LOGGER.warn("Skipping duplicate normalized table sampler name: '{}' for table: {}", samplerName,
+ tableNameWithType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+ Set<String> samplerPreSelectedOnlineSegments = sampler.sampleSegments(preSelectedOnlineSegments);
+ if (samplerPreSelectedOnlineSegments == null) {
+ LOGGER.warn("Sampler '{}' returned null segments for table '{}', treating it as empty", samplerName,
+ tableNameWithType);
+ samplerPreSelectedOnlineSegments = Collections.emptySet();
Review Comment:
(nit) `Set.of()`
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,9 +1146,30 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
- return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
+ if (routingEntry == null) {
+ return null;
+ }
+ return routingEntry.getSegments(brokerRequest, samplerName);
+ }
+
+ private static String normalizeSamplerName(String samplerName) {
+ return samplerName.trim().toLowerCase(Locale.ROOT);
+ }
+
+ @Nullable
+ static String extractSamplerName(BrokerRequest brokerRequest) {
+ if (!brokerRequest.isSetPinotQuery() || !brokerRequest.getPinotQuery().isSetQueryOptions()) {
+ return null;
+ }
Review Comment:
(minor) This check shouldn't be required
##########
pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java:
##########
@@ -111,6 +113,11 @@ public static Long getTimeoutMs(Map<String, String> queryOptions) {
return checkedParseLongPositive(QueryOptionKey.TIMEOUT_MS, timeoutMsString);
}
+ @Nullable
+ public static String getTableSampler(@Nullable Map<String, String> queryOptions) {
Review Comment:
(minor) To be consistent with other methods, remove `@Nullable` for `queryOptions`. It is never null in `BrokerRequest`
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -810,10 +820,58 @@ private void buildRoutingInternal(String tableNameWithType) {
}
segmentZkMetadataFetcher.init(idealState, externalView, preSelectedOnlineSegments);
+ // Build table sampler contexts keyed by normalized sampler name.
+ Map<String, SamplerInfo> samplerInfos = Map.of();
+ List<TableSamplerConfig> tableSamplerConfigs = tableConfig.getTableSamplers();
+ if (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, SamplerInfo> configuredSamplerInfos =
+ new HashMap<>(HashUtil.getHashMapCapacity(tableSamplerConfigs.size()));
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ String normalizedSamplerName = normalizeSamplerName(samplerName);
+ if (configuredSamplerInfos.containsKey(normalizedSamplerName)) {
+ LOGGER.warn("Skipping duplicate normalized table sampler name: '{}' for table: {}", samplerName,
+ tableNameWithType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+ Set<String> samplerPreSelectedOnlineSegments = sampler.sampleSegments(preSelectedOnlineSegments);
+ if (samplerPreSelectedOnlineSegments == null) {
Review Comment:
(minor) What is the contract for `sampleSegments()`? Do we allow it returning `null`?
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "xiangfu0 (via GitHub)" <gi...@apache.org>.
xiangfu0 commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2796336219
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,9 +1146,30 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
- return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
+ if (routingEntry == null) {
+ return null;
+ }
+ return routingEntry.getSegments(brokerRequest, samplerName);
+ }
+
+ private static String normalizeSamplerName(String samplerName) {
+ return samplerName.trim().toLowerCase(Locale.ROOT);
+ }
+
+ @Nullable
+ static String extractSamplerName(BrokerRequest brokerRequest) {
+ if (!brokerRequest.isSetPinotQuery() || !brokerRequest.getPinotQuery().isSetQueryOptions()) {
+ return null;
+ }
Review Comment:
this `!brokerRequest.isSetPinotQuery() ` is still needed, some tests will fail this check.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2768450026
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TableSamplerIntegrationTest.java:
##########
@@ -0,0 +1,128 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class TableSamplerIntegrationTest extends CustomDataQueryClusterIntegrationTest {
+ private static final int DAYS = 7;
+ private static final int SEGMENTS_PER_DAY = 10;
+ private static final int RECORDS_PER_SEGMENT = 1;
+ private static final int BASE_DAY = 20000;
+
+ private static final String DAYS_SINCE_EPOCH_COL = "DaysSinceEpoch";
+
+ @Override
+ public String getTableName() {
+ return "TableSamplerIntegrationTest";
+ }
+
+ @Override
+ protected long getCountStarResult() {
+ return (long) DAYS * SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT;
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder().setSchemaName(getTableName())
+ .addDateTime(DAYS_SINCE_EPOCH_COL, FieldSpec.DataType.INT, "1:DAYS:EPOCH", "1:DAYS")
+ .build();
+ }
+
+ @Override
+ public TableConfig createOfflineTableConfig() {
+ TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName())
+ .setTimeColumnName(DAYS_SINCE_EPOCH_COL)
+ .setTimeType("DAYS")
+ .build();
+ tableConfig.setTableSamplers(List.of(
+ new TableSamplerConfig("firstOnly", "firstN", Map.of("numSegments", "1"))));
+ return tableConfig;
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws Exception {
+ org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("myRecord", null, null, false);
+ avroSchema.setFields(List.of(
+ new org.apache.avro.Schema.Field(DAYS_SINCE_EPOCH_COL,
+ org.apache.avro.Schema.create(org.apache.avro.Schema.Type.INT), null, null)));
+
+ List<File> files = new ArrayList<>();
+ for (int day = 0; day < DAYS; day++) {
+ int dayValue = BASE_DAY + day;
+ for (int seg = 0; seg < SEGMENTS_PER_DAY; seg++) {
+ File avroFile = new File(_tempDir, "data_day_" + day + "_seg_" + seg + ".avro");
+ try (DataFileWriter<GenericData.Record> fileWriter =
+ new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) {
+ fileWriter.create(avroSchema, avroFile);
+ for (int docId = 0; docId < RECORDS_PER_SEGMENT; docId++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ record.put(DAYS_SINCE_EPOCH_COL, dayValue);
+ fileWriter.append(record);
+ }
+ }
+ files.add(avroFile);
+ }
+ }
+ return files;
+ }
+
+ @Test(dataProvider = "useBothQueryEngines")
+ public void testFirstNSamplerForGroupByDay(boolean useMultiStageQueryEngine)
+ throws Exception {
+ setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+
+ JsonNode full = postQuery("SELECT DaysSinceEpoch, COUNT(*) AS cnt FROM " + getTableName()
+ + " GROUP BY DaysSinceEpoch ORDER BY DaysSinceEpoch");
+ JsonNode fullRows = full.path("resultTable").path("rows");
+ Assert.assertEquals(fullRows.size(), DAYS);
+ for (int i = 0; i < DAYS; i++) {
+ Assert.assertEquals(fullRows.get(i).get(0).asInt(), BASE_DAY + i);
+ Assert.assertEquals(fullRows.get(i).get(1).asLong(), (long) SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT);
+ }
+ Assert.assertEquals(Integer.parseInt(full.path("numSegmentsQueried").asText()), DAYS * SEGMENTS_PER_DAY);
+
+ JsonNode sampled = postQueryWithOptions("SELECT DaysSinceEpoch, COUNT(*) AS cnt FROM " + getTableName()
+ + " GROUP BY DaysSinceEpoch ORDER BY DaysSinceEpoch",
+ "tableSampler=firstOnly");
Review Comment:
The test only validates the `firstN` sampler with `numSegments=1`. Consider adding test coverage for other scenarios such as numSegments > 1, and potentially for the `timeBucket` sampler mentioned in the PR description.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1277,37 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
+ LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
+ samplerName, tableNameWithType);
+ }
+ return _routingEntryMap.get(tableNameWithType);
+ }
+
Review Comment:
This warning is logged on every query when a sampler is not found, which could generate excessive logs if a misconfigured sampler name is used repeatedly. Consider rate-limiting this warning or logging it at debug level after the first occurrence.
```suggestion
if (MissingSamplerWarningTracker.shouldLogWarn(tableNameWithType, samplerName)) {
LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
samplerName, tableNameWithType);
} else if (LOGGER.isDebugEnabled()) {
LOGGER.debug(
"Requested sampler '{}' not found for table '{}'; falling back to default routing entry "
+ "(subsequent occurrences logged at DEBUG level only)",
samplerName, tableNameWithType);
}
}
return _routingEntryMap.get(tableNameWithType);
}
/**
* Tracks sampler/table combinations that have already emitted a warning, to avoid excessive logs
* when a misconfigured sampler name is used repeatedly.
*/
private static final class MissingSamplerWarningTracker {
private static final Map<String, Set<String>> TABLE_TO_SAMPLERS = new ConcurrentHashMap<>();
private MissingSamplerWarningTracker() {
}
static boolean shouldLogWarn(String tableNameWithType, String samplerName) {
Set<String> samplers =
TABLE_TO_SAMPLERS.computeIfAbsent(tableNameWithType, k -> ConcurrentHashMap.newKeySet());
// Returns true only the first time we see this sampler for the given table.
return samplers.add(samplerName);
}
}
```
##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -254,6 +265,38 @@ public void setCustomConfig(TableCustomConfig customConfig) {
_customConfig = customConfig;
}
+ @JsonProperty(TABLE_SAMPLERS_KEY)
+ @Nullable
+ public List<TableSamplerConfig> getTableSamplers() {
+ return _tableSamplers;
+ }
+
+ public void setTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ _tableSamplers = sanitizeTableSamplers(tableSamplers);
+ }
+
+ @Nullable
+ private static List<TableSamplerConfig> sanitizeTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
Review Comment:
The method name `sanitizeTableSamplers` doesn't clearly indicate that it also validates for duplicates. Consider renaming to `validateAndSanitizeTableSamplers` or `sanitizeAndValidateTableSamplers` to better reflect its dual purpose.
```suggestion
_tableSamplers = validateAndSanitizeTableSamplers(tableSamplers);
}
@Nullable
private static List<TableSamplerConfig> validateAndSanitizeTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -820,6 +940,69 @@ 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) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableNameWithType, tableConfig, samplerConfig, _propertyStore);
+
+ SegmentPreSelector samplerSegmentPreSelector =
+ new TableSamplerSegmentPreSelector(segmentPreSelector, sampler);
+ Set<String> samplerPreSelectedOnlineSegments =
+ // SegmentPreSelector contract allows mutation of the input set. Avoid passing the shared
+ // onlineSegments.
Review Comment:
The comment describes the contract about mutation but doesn't explain why we're creating a copy. Consider expanding the comment to clarify that we're creating a defensive copy to protect the shared `onlineSegments` set from being mutated by the sampler, as this set is reused for other samplers.
```suggestion
// SegmentPreSelector contract allows mutation of the input set. Create a defensive copy here so
// the sampler cannot mutate the shared onlineSegments set, which is reused by other samplers.
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSamplerFactory.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 com.google.common.annotations.VisibleForTesting;
+import java.lang.reflect.Modifier;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.annotations.tablesampler.TableSamplerProvider;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.plugin.PluginManager;
+import org.apache.pinot.spi.utils.PinotReflectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class TableSamplerFactory {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TableSamplerFactory.class);
+ private static final String ANNOTATION_PACKAGES_KEY = "annotation.packages";
+ private static final List<String> DEFAULT_ANNOTATION_PACKAGES =
+ List.of("org.apache.pinot.broker.routing.tablesampler");
+ private static final Map<String, String> REGISTERED_TABLE_SAMPLER_CLASS_MAP = new ConcurrentHashMap<>();
+
+ private TableSamplerFactory() {
+ }
+
+ public static void init(PinotConfiguration tableSamplerConfig) {
+ if (tableSamplerConfig == null) {
+ return;
+ }
+ registerAnnotatedTableSamplers(tableSamplerConfig);
+ }
+
+ public static void register(String alias, String className) {
+ if (StringUtils.isBlank(alias)) {
+ LOGGER.warn("Skipping table sampler registration because alias is blank");
+ return;
+ }
+ if (StringUtils.isBlank(className)) {
+ LOGGER.warn("Skipping table sampler registration for alias '{}' because class name is blank", alias);
+ return;
+ }
+ String normalizedAlias = normalizeType(alias);
+ String trimmedClassName = className.trim();
+ String previousClassName = REGISTERED_TABLE_SAMPLER_CLASS_MAP.put(normalizedAlias, trimmedClassName);
+ if (previousClassName == null) {
+ LOGGER.info("Registered table sampler alias '{}' -> '{}'", alias, trimmedClassName);
+ } else if (!previousClassName.equals(trimmedClassName)) {
+ LOGGER.warn("Overriding table sampler alias '{}' from '{}' to '{}'", alias, previousClassName, trimmedClassName);
+ }
+ }
+
+ public static TableSampler create(String type) {
+ String resolvedClassName = resolveClassName(type);
+ String classNameToLoad = resolvedClassName != null ? resolvedClassName : type;
+ try {
+ return PluginManager.get().createInstance(classNameToLoad);
+ } catch (Exception e) {
+ String errorMessage = resolvedClassName != null
+ ? String.format("Failed to create TableSampler for alias '%s' mapped to class '%s'", type,
+ resolvedClassName)
+ : "Failed to create TableSampler for type: " + type;
+ throw new RuntimeException(errorMessage, e);
+ }
+ }
+
+ @VisibleForTesting
+ static void clearRegistry() {
+ REGISTERED_TABLE_SAMPLER_CLASS_MAP.clear();
+ }
+
+ private static void registerAnnotatedTableSamplers(PinotConfiguration tableSamplerConfig) {
+ List<String> configuredPackages = tableSamplerConfig.getProperty(ANNOTATION_PACKAGES_KEY, List.of());
+ LinkedHashSet<String> combinedPackages = new LinkedHashSet<>(DEFAULT_ANNOTATION_PACKAGES);
+ if (configuredPackages != null) {
+ for (String packageName : configuredPackages) {
+ if (StringUtils.isNotBlank(packageName)) {
+ combinedPackages.add(packageName.trim());
+ }
+ }
+ }
+ List<String> sanitizedPackages = combinedPackages.stream().collect(Collectors.toList());
Review Comment:
The `combinedPackages` is already a LinkedHashSet which maintains insertion order. Converting to a List via stream is unnecessary. Consider using `new ArrayList<>(combinedPackages)` for a more direct conversion.
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2768928080
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/FirstNSegmentsTableSampler.java:
##########
@@ -0,0 +1,76 @@
+/**
+ * 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.util.ArrayList;
+import java.util.Collections;
+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.spi.annotations.tablesampler.TableSamplerProvider;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+
+
+/**
+ * Selects the first N segments after sorting segment names lexicographically.
+ *
+ * <p>Config:
+ * <ul>
+ * <li>{@code properties.numSegments}: positive integer</li>
+ * </ul>
+ */
+@TableSamplerProvider(name = FirstNSegmentsTableSampler.TYPE)
+public class FirstNSegmentsTableSampler implements TableSampler {
+ public static final String TYPE = "firstN";
+ public static final String PROP_NUM_SEGMENTS = "numSegments";
+
+ private int _numSegments;
+
+ @Override
+ public void init(String tableNameWithType, TableConfig tableConfig, TableSamplerConfig samplerConfig,
+ ZkHelixPropertyStore<ZNRecord> propertyStore) {
+ Map<String, String> props = samplerConfig.getProperties();
+ if (MapUtils.isEmpty(props) || !props.containsKey(PROP_NUM_SEGMENTS)) {
+ throw new IllegalArgumentException(
+ "Missing required property '" + PROP_NUM_SEGMENTS + "' for table sampler type '" + TYPE + "'");
+ }
+ _numSegments = Integer.parseInt(props.get(PROP_NUM_SEGMENTS));
Review Comment:
The NumberFormatException from parseInt is not caught, resulting in an unclear error message. Wrap this in a try-catch and provide a clear message indicating the invalid value and expected format.
```suggestion
String numSegmentsString = props.get(PROP_NUM_SEGMENTS);
try {
_numSegments = Integer.parseInt(numSegmentsString);
} catch (NumberFormatException e) {
throw new IllegalArgumentException(
"Invalid value for '" + PROP_NUM_SEGMENTS + "': '" + numSegmentsString + "'. Expected a positive integer.",
e);
}
```
##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TableSamplerIntegrationTest.java:
##########
@@ -0,0 +1,140 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.SchemaBuilder;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class TableSamplerIntegrationTest extends CustomDataQueryClusterIntegrationTest {
+ private static final int DAYS = 7;
+ private static final int SEGMENTS_PER_DAY = 10;
+ private static final int RECORDS_PER_SEGMENT = 1;
+ private static final int BASE_DAY = 20000;
+
+ private static final String DAYS_SINCE_EPOCH_COL = "DaysSinceEpoch";
+
+ @Override
+ public String getTableName() {
+ return "TableSamplerIntegrationTest";
+ }
+
+ @Override
+ protected long getCountStarResult() {
+ return (long) DAYS * SEGMENTS_PER_DAY * RECORDS_PER_SEGMENT;
+ }
+
+ @Override
+ public Schema createSchema() {
+ return new Schema.SchemaBuilder().setSchemaName(getTableName())
+ .addDateTime(DAYS_SINCE_EPOCH_COL, FieldSpec.DataType.INT, "1:DAYS:EPOCH", "1:DAYS")
+ .build();
+ }
+
+ @Override
+ public TableConfig createOfflineTableConfig() {
+ TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName())
+ .setTimeColumnName(DAYS_SINCE_EPOCH_COL)
+ .setTimeType("DAYS")
+ .build();
+ tableConfig.setTableSamplers(List.of(
+ new TableSamplerConfig("firstOnly", "firstN", Map.of("numSegments", "1")),
+ new TableSamplerConfig("firstTwo", "firstN", Map.of("numSegments", "2"))));
+ return tableConfig;
+ }
+
+ @Override
+ public List<File> createAvroFiles()
+ throws Exception {
+ var fieldAssembler = SchemaBuilder.record("myRecord").fields();
+ fieldAssembler.name(DAYS_SINCE_EPOCH_COL).type().intType().noDefault();
+ var avroSchema = fieldAssembler.endRecord();
+
+ List<File> files = new ArrayList<>();
+ for (int day = 0; day < DAYS; day++) {
+ int dayValue = BASE_DAY + day;
+ for (int seg = 0; seg < SEGMENTS_PER_DAY; seg++) {
Review Comment:
The file naming pattern "data_day_X_seg_Y.avro" creates segments that sort lexicographically as intended, but this dependency on naming for test correctness is fragile. Consider adding a comment explaining this requirement or using a more explicit segment ordering mechanism.
```suggestion
for (int seg = 0; seg < SEGMENTS_PER_DAY; seg++) {
// NOTE: The table sampler tests rely on the lexicographic ordering of segment names.
// Keep the "data_day_<day>_seg_<seg>.avro" pattern so segments sort by day, then by segment index.
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1277,59 @@ 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 = extractSamplerName(brokerRequest);
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ RoutingEntry samplerEntry = samplerEntries != null ? samplerEntries.get(samplerName) : null;
Review Comment:
The sampler name lookup is case-sensitive, but sampler names in table config may vary in case. Consider normalizing the sampler name (e.g., to lowercase) before lookup to avoid mismatches.
##########
pinot-spi/src/main/java/org/apache/pinot/spi/config/table/TableConfig.java:
##########
@@ -254,6 +265,39 @@ public void setCustomConfig(TableCustomConfig customConfig) {
_customConfig = customConfig;
}
+ @JsonProperty(TABLE_SAMPLERS_KEY)
+ @Nullable
+ public List<TableSamplerConfig> getTableSamplers() {
+ return _tableSamplers;
+ }
+
+ public void setTableSamplers(@Nullable List<TableSamplerConfig> tableSamplers) {
+ _tableSamplers = sanitizeAndValidateTableSamplers(tableSamplers);
+ }
+
+ @Nullable
+ private static List<TableSamplerConfig> sanitizeAndValidateTableSamplers(
+ @Nullable List<TableSamplerConfig> tableSamplers) {
+ if (tableSamplers == null || tableSamplers.isEmpty()) {
+ return null;
+ }
+ List<TableSamplerConfig> sanitized = new ArrayList<>(tableSamplers.size());
+ Set<String> seenNames = new HashSet<>();
+ for (TableSamplerConfig config : tableSamplers) {
+ if (config != null) {
+ String name = config.getName();
+ if (name != null && !seenNames.add(name)) {
+ throw new IllegalArgumentException("Duplicate table sampler name: " + name);
Review Comment:
Validate that the sampler name does not contain only whitespace. A name consisting of only spaces could pass the null check but would be effectively empty, potentially allowing duplicate "blank" samplers.
```suggestion
if (name != null) {
if (name.trim().isEmpty()) {
throw new IllegalArgumentException("Table sampler name cannot be blank");
}
if (!seenNames.add(name)) {
throw new IllegalArgumentException("Duplicate table sampler name: " + name);
}
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSamplerFactory.java:
##########
@@ -0,0 +1,147 @@
+/**
+ * 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 com.google.common.annotations.VisibleForTesting;
+import java.lang.reflect.Modifier;
+import java.util.ArrayList;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.annotations.tablesampler.TableSamplerProvider;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.plugin.PluginManager;
+import org.apache.pinot.spi.utils.PinotReflectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class TableSamplerFactory {
+ private static final Logger LOGGER = LoggerFactory.getLogger(TableSamplerFactory.class);
+ private static final String ANNOTATION_PACKAGES_KEY = "annotation.packages";
+ // Keep this list in sync with built-in TableSampler locations; additional packages can be configured via broker
+ // config, but defaults should always include built-ins.
+ private static final List<String> DEFAULT_ANNOTATION_PACKAGES =
+ List.of("org.apache.pinot.broker.routing.tablesampler");
+ private static final Map<String, String> REGISTERED_TABLE_SAMPLER_CLASS_MAP = new ConcurrentHashMap<>();
+
+ private TableSamplerFactory() {
+ }
+
+ public static void init(PinotConfiguration tableSamplerConfig) {
+ if (tableSamplerConfig == null) {
+ return;
+ }
+ registerAnnotatedTableSamplers(tableSamplerConfig);
+ }
+
+ public static void register(String alias, String className) {
+ if (StringUtils.isBlank(alias)) {
+ LOGGER.warn("Skipping table sampler registration because alias is blank");
+ return;
+ }
+ if (StringUtils.isBlank(className)) {
+ LOGGER.warn("Skipping table sampler registration for alias '{}' because class name is blank", alias);
+ return;
+ }
+ String normalizedAlias = normalizeType(alias);
+ String trimmedClassName = className.trim();
+ String previousClassName = REGISTERED_TABLE_SAMPLER_CLASS_MAP.put(normalizedAlias, trimmedClassName);
+ if (previousClassName == null) {
+ LOGGER.info("Registered table sampler alias '{}' -> '{}'", alias, trimmedClassName);
+ } else if (!previousClassName.equals(trimmedClassName)) {
+ LOGGER.warn("Overriding table sampler alias '{}' from '{}' to '{}'", alias, previousClassName, trimmedClassName);
+ }
+ }
+
+ public static TableSampler create(String type) {
+ String resolvedClassName = resolveClassName(type);
+ String classNameToLoad = resolvedClassName != null ? resolvedClassName : type;
+ try {
+ return PluginManager.get().createInstance(classNameToLoad);
+ } catch (Exception e) {
+ String errorMessage = resolvedClassName != null
+ ? String.format("Failed to create TableSampler for alias '%s' mapped to class '%s'", type,
+ resolvedClassName)
+ : "Failed to create TableSampler for type: " + type;
+ throw new RuntimeException(errorMessage, e);
+ }
+ }
+
+ @VisibleForTesting
+ static void clearRegistry() {
+ REGISTERED_TABLE_SAMPLER_CLASS_MAP.clear();
+ }
+
+ private static void registerAnnotatedTableSamplers(PinotConfiguration tableSamplerConfig) {
+ List<String> configuredPackages = tableSamplerConfig.getProperty(ANNOTATION_PACKAGES_KEY, List.of());
+ LinkedHashSet<String> combinedPackages = new LinkedHashSet<>(DEFAULT_ANNOTATION_PACKAGES);
+ if (configuredPackages != null) {
+ for (String packageName : configuredPackages) {
+ if (StringUtils.isNotBlank(packageName)) {
+ combinedPackages.add(packageName.trim());
+ }
+ }
+ }
+ List<String> sanitizedPackages = new ArrayList<>(combinedPackages);
+ if (sanitizedPackages.isEmpty()) {
+ LOGGER.info("No table sampler annotation packages configured");
+ return;
+ }
+ Set<Class<?>> samplerClasses =
+ PinotReflectionUtils.getClassesThroughReflection(sanitizedPackages, ".*", TableSamplerProvider.class);
+ for (Class<?> samplerClass : samplerClasses) {
+ TableSamplerProvider annotation = samplerClass.getAnnotation(TableSamplerProvider.class);
+ if (annotation == null || !annotation.enabled()) {
+ continue;
+ }
+ if (!TableSampler.class.isAssignableFrom(samplerClass)) {
+ LOGGER.warn("Skipping table sampler class '{}' because it does not implement TableSampler",
+ samplerClass.getName());
+ continue;
+ }
+ if (!Modifier.isPublic(samplerClass.getModifiers()) || Modifier.isAbstract(samplerClass.getModifiers())) {
+ LOGGER.warn("Skipping table sampler class '{}' because it is not a public concrete class",
+ samplerClass.getName());
+ continue;
+ }
+ String alias = annotation.name();
+ if (StringUtils.isBlank(alias)) {
+ LOGGER.warn("Skipping table sampler class '{}' because annotation name is blank", samplerClass.getName());
+ continue;
+ }
+ register(alias, samplerClass.getName());
+ }
+ }
+
+ private static String resolveClassName(String type) {
+ if (StringUtils.isBlank(type)) {
+ return null;
+ }
+ return REGISTERED_TABLE_SAMPLER_CLASS_MAP.get(normalizeType(type));
+ }
+
Review Comment:
The normalizeType method converts sampler names to lowercase, but this behavior is not documented in TableSamplerConfig or TableSamplerProvider. Add documentation explaining that sampler name matching is case-insensitive.
```suggestion
/**
* Normalizes a table sampler type/alias for use in the internal registry.
* <p>
* Both registration and lookup go through this method, so sampler names are effectively
* matched in a case-insensitive manner: the configured name is trimmed and converted to
* lower case using {@link Locale#ROOT} before being used as a key.
*/
```
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add table sampler routing entries (precomputed segment subsets) with nPerDay sampler and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2707307061
##########
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;
+ }
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableNameWithType, tableConfig, samplerConfig, _propertyStore);
+
+ SegmentPreSelector samplerSegmentPreSelector =
+ new TableSamplerSegmentPreSelector(segmentPreSelector, sampler);
+ Set<String> samplerPreSelectedOnlineSegments =
+ samplerSegmentPreSelector.preSelect(new HashSet<>(onlineSegments));
Review Comment:
Creating a defensive copy of `onlineSegments` via `new HashSet<>(onlineSegments)` may be unnecessary if the preSelect method doesn't modify the input. If the preSelect contract guarantees the input is not mutated, pass `onlineSegments` directly to avoid the copy overhead.
```suggestion
samplerSegmentPreSelector.preSelect(onlineSegments);
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1210,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);
+ }
+ }
+ if (StringUtils.isNotBlank(samplerName)) {
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ if (samplerEntries != null) {
+ RoutingEntry samplerEntry = samplerEntries.get(samplerName);
+ if (samplerEntry != null) {
+ return samplerEntry;
+ }
Review Comment:
The nested null checks create unnecessary nesting. Consider using a more streamlined approach: `Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType); if (samplerEntries != null) { RoutingEntry samplerEntry = samplerEntries.get(samplerName); if (samplerEntry != null) return samplerEntry; }`
```suggestion
RoutingEntry samplerEntry =
samplerEntries != null ? samplerEntries.get(samplerName) : null;
if (samplerEntry != null) {
return samplerEntry;
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/NPerDaySegmentsTableSampler.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.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;
+
+
+/**
+ * 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 = zkMetadata.getStartTimeMs();
+ if (startTimeMs <= 0) {
+ unknownTimeSegments.add(segmentName);
Review Comment:
Fetching segment ZK metadata individually for each segment (N calls) creates an N+1 query pattern. Consider batch-fetching all segment metadata upfront using `ZKMetadataProvider.getSegmentZKMetadataForTable()` or similar bulk API to reduce ZooKeeper round trips, especially for tables with many segments.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1087,10 +1210,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()) {
Review Comment:
There are redundant null checks. If `isSetPinotQuery()` returns true, `getPinotQuery()` should not return null. Similarly, `isSetQueryOptions()` implies the result is non-null. Simplify by removing the redundant null checks after the `isSet` checks.
```suggestion
if (pinotQuery.isSetQueryOptions()) {
```
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add table sampler routing entries (precomputed segment subsets) with nPerDay sampler and tableSampler query option [pinot]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2784473354
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -351,6 +359,14 @@ private boolean processAssignmentChangeForTable(int idealStateVersion, int exter
return true;
}
routingEntry.onAssignmentChange(idealState, externalView);
+
+ // Handle sampler routing entries
+ Map<String, RoutingEntry> samplerEntries = _samplerRoutingEntryMap.get(tableNameWithType);
+ if (samplerEntries != null && !samplerEntries.isEmpty()) {
Review Comment:
(nit) `samplerEntries` can never be empty
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -812,14 +852,74 @@ private void buildRoutingInternal(String tableNameWithType) {
RoutingEntry routingEntry =
new RoutingEntry(tableNameWithType, idealStatePath, externalViewPath, segmentPreSelector, segmentSelector,
- segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
+ null, segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
timeBoundaryManager, partitionMetadataManager, queryTimeoutMs, !idealState.isEnabled());
if (_routingEntryMap.put(tableNameWithType, routingEntry) == null) {
LOGGER.info("Built routing for table: {}", tableNameWithType);
} else {
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 (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, RoutingEntry> samplerRoutingEntries = new HashMap<>();
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+
+ // Apply sampler after base pre-selection (e.g. segment lineage).
+ Set<String> samplerPreSelectedOnlineSegments =
+ sampler.selectSegments(new HashSet<>(preSelectedOnlineSegments));
+ SegmentSelector samplerSegmentSelector = SegmentSelectorFactory.getSegmentSelector(tableConfig);
+ samplerSegmentSelector.init(idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ List<SegmentPruner> samplerSegmentPruners =
+ SegmentPrunerFactory.getSegmentPruners(tableConfig, _propertyStore);
+ AdaptiveServerSelector samplerAdaptiveServerSelector =
+ AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager, _pinotConfig);
+ InstanceSelector samplerInstanceSelector =
+ InstanceSelectorFactory.getInstanceSelector(tableConfig, _propertyStore, _brokerMetrics,
+ samplerAdaptiveServerSelector, _pinotConfig, _routableServers, _enabledServerInstanceMap,
+ idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ SegmentZkMetadataFetcher samplerSegmentZkMetadataFetcher =
+ new SegmentZkMetadataFetcher(tableNameWithType, _propertyStore);
+ for (SegmentZkMetadataFetchListener listener : samplerSegmentPruners) {
+ samplerSegmentZkMetadataFetcher.register(listener);
+ }
+ samplerSegmentZkMetadataFetcher.init(idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ RoutingEntry samplerRoutingEntry =
+ new RoutingEntry(tableNameWithType, idealStatePath, externalViewPath, segmentPreSelector,
+ samplerSegmentSelector, sampler, samplerSegmentPruners, samplerInstanceSelector, idealStateVersion,
+ externalViewVersion, samplerSegmentZkMetadataFetcher, null, null, queryTimeoutMs,
Review Comment:
Seems time boundary and partitioning is not handled with sampler?
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1177,6 +1311,8 @@ private static class RoutingEntry {
final String _idealStatePath;
final String _externalViewPath;
final SegmentPreSelector _segmentPreSelector;
+ @Nullable
+ final TableSampler _tableSampler;
Review Comment:
Most of the content in the sampler can be shared with the main entry. I think an easier way to handle sampling is to add a SamplerInfo class for sampler specific context (probably only `TableSampler`, `SegmentSelector` and `InstanceSelector`). This way most of the repeated work can be avoided, and time boundary and partitioning info can be automatically supported.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/tablesampler/TableSampler.java:
##########
@@ -0,0 +1,45 @@
+/**
+ * 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.util.Set;
+import org.apache.helix.store.zk.ZkHelixPropertyStore;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.sampler.TableSamplerConfig;
+
+
+/**
+ * A {@code TableSampler} deterministically selects a subset of segments from the set of online segments for a table.
+ *
+ * <p>Selection is performed during routing table build/update so there is no additional per-query overhead beyond
+ * selecting the pre-built routing entry.
+ */
+public interface TableSampler {
+
+ /**
+ * Initializes the sampler for a specific table and sampler config.
+ */
+ void init(TableConfig tableConfig, TableSamplerConfig samplerConfig, ZkHelixPropertyStore<ZNRecord> propertyStore);
+
+ /**
+ * Selects a subset of segments from the provided online segments.
+ */
+ Set<String> selectSegments(Set<String> onlineSegments);
Review Comment:
```suggestion
Set<String> sampleSegments(Set<String> onlineSegments);
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/MultiClusterRoutingManager.java:
##########
Review Comment:
Avoid the repeated logic of parsing sampler name
```suggestion
return getSegments(brokerRequest, extractSamplerName(brokerRequest));
```
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -812,14 +852,74 @@ private void buildRoutingInternal(String tableNameWithType) {
RoutingEntry routingEntry =
new RoutingEntry(tableNameWithType, idealStatePath, externalViewPath, segmentPreSelector, segmentSelector,
- segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
+ null, segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
timeBoundaryManager, partitionMetadataManager, queryTimeoutMs, !idealState.isEnabled());
if (_routingEntryMap.put(tableNameWithType, routingEntry) == null) {
LOGGER.info("Built routing for table: {}", tableNameWithType);
} else {
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 (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, RoutingEntry> samplerRoutingEntries = new HashMap<>();
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+
+ // Apply sampler after base pre-selection (e.g. segment lineage).
+ Set<String> samplerPreSelectedOnlineSegments =
+ sampler.selectSegments(new HashSet<>(preSelectedOnlineSegments));
Review Comment:
Cloning a set upfront involves constant overhead. Does it work if we make the contract that sampler should not modify the input set?
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -812,14 +852,74 @@ private void buildRoutingInternal(String tableNameWithType) {
RoutingEntry routingEntry =
new RoutingEntry(tableNameWithType, idealStatePath, externalViewPath, segmentPreSelector, segmentSelector,
- segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
+ null, segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
timeBoundaryManager, partitionMetadataManager, queryTimeoutMs, !idealState.isEnabled());
if (_routingEntryMap.put(tableNameWithType, routingEntry) == null) {
LOGGER.info("Built routing for table: {}", tableNameWithType);
} else {
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 (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, RoutingEntry> samplerRoutingEntries = new HashMap<>();
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+
+ // Apply sampler after base pre-selection (e.g. segment lineage).
+ Set<String> samplerPreSelectedOnlineSegments =
+ sampler.selectSegments(new HashSet<>(preSelectedOnlineSegments));
+ SegmentSelector samplerSegmentSelector = SegmentSelectorFactory.getSegmentSelector(tableConfig);
+ samplerSegmentSelector.init(idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ List<SegmentPruner> samplerSegmentPruners =
+ SegmentPrunerFactory.getSegmentPruners(tableConfig, _propertyStore);
+ AdaptiveServerSelector samplerAdaptiveServerSelector =
+ AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager, _pinotConfig);
+ InstanceSelector samplerInstanceSelector =
+ InstanceSelectorFactory.getInstanceSelector(tableConfig, _propertyStore, _brokerMetrics,
+ samplerAdaptiveServerSelector, _pinotConfig, _routableServers, _enabledServerInstanceMap,
+ idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ SegmentZkMetadataFetcher samplerSegmentZkMetadataFetcher =
+ new SegmentZkMetadataFetcher(tableNameWithType, _propertyStore);
+ for (SegmentZkMetadataFetchListener listener : samplerSegmentPruners) {
+ samplerSegmentZkMetadataFetcher.register(listener);
+ }
+ samplerSegmentZkMetadataFetcher.init(idealState, externalView, samplerPreSelectedOnlineSegments);
Review Comment:
This can be shared with the main routing
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -812,14 +852,74 @@ private void buildRoutingInternal(String tableNameWithType) {
RoutingEntry routingEntry =
new RoutingEntry(tableNameWithType, idealStatePath, externalViewPath, segmentPreSelector, segmentSelector,
- segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
+ null, segmentPruners, instanceSelector, idealStateVersion, externalViewVersion, segmentZkMetadataFetcher,
timeBoundaryManager, partitionMetadataManager, queryTimeoutMs, !idealState.isEnabled());
if (_routingEntryMap.put(tableNameWithType, routingEntry) == null) {
LOGGER.info("Built routing for table: {}", tableNameWithType);
} else {
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 (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, RoutingEntry> samplerRoutingEntries = new HashMap<>();
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+
+ // Apply sampler after base pre-selection (e.g. segment lineage).
+ Set<String> samplerPreSelectedOnlineSegments =
+ sampler.selectSegments(new HashSet<>(preSelectedOnlineSegments));
+ SegmentSelector samplerSegmentSelector = SegmentSelectorFactory.getSegmentSelector(tableConfig);
+ samplerSegmentSelector.init(idealState, externalView, samplerPreSelectedOnlineSegments);
+
+ List<SegmentPruner> samplerSegmentPruners =
+ SegmentPrunerFactory.getSegmentPruners(tableConfig, _propertyStore);
+ AdaptiveServerSelector samplerAdaptiveServerSelector =
+ AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager, _pinotConfig);
Review Comment:
They can be shared with the main routing
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org
Re: [PR] Add pluggable table samplers with precomputed broker routing entries and tableSampler query option [pinot]
Posted by "Jackie-Jiang (via GitHub)" <gi...@apache.org>.
Jackie-Jiang commented on code in PR #17532:
URL: https://github.com/apache/pinot/pull/17532#discussion_r2791163442
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1248,6 +1324,56 @@ boolean isDisabled() {
return _disabled;
}
+ private void updateSamplerSegments(Set<String> preSelectedOnlineSegments) {
+ if (_samplers.isEmpty()) {
+ _samplerSegments = Collections.emptyMap();
+ return;
+ }
+ Set<String> readOnlyPreSelectedOnlineSegments =
+ Collections.unmodifiableSet(new HashSet<>(preSelectedOnlineSegments));
+ Map<String, Set<String>> samplerSegments = new HashMap<>();
+ for (Map.Entry<String, TableSampler> entry : _samplers.entrySet()) {
+ String samplerName = entry.getKey();
+ try {
+ Set<String> sampledSegments = entry.getValue().sampleSegments(readOnlyPreSelectedOnlineSegments);
+ if (sampledSegments != null) {
+ samplerSegments.put(samplerName, Set.copyOf(sampledSegments));
+ } else {
+ LOGGER.warn("Sampler '{}' returned null segments for table '{}', treating it as empty", samplerName,
+ _tableNameWithType);
+ }
+ } catch (Exception e) {
+ LOGGER.error("Caught unexpected exception while sampling segments for sampler: {} for table: {}", samplerName,
+ _tableNameWithType, e);
+ }
+ }
+ _samplerSegments = samplerSegments.isEmpty() ? Collections.emptyMap() : Map.copyOf(samplerSegments);
+ }
+
+ @Nullable
+ private Set<String> getSampledSegments(@Nullable String samplerName) {
+ if (StringUtils.isNotBlank(samplerName)) {
+ String normalizedSamplerName = normalizeSamplerName(samplerName);
+ Set<String> sampledSegments = _samplerSegments.get(normalizedSamplerName);
+ if (sampledSegments != null) {
+ return sampledSegments;
+ }
+ LOGGER.warn("Requested sampler '{}' not found for table '{}'; falling back to default routing entry",
+ samplerName, _tableNameWithType);
+ }
+ return null;
+ }
+
+ private Set<String> maybeApplySampler(Set<String> selectedSegments, @Nullable String samplerName) {
+ Set<String> sampledSegments = getSampledSegments(samplerName);
+ if (sampledSegments != null && !selectedSegments.isEmpty()) {
+ Set<String> sampledSelectedSegments = new HashSet<>(selectedSegments);
+ sampledSelectedSegments.retainAll(sampledSegments);
Review Comment:
This is expensive. For each sampler, we should have a dedicate `SegmentSelector`. See the previous comment about `SamplerInfo` wrapper class
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1210,7 +1283,10 @@ private static class RoutingEntry {
_timeBoundaryManager = timeBoundaryManager;
_partitionMetadataManager = partitionMetadataManager;
_queryTimeoutMs = queryTimeoutMs;
+ _samplers = samplers.isEmpty() ? Collections.emptyMap() : Map.copyOf(samplers);
Review Comment:
(nit) Use `Map.of()` for empty map; No need to copy a map
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -810,10 +820,44 @@ private void buildRoutingInternal(String tableNameWithType) {
}
segmentZkMetadataFetcher.init(idealState, externalView, preSelectedOnlineSegments);
+ // Build table samplers keyed by normalized sampler name.
+ Map<String, TableSampler> samplers = Collections.emptyMap();
+ List<TableSamplerConfig> tableSamplerConfigs = tableConfig.getTableSamplers();
+ if (CollectionUtils.isNotEmpty(tableSamplerConfigs)) {
+ Map<String, TableSampler> configuredSamplers = new HashMap<>();
+ for (TableSamplerConfig samplerConfig : tableSamplerConfigs) {
+ String samplerName = samplerConfig.getName();
+ String samplerType = samplerConfig.getType();
+ if (StringUtils.isBlank(samplerName) || StringUtils.isBlank(samplerType)) {
+ LOGGER.warn("Skipping invalid table sampler config for table: {}, samplerName: {}, samplerType: {}",
+ tableNameWithType, samplerName, samplerType);
+ continue;
+ }
+ String normalizedSamplerName = normalizeSamplerName(samplerName);
+ if (configuredSamplers.containsKey(normalizedSamplerName)) {
+ LOGGER.warn("Skipping duplicate normalized table sampler name: '{}' for table: {}", samplerName,
+ tableNameWithType);
+ continue;
+ }
+ try {
+ TableSampler sampler = TableSamplerFactory.create(samplerType);
+ sampler.init(tableConfig, samplerConfig, _propertyStore);
+ configuredSamplers.put(normalizedSamplerName, sampler);
+ } catch (Exception e) {
+ LOGGER.error("Caught unexpected exception while building routing for table sampler: {} for table: {}",
+ samplerName, tableNameWithType, e);
+ }
+ }
+ if (!configuredSamplers.isEmpty()) {
+ samplers = Map.copyOf(configuredSamplers);
Review Comment:
(minor) Seems you are making multiple copies. I don't think you need to make any copy. You can create a pre-sized hash map
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,9 +1132,34 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
- return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
+ if (routingEntry == null) {
+ return null;
+ }
+ return routingEntry.getSegments(brokerRequest, samplerName);
+ }
+
+ private static String normalizeSamplerName(String samplerName) {
+ return samplerName.trim().toLowerCase(Locale.ROOT);
+ }
+
+ @Nullable
+ static String extractSamplerName(BrokerRequest brokerRequest) {
+ if (!brokerRequest.isSetPinotQuery()) {
+ return null;
+ }
+ PinotQuery pinotQuery = brokerRequest.getPinotQuery();
+ if (!pinotQuery.isSetQueryOptions()) {
+ return null;
+ }
Review Comment:
Both the checks can be removed
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1248,6 +1324,56 @@ boolean isDisabled() {
return _disabled;
}
+ private void updateSamplerSegments(Set<String> preSelectedOnlineSegments) {
+ if (_samplers.isEmpty()) {
+ _samplerSegments = Collections.emptyMap();
+ return;
+ }
+ Set<String> readOnlyPreSelectedOnlineSegments =
+ Collections.unmodifiableSet(new HashSet<>(preSelectedOnlineSegments));
+ Map<String, Set<String>> samplerSegments = new HashMap<>();
+ for (Map.Entry<String, TableSampler> entry : _samplers.entrySet()) {
+ String samplerName = entry.getKey();
+ try {
+ Set<String> sampledSegments = entry.getValue().sampleSegments(readOnlyPreSelectedOnlineSegments);
+ if (sampledSegments != null) {
+ samplerSegments.put(samplerName, Set.copyOf(sampledSegments));
+ } else {
+ LOGGER.warn("Sampler '{}' returned null segments for table '{}', treating it as empty", samplerName,
+ _tableNameWithType);
+ }
+ } catch (Exception e) {
+ LOGGER.error("Caught unexpected exception while sampling segments for sampler: {} for table: {}", samplerName,
+ _tableNameWithType, e);
+ }
+ }
+ _samplerSegments = samplerSegments.isEmpty() ? Collections.emptyMap() : Map.copyOf(samplerSegments);
Review Comment:
Reduce redundant copy
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1086,9 +1132,34 @@ private Map<ServerInstance, SegmentsToQuery> getServerInstanceToSegmentsMap(Stri
@Nullable
@Override
public List<String> getSegments(BrokerRequest brokerRequest) {
+ return getSegments(brokerRequest, extractSamplerName(brokerRequest));
+ }
+
+ @Nullable
+ @Override
+ public List<String> getSegments(BrokerRequest brokerRequest, @Nullable String samplerName) {
String tableNameWithType = brokerRequest.getQuerySource().getTableName();
RoutingEntry routingEntry = _routingEntryMap.get(tableNameWithType);
- return routingEntry != null ? routingEntry.getSegments(brokerRequest) : null;
+ if (routingEntry == null) {
+ return null;
+ }
+ return routingEntry.getSegments(brokerRequest, samplerName);
+ }
+
+ private static String normalizeSamplerName(String samplerName) {
+ return samplerName.trim().toLowerCase(Locale.ROOT);
+ }
+
+ @Nullable
+ static String extractSamplerName(BrokerRequest brokerRequest) {
+ if (!brokerRequest.isSetPinotQuery()) {
+ return null;
+ }
+ PinotQuery pinotQuery = brokerRequest.getPinotQuery();
+ if (!pinotQuery.isSetQueryOptions()) {
+ return null;
+ }
+ return pinotQuery.getQueryOptions().get(CommonConstants.Broker.Request.QueryOptionKey.TABLE_SAMPLER);
Review Comment:
Move this into `QueryOptionsUtils`
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1248,6 +1324,56 @@ boolean isDisabled() {
return _disabled;
}
+ private void updateSamplerSegments(Set<String> preSelectedOnlineSegments) {
+ if (_samplers.isEmpty()) {
+ _samplerSegments = Collections.emptyMap();
+ return;
+ }
+ Set<String> readOnlyPreSelectedOnlineSegments =
+ Collections.unmodifiableSet(new HashSet<>(preSelectedOnlineSegments));
Review Comment:
Reduce copies. Just make the contract that sampler shouldn't modify the input
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java:
##########
@@ -1044,11 +1088,13 @@ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, long requestId)
@Nullable
@Override
public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String tableNameWithType, long requestId) {
+ String samplerName = extractSamplerName(brokerRequest);
Review Comment:
（minor) Move this after reading `RoutingEntry`
--
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: commits-unsubscribe@pinot.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org