You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by cv...@apache.org on 2025/12/05 20:27:22 UTC
(beam) branch revert-36897-optimize_weigh created (now 63177cb5db6)
This is an automated email from the ASF dual-hosted git repository.
cvandermerwe pushed a change to branch revert-36897-optimize_weigh
in repository https://gitbox.apache.org/repos/asf/beam.git
at 63177cb5db6 Revert "Reduce weighing overhead for caching blocks (#36897)"
This branch includes the following new commits:
new 63177cb5db6 Revert "Reduce weighing overhead for caching blocks (#36897)"
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(beam) 01/01: Revert "Reduce weighing overhead for caching blocks (#36897)"
Posted by cv...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
cvandermerwe pushed a commit to branch revert-36897-optimize_weigh
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 63177cb5db6623bbfe9ba7ede1041f58dd4ddb93
Author: claudevdm <33...@users.noreply.github.com>
AuthorDate: Fri Dec 5 15:27:11 2025 -0500
Revert "Reduce weighing overhead for caching blocks (#36897)"
This reverts commit 84684aef490e03163295050e3b6e51aa6c6a2ee6.
---
.../org/apache/beam/sdk/fn/data/WeightedList.java | 11 ++-
.../fn/harness/state/StateFetchingIterators.java | 83 ++++++++--------------
2 files changed, 39 insertions(+), 55 deletions(-)
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/WeightedList.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/WeightedList.java
index 5eb317fc287..ad5e131cb2d 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/WeightedList.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/WeightedList.java
@@ -20,7 +20,6 @@ package org.apache.beam.sdk.fn.data;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.sdk.util.Weighted;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.LongMath;
/** Facade for a {@link List<T>} that keeps track of weight, for cache limit reasons. */
public class WeightedList<T> implements Weighted {
@@ -72,6 +71,14 @@ public class WeightedList<T> implements Weighted {
}
public void accumulateWeight(long weight) {
- this.weight.accumulateAndGet(weight, LongMath::saturatedAdd);
+ this.weight.accumulateAndGet(
+ weight,
+ (first, second) -> {
+ try {
+ return Math.addExact(first, second);
+ } catch (ArithmeticException e) {
+ return Long.MAX_VALUE;
+ }
+ });
}
}
diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
index 339ddad4061..1e06c98f2e3 100644
--- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
+++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/StateFetchingIterators.java
@@ -49,8 +49,6 @@ import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.math.LongMath;
/**
* Adapters which convert a logical series of chunks using continuation tokens over the Beam Fn
@@ -251,11 +249,15 @@ public class StateFetchingIterators {
@Override
public long getWeight() {
- long sum = 8 + blocks.size() * 8L;
- for (Block<T> block : blocks) {
- sum = LongMath.saturatedAdd(sum, block.getWeight());
+ try {
+ long sum = 8 + blocks.size() * 8L;
+ for (Block<T> block : blocks) {
+ sum = Math.addExact(sum, block.getWeight());
+ }
+ return sum;
+ } catch (ArithmeticException e) {
+ return Long.MAX_VALUE;
}
- return sum;
}
BlocksPrefix(List<Block<T>> blocks) {
@@ -280,7 +282,8 @@ public class StateFetchingIterators {
@AutoValue
abstract static class Block<T> implements Weighted {
- private static final Block<Void> EMPTY = fromValues(ImmutableList.of(), 0, null);
+ private static final Block<Void> EMPTY =
+ fromValues(WeightedList.of(Collections.emptyList(), 0), null);
@SuppressWarnings("unchecked") // Based upon as Collections.emptyList()
public static <T> Block<T> emptyBlock() {
@@ -296,37 +299,21 @@ public class StateFetchingIterators {
}
public static <T> Block<T> fromValues(List<T> values, @Nullable ByteString nextToken) {
- if (values.isEmpty() && nextToken == null) {
- return emptyBlock();
- }
- ImmutableList<T> immutableValues = ImmutableList.copyOf(values);
- long listWeight = immutableValues.size() * Caches.REFERENCE_SIZE;
- for (T value : immutableValues) {
- listWeight = LongMath.saturatedAdd(listWeight, Caches.weigh(value));
- }
- return fromValues(immutableValues, listWeight, nextToken);
+ return fromValues(WeightedList.of(values, Caches.weigh(values)), nextToken);
}
public static <T> Block<T> fromValues(
WeightedList<T> values, @Nullable ByteString nextToken) {
- if (values.isEmpty() && nextToken == null) {
- return emptyBlock();
- }
- return fromValues(ImmutableList.copyOf(values.getBacking()), values.getWeight(), nextToken);
- }
-
- private static <T> Block<T> fromValues(
- ImmutableList<T> values, long listWeight, @Nullable ByteString nextToken) {
- long weight = LongMath.saturatedAdd(listWeight, 24);
+ long weight = values.getWeight() + 24;
if (nextToken != null) {
if (nextToken.isEmpty()) {
nextToken = ByteString.EMPTY;
} else {
- weight = LongMath.saturatedAdd(weight, Caches.weigh(nextToken));
+ weight += Caches.weigh(nextToken);
}
}
return new AutoValue_StateFetchingIterators_CachingStateIterable_Block<>(
- values, nextToken, weight);
+ values.getBacking(), nextToken, weight);
}
abstract List<T> getValues();
@@ -385,12 +372,10 @@ public class StateFetchingIterators {
totalSize += tBlock.getValues().size();
}
- ImmutableList.Builder<T> allValues = ImmutableList.builderWithExpectedSize(totalSize);
- long weight = 0;
- List<T> blockValuesToKeep = new ArrayList<>();
+ WeightedList<T> allValues = WeightedList.of(new ArrayList<>(totalSize), 0L);
for (Block<T> block : blocks) {
- blockValuesToKeep.clear();
boolean valueRemovedFromBlock = false;
+ List<T> blockValuesToKeep = new ArrayList<>();
for (T value : block.getValues()) {
if (!toRemoveStructuralValues.contains(valueCoder.structuralValue(value))) {
blockValuesToKeep.add(value);
@@ -402,19 +387,13 @@ public class StateFetchingIterators {
// If any value was removed from this block, need to estimate the weight again.
// Otherwise, just reuse the block's weight.
if (valueRemovedFromBlock) {
- allValues.addAll(blockValuesToKeep);
- for (T value : blockValuesToKeep) {
- weight = LongMath.saturatedAdd(weight, Caches.weigh(value));
- }
+ allValues.addAll(blockValuesToKeep, Caches.weigh(block.getValues()));
} else {
- allValues.addAll(block.getValues());
- weight = LongMath.saturatedAdd(weight, block.getWeight());
+ allValues.addAll(block.getValues(), block.getWeight());
}
}
- cache.put(
- IterableCacheKey.INSTANCE,
- new MutatedBlocks<>(Block.fromValues(allValues.build(), weight, null)));
+ cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues)));
}
/**
@@ -505,24 +484,21 @@ public class StateFetchingIterators {
for (Block<T> block : blocks) {
totalSize += block.getValues().size();
}
- ImmutableList.Builder<T> allValues = ImmutableList.builderWithExpectedSize(totalSize);
- long weight = 0;
+ WeightedList<T> allValues = WeightedList.of(new ArrayList<>(totalSize), 0L);
for (Block<T> block : blocks) {
- allValues.addAll(block.getValues());
- weight = LongMath.saturatedAdd(weight, block.getWeight());
+ allValues.addAll(block.getValues(), block.getWeight());
}
if (newWeight < 0) {
- newWeight = 0;
- for (T value : newValues) {
- newWeight = LongMath.saturatedAdd(newWeight, Caches.weigh(value));
+ if (newValues.size() == 1) {
+ // Optimize weighing of the common value state as single single-element bag state.
+ newWeight = Caches.weigh(newValues.get(0));
+ } else {
+ newWeight = Caches.weigh(newValues);
}
}
- allValues.addAll(newValues);
- weight = LongMath.saturatedAdd(weight, newWeight);
+ allValues.addAll(newValues, newWeight);
- cache.put(
- IterableCacheKey.INSTANCE,
- new MutatedBlocks<>(Block.fromValues(allValues.build(), weight, null)));
+ cache.put(IterableCacheKey.INSTANCE, new MutatedBlocks<>(Block.mutatedBlock(allValues)));
}
class CachingStateIterator implements PrefetchableIterator<T> {
@@ -604,7 +580,8 @@ public class StateFetchingIterators {
return false;
}
// Release the block while we are loading the next one.
- currentBlock = Block.emptyBlock();
+ currentBlock =
+ Block.fromValues(WeightedList.of(Collections.emptyList(), 0L), ByteString.EMPTY);
@Nullable Blocks<T> existing = cache.peek(IterableCacheKey.INSTANCE);
boolean isFirstBlock = ByteString.EMPTY.equals(nextToken);