Merge pull request apache#503 from apache/ebpps

use partial item flag in serialization
anirudhacharya · Feb 6, 2024 · 5e7174f · 5e7174f
2 parents a8231cb + 2d5769a
commit 5e7174f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/main/java/org/apache/datasketches/sampling/EbppsItemsSketch.java b/src/main/java/org/apache/datasketches/sampling/EbppsItemsSketch.java
@@ -21,6 +21,7 @@
 
 import static org.apache.datasketches.sampling.PreambleUtil.EBPPS_SER_VER;
 import static org.apache.datasketches.sampling.PreambleUtil.EMPTY_FLAG_MASK;
+import static org.apache.datasketches.sampling.PreambleUtil.HAS_PARTIAL_ITEM_MASK;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -121,6 +122,7 @@ public static <T> EbppsItemsSketch<T> heapify(final Memory srcMem,
     final int familyId = PreambleUtil.extractFamilyID(srcMem);
     final int flags = PreambleUtil.extractFlags(srcMem);
     final boolean isEmpty = (flags & EMPTY_FLAG_MASK) != 0;
+    final boolean hasPartialItem = (flags & HAS_PARTIAL_ITEM_MASK) != 0;
 
     // Check values
     if (isEmpty) {
@@ -189,8 +191,11 @@ public static <T> EbppsItemsSketch<T> heapify(final Memory srcMem,
     final List<T> itemsList = Arrays.asList(rawItems);
     final ArrayList<T> data;
     final T partialItem;
-    if (numFullItems < numTotalItems) {
-      data = new ArrayList<>(itemsList.subList(0, numFullItems));
+    if (hasPartialItem) {
+      if (numFullItems >= numTotalItems)
+        throw new SketchesArgumentException("Possible Corruption: Expected partial item but none found");
+
+        data = new ArrayList<>(itemsList.subList(0, numFullItems));
       partialItem = itemsList.get(numFullItems); // 0-based, so last item
     } else {
       data = new ArrayList<>(itemsList);
@@ -487,7 +492,7 @@ public byte[] toByteArray(final ArrayOfItemsSerDe<? super T> serDe, final Class<
     if (empty) {
       PreambleUtil.insertFlags(mem, EMPTY_FLAG_MASK);        // Byte 3
     } else {
-      PreambleUtil.insertFlags(mem, 0);
+      PreambleUtil.insertFlags(mem, sample_.hasPartialItem() ? HAS_PARTIAL_ITEM_MASK : 0);
     }
     PreambleUtil.insertK(mem, k_);                           // Bytes 4-7
 

diff --git a/src/main/java/org/apache/datasketches/sampling/PreambleUtil.java b/src/main/java/org/apache/datasketches/sampling/PreambleUtil.java
@@ -209,6 +209,7 @@ private PreambleUtil() {}
   //static final int BIG_ENDIAN_FLAG_MASK = 1;
   //static final int READ_ONLY_FLAG_MASK  = 2;
   static final int EMPTY_FLAG_MASK      = 4;
+  static final int HAS_PARTIAL_ITEM_MASK = 8; // EBPPS only
   static final int GADGET_FLAG_MASK     = 128;
 
   //Other constants