diff --git a/src/main/java/org/apache/datasketches/sampling/EbppsItemsSketch.java b/src/main/java/org/apache/datasketches/sampling/EbppsItemsSketch.java index 8b44aa086..7a1a87574 100644 --- 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 EbppsItemsSketch 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 EbppsItemsSketch heapify(final Memory srcMem, final List itemsList = Arrays.asList(rawItems); final ArrayList 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 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 index 454514fed..19faa8d38 100644 --- 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