Resolve a bug where datasketches would not downsample sketches suffic…

…iently (#16119) * Fix sketch memory issue * Rename function * Add unit test * Revert downsampling change
apache · May 14, 2024 · 18a4722 · 18a4722
1 parent b8dd747
commit 18a4722
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/...query/src/main/java/org/apache/druid/msq/statistics/ClusterByStatisticsCollectorImpl.java b/...query/src/main/java/org/apache/druid/msq/statistics/ClusterByStatisticsCollectorImpl.java
@@ -212,6 +212,12 @@ public long estimatedTotalWeight()
     return count;
   }
 
+  @VisibleForTesting
+  long getTotalRetainedBytes()
+  {
+    return totalRetainedBytes;
+  }
+
   @Override
   public boolean hasMultipleValues(final int keyPosition)
   {
@@ -414,7 +420,7 @@ private BucketHolder getOrCreateBucketHolder(final RowKey bucketKey)
   void downSample()
   {
     long newTotalRetainedBytes = totalRetainedBytes;
-    final long targetTotalRetainedBytes = totalRetainedBytes / 2;
+    final long targetTotalRetainedBytes = Math.min(totalRetainedBytes / 2, maxRetainedBytes);
 
     final List<Pair<Long, BucketHolder>> sortedHolders = new ArrayList<>(buckets.size());
     final RowKeyReader trimmedRowReader = keyReader.trimmedKeyReader(clusterBy.getBucketByCount());

diff --git a/...y/src/test/java/org/apache/druid/msq/statistics/ClusterByStatisticsCollectorImplTest.java b/...y/src/test/java/org/apache/druid/msq/statistics/ClusterByStatisticsCollectorImplTest.java
@@ -451,6 +451,26 @@ public void test_clusterByXYbucketByX_maxX_lowCardinalityY_withAggregation()
     );
   }
 
+  @Test
+  public void testShouldDownsampleSingleBucket()
+  {
+    ClusterByStatisticsCollectorImpl clusterByStatisticsCollector =
+        (ClusterByStatisticsCollectorImpl) ClusterByStatisticsCollectorImpl.create(
+            CLUSTER_BY_XYZ_BUCKET_BY_X,
+            SIGNATURE,
+            35000,
+            500,
+            false,
+            false
+        );
+
+    clusterByStatisticsCollector.add(createKey(CLUSTER_BY_XYZ_BUCKET_BY_X, 2, 1, "value1"), 1);
+    clusterByStatisticsCollector.add(createKey(CLUSTER_BY_XYZ_BUCKET_BY_X, 2, 3, "value2"), 1);
+    clusterByStatisticsCollector.add(createKey(CLUSTER_BY_XYZ_BUCKET_BY_X, 1, 1, "Extremely long key string for unit test; Extremely long key string for unit test;"), 500);
+
+    Assert.assertTrue(clusterByStatisticsCollector.getTotalRetainedBytes() <= 35000);
+  }
+
   @Test
   public void testBucketDownsampledToSingleKeyFinishesCorrectly()
   {