chenboat commented on a change in pull request #6567:
URL: https://github.com/apache/incubator-pinot/pull/6567#discussion_r596469843



##########
File path: 
pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java
##########
@@ -528,4 +542,41 @@ private boolean isValid(Schema schema, IndexingConfig 
indexingConfig) {
     }
     return isValid;
   }
+
+  private int getSegmentPartitionId(String segmentName, String tableName) {
+    // A fast path if the segmentName is a LLC segment name and we can get the 
partition id from the name directly.
+    if (LLCSegmentName.isLowLevelConsumerSegmentName(segmentName)) {
+      return new LLCSegmentName(segmentName).getPartitionId();
+    }
+    // Otherwise, retrieve the partition id from the segment zk metadata. 
Currently only realtime segments from upsert
+    // enabled tables have partition ids in their segment metadata.
+    RealtimeSegmentZKMetadata segmentZKMetadata =
+        
ZKMetadataProvider.getRealtimeSegmentZKMetadata(_helixManager.getHelixPropertyStore(),
 tableName, segmentName);
+    Preconditions.checkState(isUpsertEnabled(),
+        "Only upsert enabled table has partition ids in its segment metadata: 
seg %s of table %s", segmentName,
+        tableName);
+    Preconditions
+        .checkState(segmentZKMetadata != null, "Failed to find segment ZK 
metadata for segment: %s of table: %s",
+            segmentName, tableName);
+    return getSegmentPartitionIdFromZkMetaData(segmentZKMetadata, tableName);
+  }
+
+  private int getSegmentPartitionIdFromZkMetaData(RealtimeSegmentZKMetadata 
segmentZKMetadata, String tableName) {
+    String segmentName = segmentZKMetadata.getSegmentName();
+    Preconditions.checkState(segmentZKMetadata.getPartitionMetadata() != null,
+        "Segment ZK metadata for segment: %s of table: %s does not contain 
partition metadata", segmentName,
+        tableName);
+
+    // Use any primary key column to fetch the partition metadata
+    ColumnPartitionMetadata partitionMetadata =
+        
segmentZKMetadata.getPartitionMetadata().getColumnPartitionMap().get(_primaryKeyColumns.get(0));
+    Preconditions.checkState(partitionMetadata != null,
+        "Segment ZK metadata for segment: %s of table: %s does not contain 
partition metadata for column: %s",
+        segmentName, tableName, _primaryKeyColumns.get(0));
+    Set<Integer> partitions = partitionMetadata.getPartitions();
+    Preconditions.checkState(partitions.size() == 1,
+        "Segment ZK metadata for segment: %s of table: %s contains multiple 
partitions for column: %s", segmentName,

Review comment:
       done.




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

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org

Reply via email to