xiangfu0 commented on code in PR #20248:
URL: https://github.com/apache/pinot/pull/20248#discussion_r4559555462


##########
pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/IntermediateResultByteBufferSerDeTest.java:
##########
@@ -487,12 +628,54 @@ public void testParentExprMinMax()
     
assertEquals(function.serializeIntermediateResult(deserialized).getBytes(), 
expectedBytes);
   }
 
+  @DataProvider
+  public static Object[][] offHeapSets() {
+    OffHeap32BitSet set32 = new OffHeap32BitSet(4);
+    OffHeap64BitSet set64 = new OffHeap64BitSet(4);
+    OffHeap128BitSet set128 = new OffHeap128BitSet(4);
+    // Includes zero, which the sets hold outside their hash tables
+    for (int value = 0; value < 3; value++) {
+      set32.add(value);
+      set64.add(value);
+      set128.add(value, -value);
+    }
+    return new Object[][]{
+        {set32, 0},
+        {set64, 1},
+        {set128, 2}
+    };
+  }
+
+  /// The `DISTINCTCOUNTOFFHEAP` function, checked on one serialization of 
each set, since serializing a set frees its
+  /// memory: its wire type value, the payload size, and a payload that 
deserializes to a set with the same final
+  /// result after its source buffer is overwritten.
+  @Test(dataProvider = "offHeapSets")
+  public void testDistinctCountOffHeap(BaseOffHeapSet set, int expectedType) {
+    DistinctCountOffHeapAggregationFunction function =
+        new DistinctCountOffHeapAggregationFunction(List.of(COLUMN), false);
+    int expectedFinalResult = set.size();
+    SerializedResult serialized = function.serializeIntermediateResult(set);
+    assertEquals(serialized.getType(), expectedType);
+    byte[] bytes = serialized.getBytes();
+    assertEquals(serialized.getSize(), bytes.length);

Review Comment:
   This always passes: `SerializedBytes.getSize()` is just `_bytes.length`, so 
it wouldn't catch #20203, where the array itself was oversized. I checked by 
reverting the parentheses fix in `OffHeap32BitSet.serialize()`, and 
`testDistinctCountOffHeap[0]` still passes, because `deserialize` ignores the 
trailing bytes.
   
   If you want this to guard the size, you could add an expected payload length 
to `offHeapSets()` (`4 + 3 * 4`, `4 + 3 * 8`, `4 + 3 * 16`) and assert 
`bytes.length` against it. Otherwise, consider dropping the #20203 claim from 
the PR description. Non-blocking, since `OffHeapSetTest` already covers it.



-- 
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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to