Re: [PR] Support partitionExpression for custom segment partitioning [pinot]

["Copilot (via GitHub)" <gi...@apache.org>]

Copilot commented on code in PR #18449:
URL: https://github.com/apache/pinot/pull/18449#discussion_r3214689183


##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionUtils.java:
##########
@@ -124,7 +124,7 @@ public static Map<String, SegmentPartitionInfo> extractPartitionInfoMap(String t
         continue;
       }
       SegmentPartitionInfo segmentPartitionInfo = new SegmentPartitionInfo(partitionColumn,
-          PartitionFunctionFactory.getPartitionFunction(columnPartitionMetadata),
+          PartitionFunctionFactory.getPartitionFunction(partitionColumn, columnPartitionMetadata),
           columnPartitionMetadata.getPartitions());
       columnSegmentPartitionInfoMap.put(partitionColumn, segmentPartitionInfo);

Review Comment:
   Same as `extractPartitionInfo()`: `extractPartitionInfoMap()` directly calls `PartitionFunctionFactory.getPartitionFunction(...)` which can throw for invalid partition metadata. To avoid a single bad segment metadata record destabilizing routing updates, catch the exception per-column and treat that column as invalid partition info (skip/continue) instead of throwing.



##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionUtils.java:
##########
@@ -80,7 +80,7 @@ public static SegmentPartitionInfo extractPartitionInfo(String tableNameWithType
     }
 
     return new SegmentPartitionInfo(partitionColumn,
-        PartitionFunctionFactory.getPartitionFunction(columnPartitionMetadata),
+        PartitionFunctionFactory.getPartitionFunction(partitionColumn, columnPartitionMetadata),
         columnPartitionMetadata.getPartitions());

Review Comment:
   `extractPartitionInfo()` constructs the partition function without guarding against `PartitionFunctionFactory.getPartitionFunction(...)` throwing (e.g., malformed `partitionExpression` or mismatched partition column in segment metadata). This can propagate out and break routing-table maintenance for the whole table; consider catching `RuntimeException` here and returning `INVALID_PARTITION_INFO` (with a warn log) similar to the JSON parse failure handling above.



##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/metadata/ColumnMetadataImpl.java:
##########
@@ -356,9 +361,12 @@ public static ColumnMetadataImpl fromPropertiesConfiguration(PropertiesConfigura
       } else {
         partitionFunctionConfigMap = null;
       }
-      PartitionFunction partitionFunction =
-          PartitionFunctionFactory.getPartitionFunction(partitionFunctionName, numPartitions,
+      ColumnPartitionConfig columnPartitionMetadata = partitionExpression != null
+          ? ColumnPartitionConfig.forPartitionExpression(partitionExpression, numPartitions)
+          : ColumnPartitionConfig.forPartitionFunction(partitionFunctionName, numPartitions,
               partitionFunctionConfigMap);
+      PartitionFunction partitionFunction =
+          PartitionFunctionFactory.getPartitionFunction(column, columnPartitionMetadata);

Review Comment:
   Variable name `columnPartitionMetadata` is misleading here because the type is `ColumnPartitionConfig` (not metadata). Renaming it (e.g., `columnPartitionConfig`) would reduce confusion, especially since this code also deals with `ColumnPartitionMetadata` elsewhere.
   



##########
pinot-common/src/main/java/org/apache/pinot/common/partition/function/CustomPartitionFunction.java:
##########
@@ -0,0 +1,274 @@
+/**
+ * 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.common.partition.function;
+
+import com.google.common.base.Preconditions;
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.evaluator.InbuiltFunctionEvaluator;
+import org.apache.pinot.segment.spi.V1Constants;
+import org.apache.pinot.segment.spi.partition.PartitionFunction;
+import org.apache.pinot.segment.spi.partition.PartitionIdNormalizer;
+
+
+/// [PartitionFunction] that evaluates a scalar-function expression to produce a partition id.
+///
+/// The expression is configured with `functionConfig.partitionExpression` and can reference the partition column
+/// plus the optional `numPartitions` identifier. It must return an integral partition id in `[0, numPartitions)`;
+/// invalid runtime results map to `-1`.
+@SuppressWarnings("serial")
+public class CustomPartitionFunction implements PartitionFunction {
+  public static final String NAME = "Custom";
+  public static final String PARTITION_COLUMN_KEY = "partitionColumn";
+  public static final String NUM_PARTITIONS_IDENTIFIER = "numPartitions";
+
+  private static final double MAX_PRECISE_DOUBLE_INTEGRAL = 1L << 53;
+  private static final int MAX_EXPRESSION_LENGTH = 512;
+
+  private final int _numPartitions;
+  @Nullable
+  private final String _partitionExpression;
+  @Nullable
+  private final String _partitionColumn;
+  private final List<String> _arguments;
+  private transient volatile ThreadLocal<InbuiltFunctionEvaluator> _evaluatorThreadLocal;
+
+  public CustomPartitionFunction(int numPartitions, @Nullable Map<String, String> functionConfig) {
+    Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0");
+    _numPartitions = numPartitions;
+    if (functionConfig == null) {
+      // Probe-only path used by PartitionFunctionFactory startup scan. Real use supplies
+      // a populated config and reaches the validation below.
+      _partitionExpression = null;
+      _partitionColumn = null;
+      _arguments = Collections.emptyList();
+      return;
+    }

Review Comment:
   The `functionConfig == null` path leaves the instance intentionally unconfigured (for factory name probing), but it also means a runtime lookup of the `Custom` partition function with a missing/absent config will succeed and later fail on `getPartition()` (IllegalStateException). Consider failing fast for non-probe usage (e.g., treat null as an error unless explicitly in probe mode, or initialize in a way that causes immediate validation failure when the function is actually selected from config).



-- 
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