Use getAndSet in GroupedAggregation states

trinodb · Aug 22, 2024 · f02a840 · f02a840
1 parent 51507e6
commit f02a840
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 25 deletions.
diff --git a/.../src/main/java/io/trino/operator/aggregation/LongApproximateMostFrequentStateFactory.java b/.../src/main/java/io/trino/operator/aggregation/LongApproximateMostFrequentStateFactory.java
@@ -74,13 +74,11 @@ public ApproximateMostFrequentHistogram<Long> get()
         @Override
         public void set(ApproximateMostFrequentHistogram<Long> histogram)
         {
-            ApproximateMostFrequentHistogram<Long> previous = get();
+            ApproximateMostFrequentHistogram<Long> previous = histograms.getAndSet(getGroupId(), histogram);
+            size += histogram.estimatedInMemorySize();
             if (previous != null) {
                 size -= previous.estimatedInMemorySize();
             }
-
-            histograms.set(getGroupId(), histogram);
-            size += histogram.estimatedInMemorySize();
         }
 
         @Override

diff --git a/.../trino-main/src/main/java/io/trino/operator/aggregation/NumericHistogramStateFactory.java b/.../trino-main/src/main/java/io/trino/operator/aggregation/NumericHistogramStateFactory.java
@@ -58,13 +58,11 @@ public void set(NumericHistogram value)
         {
             requireNonNull(value, "value is null");
 
-            NumericHistogram previous = get();
+            NumericHistogram previous = histograms.getAndSet(getGroupId(), value);
+            size += value.estimatedInMemorySize();
             if (previous != null) {
                 size -= previous.estimatedInMemorySize();
             }
-
-            histograms.set(getGroupId(), value);
-            size += value.estimatedInMemorySize();
         }
 
         @Override

diff --git a/...rc/main/java/io/trino/operator/aggregation/StringApproximateMostFrequentStateFactory.java b/...rc/main/java/io/trino/operator/aggregation/StringApproximateMostFrequentStateFactory.java
@@ -75,13 +75,11 @@ public ApproximateMostFrequentHistogram<Slice> get()
         @Override
         public void set(ApproximateMostFrequentHistogram<Slice> histogram)
         {
-            ApproximateMostFrequentHistogram<Slice> previous = get();
+            ApproximateMostFrequentHistogram<Slice> previous = histograms.getAndSet(getGroupId(), histogram);
+            size += histogram.estimatedInMemorySize();
             if (previous != null) {
                 size -= previous.estimatedInMemorySize();
             }
-
-            histograms.set(getGroupId(), histogram);
-            this.size = histogram.estimatedInMemorySize();
         }
 
         @Override

diff --git a/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestStateFactory.java b/core/trino-main/src/main/java/io/trino/type/setdigest/SetDigestStateFactory.java
@@ -66,11 +66,11 @@ public SetDigest getDigest()
         @Override
         public void setDigest(SetDigest value)
         {
-            if (getDigest() != null) {
-                size -= getDigest().estimatedInMemorySize();
-            }
+            SetDigest previous = digests.getAndSet(groupId, value);
             size += value.estimatedInMemorySize();
-            digests.set(groupId, value);
+            if (previous != null) {
+                size -= previous.estimatedInMemorySize();
+            }
         }
 
         @Override

diff --git a/...-geospatial/src/main/java/io/trino/plugin/geospatial/SpatialPartitioningStateFactory.java b/...-geospatial/src/main/java/io/trino/plugin/geospatial/SpatialPartitioningStateFactory.java
@@ -89,10 +89,10 @@ public Rectangle getExtent()
         @Override
         public void setExtent(Rectangle envelope)
         {
-            if (envelopes.get(groupId) == null) {
+            Rectangle previousEnvelope = envelopes.getAndSet(groupId, envelope);
+            if (previousEnvelope == null) {
                 envelopeCount++;
             }
-            envelopes.set(groupId, envelope);
         }
 
         @Override
@@ -104,12 +104,11 @@ public List<Rectangle> getSamples()
         @Override
         public void setSamples(List<Rectangle> samples)
         {
-            List<Rectangle> currentSamples = this.samples.get(groupId);
-            if (currentSamples != null) {
-                samplesCount -= currentSamples.size();
-            }
+            List<Rectangle> previousSamples = this.samples.getAndSet(groupId, samples);
             samplesCount += samples.size();
-            this.samples.set(groupId, samples);
+            if (previousSamples != null) {
+                samplesCount -= previousSamples.size();
+            }
         }
 
         @Override

diff --git a/...geospatial/src/main/java/io/trino/plugin/geospatial/aggregation/GeometryStateFactory.java b/...geospatial/src/main/java/io/trino/plugin/geospatial/aggregation/GeometryStateFactory.java
@@ -54,10 +54,9 @@ public OGCGeometry getGeometry()
         @Override
         public void setGeometry(OGCGeometry geometry)
         {
-            OGCGeometry previousValue = this.geometries.get(groupId);
+            OGCGeometry previousValue = this.geometries.getAndSet(groupId, geometry);
             size -= getGeometryMemorySize(previousValue);
             size += getGeometryMemorySize(geometry);
-            this.geometries.set(groupId, geometry);
         }
 
         @Override