Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
/*
* Licensed 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 io.trino.execution.scheduler.faulttolerant;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ListMultimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Sets;
import io.trino.metadata.Split;
import io.trino.sql.planner.plan.PlanNodeId;
import jakarta.annotation.Nullable;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Verify.verify;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.Iterables.getOnlyElement;
import static io.airlift.slice.SizeOf.INTEGER_INSTANCE_SIZE;
import static io.airlift.slice.SizeOf.estimatedSizeOf;
import static io.airlift.slice.SizeOf.instanceSize;
import static java.util.Objects.requireNonNull;
public final class SplitsMapping
{
private static final int INSTANCE_SIZE = instanceSize(SplitsMapping.class);
public static final SplitsMapping EMPTY = SplitsMapping.builder().build();
// not using Multimap to avoid extensive data structure copying when building updated SplitsMapping
private final Map>> splits; // plan-node -> hash-partition -> Split
private SplitsMapping(ImmutableMap>> splits)
{
// Builder implementations ensure that external map as well as Maps/Lists used in values
// are immutable.
this.splits = splits;
}
public Set getPlanNodeIds()
{
return splits.keySet();
}
public ListMultimap getSplitsFlat()
{
ImmutableListMultimap.Builder splitsFlat = ImmutableListMultimap.builder();
for (Map.Entry>> entry : splits.entrySet()) {
// TODO can we do less copying?
splitsFlat.putAll(entry.getKey(), entry.getValue().values().stream().flatMap(Collection::stream).collect(toImmutableList()));
}
return splitsFlat.build();
}
public List getSplitsFlat(PlanNodeId planNodeId)
{
Map> splits = this.splits.get(planNodeId);
if (splits == null) {
return ImmutableList.of();
}
verify(!splits.isEmpty(), "expected not empty splits list %s", splits);
if (splits.size() == 1) {
return getOnlyElement(splits.values());
}
// TODO improve to not copy here; return view instead
ImmutableList.Builder result = ImmutableList.builder();
for (List partitionSplits : splits.values()) {
result.addAll(partitionSplits);
}
return result.build();
}
@VisibleForTesting
ListMultimap getSplits(PlanNodeId planNodeId)
{
Map> splits = this.splits.get(planNodeId);
if (splits == null) {
return ImmutableListMultimap.of();
}
verify(!splits.isEmpty(), "expected not empty splits list %s", splits);
ImmutableListMultimap.Builder result = ImmutableListMultimap.builder();
for (Map.Entry> entry : splits.entrySet()) {
result.putAll(entry.getKey(), entry.getValue());
}
return result.build();
}
public long getRetainedSizeInBytes()
{
return INSTANCE_SIZE +
estimatedSizeOf(
splits,
PlanNodeId::getRetainedSizeInBytes,
planNodeSplits -> estimatedSizeOf(
planNodeSplits,
partitionId -> INTEGER_INSTANCE_SIZE,
splitList -> estimatedSizeOf(splitList, Split::getRetainedSizeInBytes)));
}
@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
SplitsMapping that = (SplitsMapping) o;
return Objects.equals(splits, that.splits);
}
@Override
public int hashCode()
{
return Objects.hash(splits);
}
@Override
public String toString()
{
return toStringHelper(this)
.add("splits", splits)
.toString();
}
public static Builder builder()
{
return new NewBuilder();
}
public static Builder builder(SplitsMapping mapping)
{
return new UpdatingBuilder(mapping);
}
public long size()
{
return splits.values().stream()
.flatMap(sourcePartitionToSplits -> sourcePartitionToSplits.values().stream())
.mapToLong(List::size)
.sum();
}
public abstract static class Builder
{
private Builder() {} // close for extension
public Builder addSplit(PlanNodeId planNodeId, int partitionId, Split split)
{
return addSplits(planNodeId, partitionId, ImmutableList.of(split));
}
public Builder addSplits(PlanNodeId planNodeId, ListMultimap splits)
{
Multimaps.asMap(splits).forEach((partitionId, partitionSplits) -> addSplits(planNodeId, partitionId, partitionSplits));
return this;
}
public Builder addMapping(SplitsMapping updatingMapping)
{
for (Map.Entry>> entry : updatingMapping.splits.entrySet()) {
PlanNodeId planNodeId = entry.getKey();
entry.getValue().forEach((partitionId, partitionSplits) -> addSplits(planNodeId, partitionId, partitionSplits));
}
return this;
}
public abstract Builder addSplits(PlanNodeId planNodeId, int partitionId, List splits);
public abstract SplitsMapping build();
}
private static class UpdatingBuilder
extends Builder
{
private final SplitsMapping originalMapping;
private final Map>> updates = new HashMap<>();
public UpdatingBuilder(SplitsMapping originalMapping)
{
this.originalMapping = requireNonNull(originalMapping, "sourceMapping is null");
}
@Override
public Builder addSplits(PlanNodeId planNodeId, int partitionId, List splits)
{
if (splits.isEmpty()) {
// ensure we do not have empty lists in result splits map.
return this;
}
updates.computeIfAbsent(planNodeId, ignored -> new HashMap<>())
.computeIfAbsent(partitionId, key -> ImmutableList.builder())
.addAll(splits);
return this;
}
@Override
public SplitsMapping build()
{
ImmutableMap.Builder>> result = ImmutableMap.builder();
for (PlanNodeId planNodeId : Sets.union(originalMapping.splits.keySet(), updates.keySet())) {
Map> planNodeOriginalMapping = originalMapping.splits.getOrDefault(planNodeId, ImmutableMap.of());
Map> planNodeUpdates = updates.getOrDefault(planNodeId, ImmutableMap.of());
if (planNodeUpdates.isEmpty()) {
// just use original splits for planNodeId
result.put(planNodeId, planNodeOriginalMapping);
continue;
}
// create new mapping for planNodeId reusing as much of source as possible
ImmutableMap.Builder> targetSplitsMapBuilder = ImmutableMap.builder();
for (Integer sourcePartitionId : Sets.union(planNodeOriginalMapping.keySet(), planNodeUpdates.keySet())) {
@Nullable List originalSplits = planNodeOriginalMapping.get(sourcePartitionId);
@Nullable ImmutableList.Builder splitUpdates = planNodeUpdates.get(sourcePartitionId);
targetSplitsMapBuilder.put(sourcePartitionId, mergeIfPresent(originalSplits, splitUpdates));
}
result.put(planNodeId, targetSplitsMapBuilder.buildOrThrow());
}
return new SplitsMapping(result.buildOrThrow());
}
private static List mergeIfPresent(@Nullable List list, @Nullable ImmutableList.Builder additionalElements)
{
if (additionalElements == null) {
// reuse source immutable split list
return requireNonNull(list, "list is null");
}
if (list == null) {
return additionalElements.build();
}
return ImmutableList.builder()
.addAll(list)
.addAll(additionalElements.build())
.build();
}
}
private static class NewBuilder
extends Builder
{
private final Map>> splitsBuilder = new HashMap<>();
@Override
public Builder addSplits(PlanNodeId planNodeId, int partitionId, List splits)
{
if (splits.isEmpty()) {
// ensure we do not have empty lists in result splits map.
return this;
}
splitsBuilder.computeIfAbsent(planNodeId, ignored -> new HashMap<>())
.computeIfAbsent(partitionId, ignored -> ImmutableList.builder())
.addAll(splits);
return this;
}
@Override
public SplitsMapping build()
{
return new SplitsMapping(splitsBuilder.entrySet().stream()
.collect(toImmutableMap(
Map.Entry::getKey,
planNodeMapping -> planNodeMapping.getValue().entrySet().stream()
.collect(toImmutableMap(
Map.Entry::getKey,
sourcePartitionMapping -> sourcePartitionMapping.getValue().build())))));
}
}
}
