zhtaoxiang commented on code in PR #13636: URL: https://github.com/apache/pinot/pull/13636#discussion_r1712305564
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/dedup/ConcurrentMapPartitionDedupMetadataManager.java: ########## @@ -19,105 +19,80 @@ package org.apache.pinot.segment.local.dedup; import com.google.common.annotations.VisibleForTesting; -import java.util.HashMap; +import com.google.common.util.concurrent.AtomicDouble; import java.util.Iterator; -import java.util.List; -import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.metrics.ServerGauge; -import org.apache.pinot.common.metrics.ServerMetrics; -import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReader; import org.apache.pinot.segment.local.utils.HashUtils; import org.apache.pinot.segment.spi.IndexSegment; -import org.apache.pinot.spi.config.table.HashFunction; -import org.apache.pinot.spi.data.readers.PrimaryKey; -import org.apache.pinot.spi.utils.ByteArray; -class ConcurrentMapPartitionDedupMetadataManager implements PartitionDedupMetadataManager { - private final String _tableNameWithType; - private final List<String> _primaryKeyColumns; - private final int _partitionId; - private final ServerMetrics _serverMetrics; - private final HashFunction _hashFunction; +class ConcurrentMapPartitionDedupMetadataManager extends BasePartitionDedupMetadataManager { @VisibleForTesting - final ConcurrentHashMap<Object, IndexSegment> _primaryKeyToSegmentMap = new ConcurrentHashMap<>(); + final AtomicDouble _largestSeenTime = new AtomicDouble(0); + @VisibleForTesting + final ConcurrentHashMap<Object, Pair<IndexSegment, Double>> _primaryKeyToSegmentAndTimeMap = + new ConcurrentHashMap<>(); - public ConcurrentMapPartitionDedupMetadataManager(String tableNameWithType, List<String> primaryKeyColumns, - int partitionId, ServerMetrics serverMetrics, HashFunction hashFunction) { - _tableNameWithType = tableNameWithType; - _primaryKeyColumns = primaryKeyColumns; - _partitionId = partitionId; - _serverMetrics = serverMetrics; - _hashFunction = hashFunction; + protected ConcurrentMapPartitionDedupMetadataManager(String tableNameWithType, int partitionId, + DedupContext dedupContext) { + super(tableNameWithType, partitionId, dedupContext); } - public void addSegment(IndexSegment segment) { - // Add all PKs to _primaryKeyToSegmentMap - Iterator<PrimaryKey> primaryKeyIterator = getPrimaryKeyIterator(segment); - while (primaryKeyIterator.hasNext()) { - PrimaryKey pk = primaryKeyIterator.next(); - _primaryKeyToSegmentMap.put(HashUtils.hashPrimaryKey(pk, _hashFunction), segment); + @Override + protected void addSegment(IndexSegment segment, Iterator<DedupRecordInfo> dedupRecordInfoIterator) { + while (dedupRecordInfoIterator.hasNext()) { + DedupRecordInfo dedupRecordInfo = dedupRecordInfoIterator.next(); + double metadataTime = dedupRecordInfo.getDedupTime(); + _largestSeenTime.getAndUpdate(time -> Math.max(time, metadataTime)); + _primaryKeyToSegmentAndTimeMap.compute(HashUtils.hashPrimaryKey(dedupRecordInfo.getPrimaryKey(), _hashFunction), + (primaryKey, segmentAndTime) -> { + if (segmentAndTime == null || segmentAndTime.getRight() < metadataTime) { + return Pair.of(segment, metadataTime); + } else { + return segmentAndTime; + } + }); } - _serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, ServerGauge.DEDUP_PRIMARY_KEYS_COUNT, - _primaryKeyToSegmentMap.size()); } - public void removeSegment(IndexSegment segment) { - // TODO(saurabh): Explain reload scenario here - Iterator<PrimaryKey> primaryKeyIterator = getPrimaryKeyIterator(segment); - while (primaryKeyIterator.hasNext()) { - PrimaryKey pk = primaryKeyIterator.next(); - _primaryKeyToSegmentMap.compute(HashUtils.hashPrimaryKey(pk, _hashFunction), (primaryKey, currentSegment) -> { - if (currentSegment == segment) { - return null; - } else { - return currentSegment; - } - }); + @Override + protected void removeSegment(IndexSegment segment, Iterator<DedupRecordInfo> dedupRecordInfoIterator) { + while (dedupRecordInfoIterator.hasNext()) { + DedupRecordInfo dedupRecordInfo = dedupRecordInfoIterator.next(); + _primaryKeyToSegmentAndTimeMap.computeIfPresent( + HashUtils.hashPrimaryKey(dedupRecordInfo.getPrimaryKey(), _hashFunction), (primaryKey, segmentAndTime) -> { + if (segmentAndTime.getLeft() == segment && segmentAndTime.getRight() == dedupRecordInfo.getDedupTime()) { + return null; + } else { + return segmentAndTime; + } + }); } - _serverMetrics.setValueOfPartitionGauge(_tableNameWithType, _partitionId, ServerGauge.DEDUP_PRIMARY_KEYS_COUNT, - _primaryKeyToSegmentMap.size()); } - @VisibleForTesting - Iterator<PrimaryKey> getPrimaryKeyIterator(IndexSegment segment) { - Map<String, PinotSegmentColumnReader> columnToReaderMap = new HashMap<>(); - for (String primaryKeyColumn : _primaryKeyColumns) { - columnToReaderMap.put(primaryKeyColumn, new PinotSegmentColumnReader(segment, primaryKeyColumn)); + @Override + public int removeExpiredPrimaryKeys() { + if (_metadataTTL > 0) { + double smallestTimeToKeep = _largestSeenTime.get() - _metadataTTL; + _primaryKeyToSegmentAndTimeMap.entrySet().removeIf(entry -> entry.getValue().getRight() < smallestTimeToKeep); } - int numTotalDocs = segment.getSegmentMetadata().getTotalDocs(); - int numPrimaryKeyColumns = _primaryKeyColumns.size(); - return new Iterator<PrimaryKey>() { - private int _docId = 0; - - @Override - public boolean hasNext() { - return _docId < numTotalDocs; - } - - @Override - public PrimaryKey next() { - Object[] values = new Object[numPrimaryKeyColumns]; - for (int i = 0; i < numPrimaryKeyColumns; i++) { - Object value = columnToReaderMap.get(_primaryKeyColumns.get(i)).getValue(_docId); - if (value instanceof byte[]) { - value = new ByteArray((byte[]) value); - } - values[i] = value; - } - _docId++; - return new PrimaryKey(values); - } - }; + return _primaryKeyToSegmentAndTimeMap.size(); } - public boolean checkRecordPresentOrUpdate(PrimaryKey pk, IndexSegment indexSegment) { - boolean present = - _primaryKeyToSegmentMap.putIfAbsent(HashUtils.hashPrimaryKey(pk, _hashFunction), indexSegment) != null; + @Override + public boolean dropOrAddRecord(DedupRecordInfo dedupRecordInfo, IndexSegment indexSegment) { + if (_metadataTTL > 0 && dedupRecordInfo.getDedupTime() < _largestSeenTime.get() - _metadataTTL) { + return true; + } + _largestSeenTime.getAndUpdate(time -> Math.max(time, dedupRecordInfo.getDedupTime())); + boolean present = _primaryKeyToSegmentAndTimeMap.putIfAbsent( + HashUtils.hashPrimaryKey(dedupRecordInfo.getPrimaryKey(), _hashFunction), + Pair.of(indexSegment, dedupRecordInfo.getDedupTime())) != null; Review Comment: oh, good catch! we don't need 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-unsubscr...@pinot.apache.org 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