Add partition function expressions for chained partitioning

Author: xiangfu0

pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java

@@ -86,6 +86,8 @@
 import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService;
 import org.apache.pinot.core.transport.ServerInstance;
 import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.segment.spi.partition.PartitionFunction;
+import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory;
 import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
 import org.apache.pinot.spi.config.table.QueryConfig;
 import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
@@ -797,9 +799,13 @@ private void buildRoutingInternal(String tableNameWithType) {
                 columnPartitionMap.entrySet().iterator().next();
             LOGGER.info("Enabling SegmentPartitionMetadataManager for table: {} on partition column: {}",
                 tableNameWithType, partitionConfig.getKey());
+            PartitionFunction partitionFunction =
+                PartitionFunctionFactory.getPartitionFunction(partitionConfig.getKey(), partitionConfig.getValue());
             partitionMetadataManager =
                 new SegmentPartitionMetadataManager(tableNameWithType, partitionConfig.getKey(),
-                    partitionConfig.getValue().getFunctionName(), partitionConfig.getValue().getNumPartitions());
+                    partitionFunction.getName(), partitionFunction.getFunctionExpr(),
+                    partitionFunction.getPartitionIdNormalizer(),
+                    partitionFunction.getNumPartitions());
           } else {
             LOGGER.warn(
                 "Cannot enable SegmentPartitionMetadataManager for table: {} with multiple partition columns: {}",

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManager.java

@@ -24,6 +24,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
@@ -62,6 +63,8 @@ public class SegmentPartitionMetadataManager implements SegmentZkMetadataFetchLi
   // static content, if anything changes for the following. a rebuild of routing table is needed.
   private final String _partitionColumn;
   private final String _partitionFunctionName;
+  private final String _partitionFunctionExpr;
+  private final String _partitionIdNormalizer;
   private final int _numPartitions;
 
   // cache-able content, only follow changes if onlineSegments list (of ideal-state) is changed.
@@ -72,10 +75,12 @@ public class SegmentPartitionMetadataManager implements SegmentZkMetadataFetchLi
   private transient TablePartitionReplicatedServersInfo _tablePartitionReplicatedServersInfo;
 
   public SegmentPartitionMetadataManager(String tableNameWithType, String partitionColumn, String partitionFunctionName,
-      int numPartitions) {
+      @Nullable String partitionFunctionExpr, @Nullable String partitionIdNormalizer, int numPartitions) {
     _tableNameWithType = tableNameWithType;
     _partitionColumn = partitionColumn;
     _partitionFunctionName = partitionFunctionName;
+    _partitionFunctionExpr = partitionFunctionExpr;
+    _partitionIdNormalizer = partitionIdNormalizer;
     _numPartitions = numPartitions;
   }
 
@@ -103,7 +108,7 @@ private int getPartitionId(String segment, @Nullable ZNRecord znRecord) {
       return INVALID_PARTITION_ID;
     }
     PartitionFunction partitionFunction = segmentPartitionInfo.getPartitionFunction();
-    if (!_partitionFunctionName.equalsIgnoreCase(partitionFunction.getName())) {
+    if (!isMatchingPartitionFunction(partitionFunction)) {
       return INVALID_PARTITION_ID;
     }
     if (_numPartitions != partitionFunction.getNumPartitions()) {
@@ -116,6 +121,14 @@ private int getPartitionId(String segment, @Nullable ZNRecord znRecord) {
     return partitions.iterator().next();
   }
 
+  private boolean isMatchingPartitionFunction(PartitionFunction partitionFunction) {
+    if (_partitionFunctionExpr != null) {
+      return Objects.equals(_partitionFunctionExpr, partitionFunction.getFunctionExpr()) && Objects.equals(
+          _partitionIdNormalizer, partitionFunction.getPartitionIdNormalizer());
+    }
+    return _partitionFunctionName.equalsIgnoreCase(partitionFunction.getName());
+  }
+
   private static long getCreationTimeMs(@Nullable ZNRecord znRecord) {
     if (znRecord == null) {
       return INVALID_CREATION_TIME_MS;

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionUtils.java

@@ -25,6 +25,7 @@
 import javax.annotation.Nullable;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.pinot.common.metadata.segment.SegmentPartitionMetadata;
+import org.apache.pinot.segment.spi.partition.PartitionFunction;
 import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory;
 import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
 import org.apache.pinot.spi.utils.CommonConstants;
@@ -79,10 +80,17 @@ public static SegmentPartitionInfo extractPartitionInfo(String tableNameWithType
       return INVALID_PARTITION_INFO;
     }
 
-    return new SegmentPartitionInfo(partitionColumn,
-        PartitionFunctionFactory.getPartitionFunction(columnPartitionMetadata.getFunctionName(),
-            columnPartitionMetadata.getNumPartitions(), columnPartitionMetadata.getFunctionConfig()),
-        columnPartitionMetadata.getPartitions());
+    try {
+      PartitionFunction partitionFunction =
+          PartitionFunctionFactory.getPartitionFunction(partitionColumn, columnPartitionMetadata.getFunctionName(),
+              columnPartitionMetadata.getNumPartitions(), columnPartitionMetadata.getFunctionConfig(),
+              columnPartitionMetadata.getFunctionExpr(), columnPartitionMetadata.getPartitionIdNormalizer());
+      return new SegmentPartitionInfo(partitionColumn, partitionFunction, columnPartitionMetadata.getPartitions());
+    } catch (Exception e) {
+      LOGGER.warn("Caught exception while constructing partition function for column: {}, segment: {}, table: {}",
+          partitionColumn, segment, tableNameWithType, e);
+      return INVALID_PARTITION_INFO;
+    }
   }
 
   /**
@@ -124,11 +132,18 @@ public static Map<String, SegmentPartitionInfo> extractPartitionInfoMap(String t
             segment, tableNameWithType);
         continue;
       }
-      SegmentPartitionInfo segmentPartitionInfo = new SegmentPartitionInfo(partitionColumn,
-          PartitionFunctionFactory.getPartitionFunction(columnPartitionMetadata.getFunctionName(),
-              columnPartitionMetadata.getNumPartitions(), columnPartitionMetadata.getFunctionConfig()),
-          columnPartitionMetadata.getPartitions());
-      columnSegmentPartitionInfoMap.put(partitionColumn, segmentPartitionInfo);
+      try {
+        PartitionFunction partitionFunction =
+            PartitionFunctionFactory.getPartitionFunction(partitionColumn, columnPartitionMetadata.getFunctionName(),
+                columnPartitionMetadata.getNumPartitions(), columnPartitionMetadata.getFunctionConfig(),
+                columnPartitionMetadata.getFunctionExpr(), columnPartitionMetadata.getPartitionIdNormalizer());
+        SegmentPartitionInfo segmentPartitionInfo =
+            new SegmentPartitionInfo(partitionColumn, partitionFunction, columnPartitionMetadata.getPartitions());
+        columnSegmentPartitionInfoMap.put(partitionColumn, segmentPartitionInfo);
+      } catch (Exception e) {
+        LOGGER.warn("Caught exception while constructing partition function for column: {}, segment: {}, table: {}",
+            partitionColumn, segment, tableNameWithType, e);
+      }
     }
     if (columnSegmentPartitionInfoMap.size() == 1) {
       String partitionColumn = columnSegmentPartitionInfoMap.keySet().iterator().next();

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/MultiPartitionColumnsSegmentPruner.java

@@ -36,13 +36,17 @@
 import org.apache.pinot.common.request.Identifier;
 import org.apache.pinot.common.request.context.RequestContextUtils;
 import org.apache.pinot.sql.FilterKind;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
  * The {@code MultiPartitionColumnsSegmentPruner} prunes segments based on their partition metadata stored in ZK. The
  * pruner supports queries with filter (or nested filter) of EQUALITY and IN predicates.
  */
 public class MultiPartitionColumnsSegmentPruner implements SegmentPruner {
+  private static final Logger LOGGER = LoggerFactory.getLogger(MultiPartitionColumnsSegmentPruner.class);
+
   private final String _tableNameWithType;
   private final Set<String> _partitionColumns;
   private final Map<String, Map<String, SegmentPartitionInfo>> _segmentColumnPartitionInfoMap =
@@ -140,8 +144,8 @@ private boolean isPartitionMatch(Expression filterExpression,
         Identifier identifier = operands.get(0).getIdentifier();
         if (identifier != null) {
           SegmentPartitionInfo partitionInfo = columnPartitionInfoMap.get(identifier.getName());
-          return partitionInfo == null || partitionInfo.getPartitions().contains(
-              partitionInfo.getPartitionFunction().getPartition(RequestContextUtils.getStringValue(operands.get(1))));
+          return partitionInfo == null || isPartitionMatch(partitionInfo, RequestContextUtils.getStringValue(
+              operands.get(1)));
         } else {
           return true;
         }
@@ -155,8 +159,7 @@ private boolean isPartitionMatch(Expression filterExpression,
           }
           int numOperands = operands.size();
           for (int i = 1; i < numOperands; i++) {
-            if (partitionInfo.getPartitions().contains(partitionInfo.getPartitionFunction()
-                .getPartition(RequestContextUtils.getStringValue(operands.get(i))))) {
+            if (isPartitionMatch(partitionInfo, RequestContextUtils.getStringValue(operands.get(i)))) {
               return true;
             }
           }
@@ -169,4 +172,14 @@ private boolean isPartitionMatch(Expression filterExpression,
         return true;
     }
   }
+
+  private boolean isPartitionMatch(SegmentPartitionInfo partitionInfo, String value) {
+    try {
+      return partitionInfo.getPartitions().contains(partitionInfo.getPartitionFunction().getPartition(value));
+    } catch (RuntimeException e) {
+      LOGGER.debug("Failed to evaluate partition function for table: {}, partition column: {}; skipping partition "
+          + "pruning", _tableNameWithType, partitionInfo.getPartitionColumn(), e);
+      return true;
+    }
+  }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java

@@ -35,13 +35,17 @@
 import org.apache.pinot.common.request.Identifier;
 import org.apache.pinot.common.request.context.RequestContextUtils;
 import org.apache.pinot.sql.FilterKind;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 
 /**
  * The {@code SinglePartitionColumnSegmentPruner} prunes segments based on their partition metadata stored in ZK. The
  * pruner supports queries with filter (or nested filter) of EQUALITY and IN predicates.
  */
 public class SinglePartitionColumnSegmentPruner implements SegmentPruner {
+  private static final Logger LOGGER = LoggerFactory.getLogger(SinglePartitionColumnSegmentPruner.class);
+
   private final String _tableNameWithType;
   private final String _partitionColumn;
   private final Map<String, SegmentPartitionInfo> _partitionInfoMap = new ConcurrentHashMap<>();
@@ -129,8 +133,7 @@ private boolean isPartitionMatch(Expression filterExpression, SegmentPartitionIn
       case EQUALS: {
         Identifier identifier = operands.get(0).getIdentifier();
         if (identifier != null && identifier.getName().equals(_partitionColumn)) {
-          return partitionInfo.getPartitions().contains(partitionInfo.getPartitionFunction()
-              .getPartition(RequestContextUtils.getStringValue(operands.get(1))));
+          return isPartitionMatch(partitionInfo, RequestContextUtils.getStringValue(operands.get(1)));
         } else {
           return true;
         }
@@ -140,8 +143,7 @@ private boolean isPartitionMatch(Expression filterExpression, SegmentPartitionIn
         if (identifier != null && identifier.getName().equals(_partitionColumn)) {
           int numOperands = operands.size();
           for (int i = 1; i < numOperands; i++) {
-            if (partitionInfo.getPartitions().contains(partitionInfo.getPartitionFunction()
-                .getPartition(RequestContextUtils.getStringValue(operands.get(i))))) {
+            if (isPartitionMatch(partitionInfo, RequestContextUtils.getStringValue(operands.get(i)))) {
               return true;
             }
           }
@@ -154,4 +156,14 @@ private boolean isPartitionMatch(Expression filterExpression, SegmentPartitionIn
         return true;
     }
   }
+
+  private boolean isPartitionMatch(SegmentPartitionInfo partitionInfo, String value) {
+    try {
+      return partitionInfo.getPartitions().contains(partitionInfo.getPartitionFunction().getPartition(value));
+    } catch (RuntimeException e) {
+      LOGGER.debug("Failed to evaluate partition function for table: {}, partition column: {}; skipping partition "
+          + "pruning", _tableNameWithType, _partitionColumn, e);
+      return true;
+    }
+  }
 }

pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpartition/SegmentPartitionMetadataManagerTest.java

@@ -23,6 +23,7 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import javax.annotation.Nullable;
 import org.apache.helix.manager.zk.ZkBaseDataAccessor;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
@@ -37,6 +38,7 @@
 import org.apache.pinot.controller.helix.ControllerTest;
 import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo;
 import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
+import org.apache.pinot.segment.spi.partition.pipeline.PartitionPipelineFunction;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -84,7 +86,7 @@ public void testPartitionMetadataManagerProcessingThroughSegmentChangesSinglePar
 
     SegmentPartitionMetadataManager partitionMetadataManager =
         new SegmentPartitionMetadataManager(OFFLINE_TABLE_NAME, PARTITION_COLUMN, PARTITION_COLUMN_FUNC,
-            NUM_PARTITIONS);
+            null, null, NUM_PARTITIONS);
     SegmentZkMetadataFetcher segmentZkMetadataFetcher =
         new SegmentZkMetadataFetcher(OFFLINE_TABLE_NAME, _propertyStore);
     segmentZkMetadataFetcher.register(partitionMetadataManager);
@@ -284,11 +286,87 @@ public void testPartitionMetadataManagerProcessingThroughSegmentChangesSinglePar
     assertEquals(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().get(0), segmentInvalid);
   }
 
+  @Test
+  public void testPartitionMetadataManagerProcessingWithFunctionExpr() {
+    ExternalView externalView = new ExternalView(OFFLINE_TABLE_NAME);
+    Map<String, Map<String, String>> segmentAssignment = externalView.getRecord().getMapFields();
+    Set<String> onlineSegments = new HashSet<>();
+    IdealState idealState = new IdealState(OFFLINE_TABLE_NAME);
+    String functionExpr = "fnv1a_32(md5(" + PARTITION_COLUMN + "))";
+    String segment = "exprSegment";
+
+    SegmentPartitionMetadataManager partitionMetadataManager =
+        new SegmentPartitionMetadataManager(OFFLINE_TABLE_NAME, PARTITION_COLUMN, PartitionPipelineFunction.NAME,
+            functionExpr, "MASK", 8);
+    SegmentZkMetadataFetcher segmentZkMetadataFetcher =
+        new SegmentZkMetadataFetcher(OFFLINE_TABLE_NAME, _propertyStore);
+    segmentZkMetadataFetcher.register(partitionMetadataManager);
+    segmentZkMetadataFetcher.init(idealState, externalView, onlineSegments);
+
+    onlineSegments.add(segment);
+    segmentAssignment.put(segment, Collections.singletonMap(SERVER_0, ONLINE));
+    setSegmentZKMetadata(segment, PartitionPipelineFunction.NAME, 8, 1, functionExpr, "MASK", 0L);
+    segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments);
+
+    TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo =
+        partitionMetadataManager.getTablePartitionReplicatedServersInfo();
+    assertEquals(tablePartitionReplicatedServersInfo.getPartitionInfoMap()[1]._segments,
+        Collections.singleton(segment));
+    assertTrue(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().isEmpty());
+
+    setSegmentZKMetadata(segment, PartitionPipelineFunction.NAME, 8, 1, functionExpr, "POSITIVE_MODULO", 0L);
+    segmentZkMetadataFetcher.refreshSegment(segment);
+
+    tablePartitionReplicatedServersInfo = partitionMetadataManager.getTablePartitionReplicatedServersInfo();
+    assertEquals(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition(),
+        Collections.singletonList(segment));
+  }
+
+  @Test
+  public void testPartitionMetadataManagerProcessingWithFunctionExprAbsNormalizer() {
+    ExternalView externalView = new ExternalView(OFFLINE_TABLE_NAME);
+    Map<String, Map<String, String>> segmentAssignment = externalView.getRecord().getMapFields();
+    Set<String> onlineSegments = new HashSet<>();
+    IdealState idealState = new IdealState(OFFLINE_TABLE_NAME);
+    String functionExpr = "murmur2(" + PARTITION_COLUMN + ")";
+    String segment = "exprAbsSegment";
+
+    SegmentPartitionMetadataManager partitionMetadataManager =
+        new SegmentPartitionMetadataManager(OFFLINE_TABLE_NAME, PARTITION_COLUMN, PartitionPipelineFunction.NAME,
+            functionExpr, "ABS", 8);
+    SegmentZkMetadataFetcher segmentZkMetadataFetcher =
+        new SegmentZkMetadataFetcher(OFFLINE_TABLE_NAME, _propertyStore);
+    segmentZkMetadataFetcher.register(partitionMetadataManager);
+    segmentZkMetadataFetcher.init(idealState, externalView, onlineSegments);
+
+    onlineSegments.add(segment);
+    segmentAssignment.put(segment, Collections.singletonMap(SERVER_0, ONLINE));
+    setSegmentZKMetadata(segment, PartitionPipelineFunction.NAME, 8, 1, functionExpr, "ABS", 0L);
+    segmentZkMetadataFetcher.onAssignmentChange(idealState, externalView, onlineSegments);
+
+    TablePartitionReplicatedServersInfo tablePartitionReplicatedServersInfo =
+        partitionMetadataManager.getTablePartitionReplicatedServersInfo();
+    assertEquals(tablePartitionReplicatedServersInfo.getPartitionInfoMap()[1]._segments,
+        Collections.singleton(segment));
+    assertTrue(tablePartitionReplicatedServersInfo.getSegmentsWithInvalidPartition().isEmpty());
+  }
+
   private void setSegmentZKMetadata(String segment, String partitionFunction, int numPartitions, int partitionId,
       long creationTimeMs) {
+    setSegmentZKMetadata(segment, partitionFunction, numPartitions, partitionId, null, null, creationTimeMs);
+  }
+
+  private void setSegmentZKMetadata(String segment, String partitionFunction, int numPartitions, int partitionId,
+      @Nullable String functionExpr, long creationTimeMs) {
+    setSegmentZKMetadata(segment, partitionFunction, numPartitions, partitionId, functionExpr, null, creationTimeMs);
+  }
+
+  private void setSegmentZKMetadata(String segment, String partitionFunction, int numPartitions, int partitionId,
+      @Nullable String functionExpr, @Nullable String partitionIdNormalizer, long creationTimeMs) {
     SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment);
     segmentZKMetadata.setPartitionMetadata(new SegmentPartitionMetadata(Collections.singletonMap(PARTITION_COLUMN,
-        new ColumnPartitionMetadata(partitionFunction, numPartitions, Collections.singleton(partitionId), null))));
+        new ColumnPartitionMetadata(partitionFunction, numPartitions, Collections.singleton(partitionId), null,
+            functionExpr, partitionIdNormalizer))));
     segmentZKMetadata.setCreationTime(creationTimeMs);
     ZKMetadataProvider.setSegmentZKMetadata(_propertyStore, OFFLINE_TABLE_NAME, segmentZKMetadata);
   }