stevenzwu commented on code in PR #9321: URL: https://github.com/apache/iceberg/pull/9321#discussion_r1434407274
########## flink/v1.17/flink/src/main/java/org/apache/iceberg/flink/sink/shuffle/MapRangePartitioner.java: ########## @@ -0,0 +1,288 @@ +/* + * 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.iceberg.flink.sink.shuffle; + +import java.util.Arrays; +import java.util.Comparator; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.NavigableMap; +import java.util.concurrent.ThreadLocalRandom; +import org.apache.flink.api.common.functions.Partitioner; +import org.apache.flink.table.data.RowData; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SortKey; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.SortOrderComparators; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.flink.FlinkSchemaUtil; +import org.apache.iceberg.flink.RowDataWrapper; +import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting; +import org.apache.iceberg.relocated.com.google.common.base.MoreObjects; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.util.ArrayUtil; +import org.apache.iceberg.util.Pair; + +/** + * Internal partitioner implementation that supports MapDataStatistics, which is typically used for + * low-cardinality use cases. While MapDataStatistics can keep accurate counters, it can't be used + * for high-cardinality use cases. Otherwise, the memory footprint is too high. + */ +class MapRangePartitioner implements Partitioner<RowData> { + private final RowDataWrapper rowDataWrapper; + private final SortKey sortKey; + private final Comparator<StructLike> comparator; + private final Map<SortKey, Long> mapStatistics; + private final double closeFileCostInWeightPercentage; + + // lazily computed due to the need of numPartitions + private Map<SortKey, KeyAssignment> assignment; + private NavigableMap<SortKey, Long> sortedStatsWithCloseFileCost; + + MapRangePartitioner( + Schema schema, + SortOrder sortOrder, + MapDataStatistics dataStatistics, + double closeFileCostInWeightPercentage) { + this.rowDataWrapper = new RowDataWrapper(FlinkSchemaUtil.convert(schema), schema.asStruct()); + this.sortKey = new SortKey(schema, sortOrder); + this.comparator = SortOrderComparators.forSchema(schema, sortOrder); + this.mapStatistics = dataStatistics.statistics(); + this.closeFileCostInWeightPercentage = closeFileCostInWeightPercentage; + } + + @Override + public int partition(RowData row, int numPartitions) { + // assignment table can only be built lazily when first referenced here, + // because number of partitions (downstream subtasks) is needed + Map<SortKey, KeyAssignment> assignmentMap = assignment(numPartitions); + // reuse the sortKey and rowDataWrapper + sortKey.wrap(rowDataWrapper.wrap(row)); + KeyAssignment keyAssignment = assignmentMap.get(sortKey); + if (keyAssignment == null) { + // haven't learned about the key before. fall back to random selection. + return ThreadLocalRandom.current().nextInt(numPartitions); + } + + return keyAssignment.select(); + } + + @VisibleForTesting + Map<SortKey, KeyAssignment> assignment(int numPartitions) { + if (assignment == null) { + long totalWeight = mapStatistics.values().stream().mapToLong(l -> l).sum(); + double targetWeightPerSubtask = ((double) totalWeight) / numPartitions; + long closeFileCostInWeight = + (long) Math.ceil(targetWeightPerSubtask * closeFileCostInWeightPercentage / 100); + + // add one close file cost for each key even if a key with large weight may be assigned to + // multiple subtasks + this.sortedStatsWithCloseFileCost = Maps.newTreeMap(comparator); + mapStatistics.forEach( + (k, v) -> { + int estimatedSplits = (int) Math.ceil(v / targetWeightPerSubtask); + long estimatedCloseFileCost = closeFileCostInWeight * estimatedSplits; + sortedStatsWithCloseFileCost.put(k, v + estimatedCloseFileCost); + }); + + long totalWeightWithCloseFileCost = + sortedStatsWithCloseFileCost.values().stream().mapToLong(l -> l).sum(); + long targetWeightPerSubtaskWithCloseFileCost = + (long) Math.ceil(((double) totalWeightWithCloseFileCost) / numPartitions); + this.assignment = + buildAssignment( + numPartitions, sortedStatsWithCloseFileCost, targetWeightPerSubtaskWithCloseFileCost); + } + + return assignment; + } + + @VisibleForTesting + NavigableMap<SortKey, Long> sortedStatsWithCloseFileCost() { + return sortedStatsWithCloseFileCost; + } + + /** + * @return assignment summary for every subtask. Key is subtaskId. Value pair is (weight assigned + * to the subtask, number of keys assigned to the subtask) + */ + Map<Integer, Pair<Long, Integer>> assignmentInfo() { + Map<Integer, Pair<Long, Integer>> assignmentInfo = Maps.newTreeMap(); + assignment.forEach( + (key, keyAssignment) -> { + for (int i = 0; i < keyAssignment.assignedSubtasks.length; ++i) { + int subtaskId = keyAssignment.assignedSubtasks[i]; + long subtaskWeight = keyAssignment.subtaskWeights[i]; + Pair<Long, Integer> oldValue = assignmentInfo.getOrDefault(subtaskId, Pair.of(0L, 0)); + assignmentInfo.put( + subtaskId, Pair.of(oldValue.first() + subtaskWeight, oldValue.second() + 1)); + } + }); + + return assignmentInfo; + } + + private Map<SortKey, KeyAssignment> buildAssignment( + int numPartitions, + NavigableMap<SortKey, Long> sortedStatistics, + long targetWeightPerSubtask) { + Map<SortKey, KeyAssignment> assignmentMap = + Maps.newHashMapWithExpectedSize(sortedStatistics.size()); + Iterator<SortKey> mapKeyIterator = sortedStatistics.keySet().iterator(); + int subtaskId = 0; + SortKey currentKey = null; + long keyRemainingWeight = 0L; + long subtaskRemainingWeight = targetWeightPerSubtask; + List<Integer> assignedSubtasks = Lists.newArrayList(); + List<Long> subtaskWeights = Lists.newArrayList(); + while (mapKeyIterator.hasNext() && subtaskId < numPartitions) { + if (currentKey == null) { + currentKey = mapKeyIterator.next(); + keyRemainingWeight = sortedStatistics.get(currentKey); + } + + assignedSubtasks.add(subtaskId); + // assign the remaining weight of key to the current subtask if it is the last subtask + // or if the subtask has more capacity than the remaining key weight + if (subtaskId == numPartitions - 1 || keyRemainingWeight < subtaskRemainingWeight) { + subtaskWeights.add(keyRemainingWeight); + subtaskRemainingWeight -= keyRemainingWeight; + keyRemainingWeight = 0; + } else { + // filled up the current subtask + subtaskWeights.add(subtaskRemainingWeight); + keyRemainingWeight -= subtaskRemainingWeight; + // move on to the next subtask + subtaskId += 1; + subtaskRemainingWeight = targetWeightPerSubtask; + } + + if (keyRemainingWeight <= 0) { + // finishing up the assignment for the current key + Preconditions.checkState( + assignedSubtasks.size() == subtaskWeights.size(), + "List size mismatch: assigned subtasks = %s, subtask weights = %s", + assignedSubtasks, + subtaskWeights); + KeyAssignment keyAssignment = + new KeyAssignment( + ArrayUtil.toIntArray(assignedSubtasks), ArrayUtil.toLongArray(subtaskWeights)); Review Comment: let me make sure I understand your example correctly - target weight before closing cost is 9. because `SORT_KEY_0` has weight 10. it should be split into 2 trunks. hence the closing weight for the key is 2 (1x2 trunks). and total weight (including closing cost) is 12 for SORT_KEY_0. - target weight with closing cost is 10 - SORT_KEY_0 will be assigned to 2 subtasks: subtask 0 with weight 10 and subtask 1 with weight 2 the bin packing in the last step is fine. but the `select` function should have used the weight of 9 and 1 (not 10 and 2) when route the traffic. I agree in this case, 9 and 1 are more accurate. I want to discuss a scenario from the unit test ``` // target subtask weight is 10 before close file cost factored in. // close file cost is 2 = 20% * 10. // key weights before and after close file cost factored in // before: 35, 23, 12, 4, 1, 1, 1, 1, 1, 1 // close-cost: 8, 6, 4, 2, 2, 2, 2, 2, 2, 2 // after: 43, 29, 16, 6, 3, 3, 3, 3, 3, 3 // target subtask weight per subtask is 14 (112/8) ``` for key0, the bin packing weight distribution will be 14, 14, 14, 1, because we use the total weight with closing cost. then what weights should `select` use in this case? assuming it should be the weight before closing cost. then it will be 12, 12, 12, -1/0/1? not sure about the last trunk. -1 definitely doesn't make sense. 0 also doesn't make sense. In your example, closing cost certainly skewed the distribution away from the ideal assignment. but it might be a little more generally applicable. -- 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: issues-unsubscr...@iceberg.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: issues-unsubscr...@iceberg.apache.org For additional commands, e-mail: issues-h...@iceberg.apache.org
