fix(deletion_vector): read Java bitmap64 deletion vectors

crates/paimon/src/deletion_vector/core.rs

@@ -32,7 +32,15 @@ pub struct DeletionVector {
 /// Magic number for BitmapDeletionVector serialization format
 /// Same as Java: 1581511376
 const MAGIC_NUMBER: u32 = 1581511376;
-const MAGIC_NUMBER_SIZE_BYTES: usize = 4;
+/// Magic number for Java `Bitmap64DeletionVector` (v2). Written little-endian,
+/// unlike the v1 magic, so the two are told apart by byte order as well as value.
+const MAGIC_NUMBER_64: u32 = 1681511377;
+const MAGIC_NUMBER_SIZE_BYTES: u64 = 4;
+/// Size of the `bitmapLength` prefix and of the trailing CRC. Mirrors Java
+/// `Bitmap64DeletionVector.LENGTH_SIZE_BYTES` / `CRC_SIZE_BYTES`; a v2 entry
+/// counts both in `DeletionFile.length()` while a v1 entry counts neither.
+const LENGTH_SIZE_BYTES: u64 = 4;
+const CRC_SIZE_BYTES: u64 = 4;
 
 impl DeletionVector {
     /// Create a new empty DeletionVector
@@ -96,12 +104,10 @@ impl DeletionVector {
                 source: Some(Box::new(e)),
             })?;
 
-        let bitmap_length =
-            i32::try_from(MAGIC_NUMBER_SIZE_BYTES + bitmap_bytes.len()).map_err(|_| {
-                crate::Error::DataInvalid {
-                    message: "Deletion vector bitmap is too large to serialize".to_string(),
-                    source: None,
-                }
+        let bitmap_length = i32::try_from(MAGIC_NUMBER_SIZE_BYTES as usize + bitmap_bytes.len())
+            .map_err(|_| crate::Error::DataInvalid {
+                message: "Deletion vector bitmap is too large to serialize".to_string(),
+                source: None,
             })?;
 
         let mut payload = Vec::with_capacity(8 + bitmap_bytes.len() + 4);
@@ -121,13 +127,24 @@ impl DeletionVector {
         &self.bitmap
     }
 
-    /// Read a DeletionVector from bytes, similar to Java DeletionVector.read(DataInputStream, length)
+    /// Read a DeletionVector from bytes, mirroring Java `DeletionVector.read(DataInputStream, length)`,
+    /// which accepts two on-disk formats and dispatches on the magic number.
+    ///
+    /// `BitmapDeletionVector` (v1):
+    /// - `bitmapLength` (i32 big-endian): magic + bitmap data
+    /// - magic [`MAGIC_NUMBER`] (i32 big-endian)
+    /// - bitmap data (`bitmapLength - 4` bytes): serialized roaring32 bitmap
+    /// - CRC (i32): not verified here, matching Java's `dis.skipBytes(4)`
+    ///
+    /// `Bitmap64DeletionVector` (v2), whose payload is little-endian:
+    /// - `bitmapLength` (i32 **big**-endian): magic + bitmap data
+    /// - magic [`MAGIC_NUMBER_64`] (i32 **little**-endian)
+    /// - bitmap data: portable 64-bit roaring (`u64` bucket count, then `u32` key
+    ///   plus a roaring32 bitmap per bucket), possibly run-length encoded
+    /// - CRC (i32): not verified here either
     ///
-    /// Format (as read by DeletionVector.read):
-    /// - bitmapLength (4 bytes int): total size including magic
-    /// - magicNumber (4 bytes int): BitmapDeletionVector.MAGIC_NUMBER
-    /// - bitmap data (bitmapLength - 4 bytes): serialized RoaringBitmap
-    /// - CRC (4 bytes): checksum (skipped during read)
+    /// `expected_length` is `DeletionFile.length()`. The two formats count it
+    /// differently: v1 excludes the length prefix and the CRC, v2 includes both.
     pub fn read_from_bytes(bytes: &[u8], expected_length: Option<u64>) -> crate::Result<Self> {
         use bytes::Buf;
         if bytes.len() < 8 {
@@ -139,59 +156,117 @@ impl DeletionVector {
 
         let mut buf = bytes;
 
-        // Read bitmapLength (total size including magic)
-        let bitmap_length = buf.get_i32() as usize;
+        // Read bitmapLength (magic + bitmap data). Both formats store it
+        // big-endian. Reject a negative value here so the size arithmetic below
+        // cannot wrap: `as usize` would turn -1 into u64::MAX and make the
+        // "data incomplete" guard compute a tiny requirement and pass.
+        let bitmap_length =
+            u64::try_from(buf.get_i32()).map_err(|_| crate::Error::DataInvalid {
+                message: "Deletion vector bitmap length is negative".to_string(),
+                source: None,
+            })?;
 
-        // Read magic number
+        // Read magic number. v1 is big-endian, v2 little-endian.
         let magic_number = buf.get_i32() as u32;
-        if magic_number != MAGIC_NUMBER {
+        let is_bitmap64 = if magic_number == MAGIC_NUMBER {
+            false
+        } else if magic_number.swap_bytes() == MAGIC_NUMBER_64 {
+            true
+        } else {
             return Err(crate::Error::DataInvalid {
                 message: format!(
-                    "Invalid magic number: expected {MAGIC_NUMBER}, got {magic_number}"
+                    "Invalid magic number: {magic_number}, \
+                     v1 dv magic number: {MAGIC_NUMBER}, v2 magic number: {MAGIC_NUMBER_64}"
                 ),
                 source: None,
             });
-        }
+        };
 
-        // Verify length if provided
+        // Verify length if provided, using each format's own convention.
         if let Some(expected) = expected_length {
-            if bitmap_length as u64 != expected {
+            let expected_bitmap_length = if is_bitmap64 {
+                expected
+                    .checked_sub(LENGTH_SIZE_BYTES + CRC_SIZE_BYTES)
+                    .ok_or_else(|| crate::Error::DataInvalid {
+                        message: format!("Deletion vector length {expected} is too small"),
+                        source: None,
+                    })?
+            } else {
+                expected
+            };
+            if bitmap_length != expected_bitmap_length {
                 return Err(crate::Error::DataInvalid {
                     message: format!(
-                        "Size not match, actual size: {bitmap_length}, expected size: {expected}"
+                        "Size not match, actual size: {bitmap_length}, expected size: {expected_bitmap_length}"
                     ),
                     source: None,
                 });
             }
         }
 
-        // Read bitmap data (bitmapLength - 4 bytes, since magic is already included in bitmapLength)
-        let bitmap_data_size = bitmap_length - MAGIC_NUMBER_SIZE_BYTES;
-        // 4(bitmap_length) + 4(magic_number) + bitmap_data_size + 4(crc)
-        if bytes.len() < 8 + bitmap_data_size + 4 {
+        // Bitmap data follows the magic, which bitmapLength counts.
+        let bitmap_data_size = bitmap_length
+            .checked_sub(MAGIC_NUMBER_SIZE_BYTES)
+            .and_then(|size| usize::try_from(size).ok())
+            .ok_or_else(|| crate::Error::DataInvalid {
+                message: format!("Deletion vector bitmap length {bitmap_length} is too small"),
+                source: None,
+            })?;
+        // The CRC is not verified, so it need not be present: a v2 entry that
+        // ends the index file cannot be over-read (see `DeletionVectorFactory::read`).
+        let needed = 8 + bitmap_data_size;
+        if bytes.len() < needed {
             return Err(crate::Error::DataInvalid {
                 message: format!(
-                    "Deletion vector data incomplete: need {} bytes, got {}",
-                    8 + bitmap_data_size + 4,
+                    "Deletion vector data incomplete: need {needed} bytes, got {}",
                     bytes.len()
                 ),
                 source: None,
             });
         }
 
-        let bitmap_data = &bytes[8..8 + bitmap_data_size];
-
-        // Skip CRC (4 bytes) - Java code does: dis.skipBytes(4)
-        // We don't need to verify it here as it's skipped
+        let bitmap_data = &bytes[8..needed];
+        if is_bitmap64 {
+            Self::from_bitmap64_bytes(bitmap_data)
+        } else {
+            let bitmap = RoaringBitmap::deserialize_from(bitmap_data).map_err(|e| {
+                crate::Error::DataInvalid {
+                    message: format!("Failed to deserialize RoaringBitmap: {e}"),
+                    source: Some(Box::new(e)),
+                }
+            })?;
+            Ok(Self::from_bitmap(bitmap))
+        }
+    }
 
-        // Deserialize RoaringBitmap
-        let bitmap = RoaringBitmap::deserialize_from(bitmap_data).map_err(|e| {
+    /// Decode a `Bitmap64DeletionVector` payload into the roaring32 representation
+    /// this type stores.
+    ///
+    /// Row positions are offsets inside a single data file, and every existing API
+    /// here is already roaring32-bound -- [`Self::is_deleted`] documents that
+    /// positions above `u32::MAX` cannot be present, and `to_bitmap` hands a
+    /// `RoaringBitmap` to the writer. A position that does not fit is therefore
+    /// rejected rather than truncated, so a 64-bit vector can never silently lose
+    /// deletes; a data file with more than `u32::MAX` rows would be needed to
+    /// reach it.
+    fn from_bitmap64_bytes(bitmap_data: &[u8]) -> crate::Result<Self> {
+        let treemap = roaring::RoaringTreemap::deserialize_from(bitmap_data).map_err(|e| {
             crate::Error::DataInvalid {
-                message: format!("Failed to deserialize RoaringBitmap: {e}"),
+                message: format!("Failed to deserialize 64-bit RoaringBitmap: {e}"),
                 source: Some(Box::new(e)),
             }
         })?;
-
+        let mut bitmap = RoaringBitmap::new();
+        for position in treemap {
+            let position = u32::try_from(position).map_err(|_| crate::Error::DataInvalid {
+                message: format!(
+                    "Deletion vector position {position} exceeds u32::MAX, \
+                     which this reader cannot represent"
+                ),
+                source: None,
+            })?;
+            bitmap.insert(position);
+        }
         Ok(Self::from_bitmap(bitmap))
     }
 }
@@ -259,6 +334,118 @@ mod tests {
         assert_eq!(dv.bitmap(), &expected_bitmap, "bitmap should be [1, 2]");
     }
 
+    /// Build the on-disk bytes Java `Bitmap64DeletionVector.serializeTo` writes:
+    /// `i32 BE bitmapLength | i32 LE magic | portable 64-bit roaring | i32 BE crc`,
+    /// where the payload is `u64 LE bucket count` then `u32 LE key` plus a
+    /// roaring32 bitmap per bucket, densely for every key from 0 upwards.
+    fn java_bitmap64_bytes(buckets: &[&[u32]], run_length_encode: bool) -> Vec<u8> {
+        let mut payload = Vec::new();
+        payload.extend_from_slice(&(buckets.len() as u64).to_le_bytes());
+        for (key, values) in buckets.iter().enumerate() {
+            payload.extend_from_slice(&(key as u32).to_le_bytes());
+            let mut bitmap = RoaringBitmap::new();
+            for value in *values {
+                bitmap.insert(*value);
+            }
+            if run_length_encode {
+                // Java calls `roaringBitmap.runLengthEncode()` before serializing.
+                bitmap.optimize();
+            }
+            bitmap.serialize_into(&mut payload).unwrap();
+        }
+
+        let bitmap_length = 4 + payload.len();
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&(bitmap_length as i32).to_be_bytes());
+        bytes.extend_from_slice(&MAGIC_NUMBER_64.to_le_bytes());
+        bytes.extend_from_slice(&payload);
+        let mut crc = crc32fast::Hasher::new();
+        crc.update(&bytes[4..]);
+        bytes.extend_from_slice(&(crc.finalize() as i32).to_be_bytes());
+        bytes
+    }
+
+    /// `DeletionFile.length()` for a v2 entry is the whole on-disk run.
+    fn bitmap64_declared_length(bytes: &[u8]) -> u64 {
+        bytes.len() as u64
+    }
+
+    #[test]
+    fn test_read_bitmap64_deletion_vector() {
+        let bytes = java_bitmap64_bytes(&[&[1u32, 2u32, 9u32]], false);
+        let dv = DeletionVector::read_from_bytes(&bytes, Some(bitmap64_declared_length(&bytes)))
+            .expect("v2 deletion vector must be readable");
+        assert_eq!(dv.bitmap(), &RoaringBitmap::from_iter([1u32, 2, 9]));
+        assert!(dv.is_deleted(2) && !dv.is_deleted(3));
+        assert_eq!(dv.cardinality(), 3);
+    }
+
+    #[test]
+    fn test_read_bitmap64_run_length_encoded() {
+        // Java run-length encodes before writing, which switches the container
+        // type and the serialization cookie.
+        let run: Vec<u32> = (100u32..400).collect();
+        let bytes = java_bitmap64_bytes(&[&run], true);
+        let dv = DeletionVector::read_from_bytes(&bytes, Some(bitmap64_declared_length(&bytes)))
+            .expect("run-length encoded v2 deletion vector must be readable");
+        assert_eq!(dv.cardinality(), run.len() as u64);
+        assert!(dv.is_deleted(100) && dv.is_deleted(399) && !dv.is_deleted(400));
+    }
+
+    #[test]
+    fn test_read_bitmap64_with_empty_leading_bucket() {
+        // Java writes every key from 0 densely, so bucket 0 is present and empty
+        // when no position below 2^32 is deleted. Rust's own writer omits it.
+        let bytes = java_bitmap64_bytes(&[&[], &[7u32]], false);
+        let err = DeletionVector::read_from_bytes(&bytes, Some(bitmap64_declared_length(&bytes)))
+            .expect_err("a position above u32::MAX must be rejected, not truncated");
+        assert!(
+            err.to_string().contains("exceeds u32::MAX"),
+            "unexpected error: {err}"
+        );
+    }
+
+    #[test]
+    fn test_read_bitmap64_uses_its_own_length_convention() {
+        let bytes = java_bitmap64_bytes(&[&[5u32]], false);
+        let declared = bitmap64_declared_length(&bytes);
+        // v1's convention (length == bitmapLength) must not be applied to v2.
+        let err = DeletionVector::read_from_bytes(&bytes, Some(declared - 8))
+            .expect_err("v2 length is the whole on-disk run, prefix and CRC included");
+        assert!(err.to_string().contains("Size not match"), "got: {err}");
+        assert!(DeletionVector::read_from_bytes(&bytes, Some(declared)).is_ok());
+    }
+
+    #[test]
+    fn test_read_rejects_unknown_magic_naming_both_formats() {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&8i32.to_be_bytes());
+        bytes.extend_from_slice(&0xDEAD_BEEFu32.to_be_bytes());
+        bytes.extend_from_slice(&[0u8; 8]);
+        let err = DeletionVector::read_from_bytes(&bytes, None).expect_err("magic must be checked");
+        let message = err.to_string();
+        assert!(
+            message.contains(&MAGIC_NUMBER.to_string()),
+            "got: {message}"
+        );
+        assert!(
+            message.contains(&MAGIC_NUMBER_64.to_string()),
+            "got: {message}"
+        );
+    }
+
+    #[test]
+    fn test_read_rejects_negative_bitmap_length() {
+        // Taken straight from the file, a negative length used to become a huge
+        // usize and wrap the "data incomplete" guard into passing.
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&(-1i32).to_be_bytes());
+        bytes.extend_from_slice(&MAGIC_NUMBER.to_be_bytes());
+        bytes.extend_from_slice(&[0u8; 8]);
+        let err = DeletionVector::read_from_bytes(&bytes, None).expect_err("must not panic");
+        assert!(err.to_string().contains("negative"), "got: {err}");
+    }
+
     #[test]
     fn test_is_deleted_reports_membership_and_guards_u32_overflow() {
         let mut bitmap = RoaringBitmap::new();