hudi-core 0.6.0

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
//! HFile block parsing and handling.

use crate::hfile::block_type::{HFileBlockType, MAGIC_LENGTH};
use crate::hfile::compression::CompressionCodec;
use crate::hfile::error::{HFileError, Result};
use crate::hfile::key::{KEY_VALUE_HEADER_SIZE, Key, KeyValue, key_content};

/// Size constants for block header
const SIZEOF_INT32: usize = 4;
const SIZEOF_INT64: usize = 8;
const SIZEOF_BYTE: usize = 1;

/// Block header size without checksum info (HFile v2)
const HEADER_SIZE_NO_CHECKSUM: usize = MAGIC_LENGTH + 2 * SIZEOF_INT32 + SIZEOF_INT64;

/// Block header size with checksum (HFile v3)
/// Header fields:
/// - 8 bytes: magic
/// - 4 bytes: on-disk size without header
/// - 4 bytes: uncompressed size without header
/// - 8 bytes: previous block offset
/// - 1 byte: checksum type
/// - 4 bytes: bytes per checksum
/// - 4 bytes: on-disk data size with header
pub const BLOCK_HEADER_SIZE: usize = HEADER_SIZE_NO_CHECKSUM + SIZEOF_BYTE + 2 * SIZEOF_INT32;

/// Parsed block header information.
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct BlockHeader {
    pub block_type: HFileBlockType,
    pub on_disk_size_without_header: usize,
    pub uncompressed_size_without_header: usize,
    pub prev_block_offset: i64,
    pub checksum_type: u8,
    pub bytes_per_checksum: usize,
    pub on_disk_data_size_with_header: usize,
}

impl BlockHeader {
    /// Parse a block header from bytes.
    pub fn parse(bytes: &[u8]) -> Result<Self> {
        if bytes.len() < BLOCK_HEADER_SIZE {
            return Err(HFileError::InvalidFormat(format!(
                "Buffer too small for block header: {} bytes, need {}",
                bytes.len(),
                BLOCK_HEADER_SIZE
            )));
        }

        let block_type = HFileBlockType::from_magic(bytes)?;

        let on_disk_size_without_header =
            i32::from_be_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize;

        let uncompressed_size_without_header =
            i32::from_be_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]) as usize;

        let prev_block_offset = i64::from_be_bytes([
            bytes[16], bytes[17], bytes[18], bytes[19], bytes[20], bytes[21], bytes[22], bytes[23],
        ]);

        let checksum_type = bytes[24];

        let bytes_per_checksum =
            i32::from_be_bytes([bytes[25], bytes[26], bytes[27], bytes[28]]) as usize;

        let on_disk_data_size_with_header =
            i32::from_be_bytes([bytes[29], bytes[30], bytes[31], bytes[32]]) as usize;

        Ok(Self {
            block_type,
            on_disk_size_without_header,
            uncompressed_size_without_header,
            prev_block_offset,
            checksum_type,
            bytes_per_checksum,
            on_disk_data_size_with_header,
        })
    }

    /// Returns the total on-disk size including header.
    pub fn on_disk_size_with_header(&self) -> usize {
        BLOCK_HEADER_SIZE + self.on_disk_size_without_header
    }
}

/// An HFile block with parsed content.
#[derive(Debug)]
pub struct HFileBlock {
    pub header: BlockHeader,
    /// Uncompressed block data (after header, before checksum)
    pub data: Vec<u8>,
}

impl HFileBlock {
    /// Parse a block from bytes, decompressing if necessary.
    pub fn parse(bytes: &[u8], codec: CompressionCodec) -> Result<Self> {
        let header = BlockHeader::parse(bytes)?;

        // Extract the compressed data (between header and checksum)
        let data_start = BLOCK_HEADER_SIZE;
        let data_end = header.on_disk_data_size_with_header;

        // For uncompressed blocks, on_disk_size == uncompressed_size
        // For compressed blocks, we need to decompress
        let data = if codec == CompressionCodec::None {
            bytes[data_start..data_start + header.uncompressed_size_without_header].to_vec()
        } else {
            let compressed_data = &bytes[data_start..data_end];
            codec.decompress(compressed_data, header.uncompressed_size_without_header)?
        };

        Ok(Self { header, data })
    }

    /// Returns the block type.
    pub fn block_type(&self) -> HFileBlockType {
        self.header.block_type
    }
}

/// A block index entry pointing to a data or meta block.
#[derive(Debug, Clone)]
pub struct BlockIndexEntry {
    /// First key in the block (may be shortened/fake for optimization)
    pub first_key: Key,
    /// First key of the next block (if present)
    pub next_block_first_key: Option<Key>,
    /// Offset of the block in the file
    pub offset: u64,
    /// On-disk size of the block
    pub size: u32,
}

impl BlockIndexEntry {
    /// Create a new block index entry.
    pub fn new(first_key: Key, next_block_first_key: Option<Key>, offset: u64, size: u32) -> Self {
        Self {
            first_key,
            next_block_first_key,
            offset,
            size,
        }
    }
}

impl PartialEq for BlockIndexEntry {
    fn eq(&self, other: &Self) -> bool {
        self.first_key == other.first_key
    }
}

impl Eq for BlockIndexEntry {}

impl PartialOrd for BlockIndexEntry {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for BlockIndexEntry {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.first_key.cmp(&other.first_key)
    }
}

/// Data block containing key-value pairs.
pub struct DataBlock {
    /// Uncompressed block data
    data: Vec<u8>,
    /// End offset of content (excluding checksum)
    content_end: usize,
}

impl DataBlock {
    /// Create a data block from an HFileBlock.
    pub fn from_block(block: HFileBlock) -> Self {
        let content_end = block.data.len();
        Self {
            data: block.data,
            content_end,
        }
    }

    /// Read a key-value at the given offset within the data block.
    pub fn read_key_value(&self, offset: usize) -> KeyValue {
        KeyValue::parse(&self.data, offset)
    }

    /// Check if the offset is within valid content range.
    pub fn is_valid_offset(&self, offset: usize) -> bool {
        offset < self.content_end
    }

    /// The record key content and value of every entry, as slices of the
    /// block.
    ///
    /// [`Self::iter`] copies each entry out of the block; a caller that only
    /// wants to find a few keys among thousands of entries can look at the
    /// keys in place and copy the entries it keeps.
    pub fn entries(&self) -> DataBlockEntries<'_> {
        DataBlockEntries {
            block: self,
            offset: 0,
        }
    }

    /// Iterate over all key-value pairs in the block.
    #[allow(dead_code)]
    pub fn iter(&self) -> DataBlockIterator<'_> {
        DataBlockIterator {
            block: self,
            offset: 0,
        }
    }

    /// Returns the raw data.
    #[allow(dead_code)]
    pub fn data(&self) -> &[u8] {
        &self.data
    }
}

/// Iterator over key-value pairs in a data block.
pub struct DataBlockIterator<'a> {
    block: &'a DataBlock,
    offset: usize,
}

impl<'a> Iterator for DataBlockIterator<'a> {
    type Item = KeyValue;

    fn next(&mut self) -> Option<Self::Item> {
        if !self.block.is_valid_offset(self.offset) {
            return None;
        }

        // Need at least 8 bytes for key/value lengths
        if self.offset + KEY_VALUE_HEADER_SIZE > self.block.content_end {
            return None;
        }

        let kv = self.block.read_key_value(self.offset);
        self.offset += kv.record_size();
        Some(kv)
    }
}

/// Borrowing walk over a data block's entries: see [`DataBlock::entries`].
pub struct DataBlockEntries<'a> {
    block: &'a DataBlock,
    offset: usize,
}

impl<'a> Iterator for DataBlockEntries<'a> {
    /// `(key content, value)`, the bytes `KeyValue::key().content()` and
    /// `KeyValue::value()` return for the entry, or an error for an entry
    /// that does not fit the block. The error ends the walk: a lookup that
    /// quietly stopped at a corrupt entry would report its keys as absent.
    type Item = Result<(&'a [u8], &'a [u8])>;

    fn next(&mut self) -> Option<Self::Item> {
        let data = &self.block.data[..self.block.content_end];
        let offset = self.offset;
        if offset + KEY_VALUE_HEADER_SIZE > data.len() {
            return None;
        }
        let truncated = || {
            HFileError::InvalidFormat(format!(
                "HFile data block entry at offset {offset} runs past the block"
            ))
        };
        let key_length = i32::from_be_bytes(data[offset..offset + 4].try_into().ok()?);
        let value_length = i32::from_be_bytes(data[offset + 4..offset + 8].try_into().ok()?);
        let (Ok(key_length), Ok(value_length)) =
            (usize::try_from(key_length), usize::try_from(value_length))
        else {
            self.offset = data.len();
            return Some(Err(truncated()));
        };
        let key_start = offset + KEY_VALUE_HEADER_SIZE;
        let value_start = key_start.saturating_add(key_length);
        let end = value_start.saturating_add(value_length);
        // One MVCC timestamp byte follows the value (`KeyValue::record_size`).
        if end.saturating_add(1) > data.len() {
            self.offset = data.len();
            return Some(Err(truncated()));
        }
        self.offset = end + 1;
        let content = key_content(&data[key_start..value_start]);
        Some(Ok((content, &data[value_start..end])))
    }
}

/// Read a Hadoop VLong encoded integer from bytes.
/// This is the encoding used in HFile root/leaf index blocks.
/// Returns (value, bytes_consumed).
///
/// Hadoop VLong format:
/// - First byte determines total size
/// - If first byte >= -112 (signed), value fits in single byte
/// - Otherwise, -111 - first_byte gives number of additional data bytes
pub fn read_var_long(bytes: &[u8], offset: usize) -> (u64, usize) {
    let first_byte = bytes[offset];
    let size = var_long_size_on_disk_single(first_byte);

    if size == 1 {
        // Single byte encoding: value is the byte itself (as signed, then cast to u64)
        return (first_byte as i8 as i64 as u64, 1);
    }

    // Multi-byte encoding: read size-1 bytes as big-endian
    let mut value: u64 = 0;
    for i in 0..size - 1 {
        value = (value << 8) | (bytes[offset + 1 + i] as u64);
    }

    // Check if negative (first byte < -120 in signed representation)
    let is_negative = (first_byte as i8) < -120;
    if is_negative {
        (!value, size)
    } else {
        (value, size)
    }
}

/// Calculate the size of a Hadoop VLong encoded integer from the first byte.
fn var_long_size_on_disk_single(first_byte: u8) -> usize {
    let signed = first_byte as i8;
    if signed >= -112 {
        1
    } else {
        (-111 - signed as i32) as usize
    }
}

/// Calculate the size of a Hadoop VLong encoded integer.
pub fn var_long_size_on_disk(bytes: &[u8], offset: usize) -> usize {
    var_long_size_on_disk_single(bytes[offset])
}

#[cfg(test)]
mod tests {
    /// One entry as the writer lays it out: lengths, structured key, value,
    /// MVCC byte.
    fn entry(content: &[u8], value: &[u8]) -> Vec<u8> {
        let mut key = (content.len() as u16).to_be_bytes().to_vec();
        key.extend_from_slice(content);
        key.push(0);
        key.extend_from_slice(&i64::MAX.to_be_bytes());
        key.push(4);
        let mut out = (key.len() as i32).to_be_bytes().to_vec();
        out.extend_from_slice(&(value.len() as i32).to_be_bytes());
        out.extend_from_slice(&key);
        out.extend_from_slice(value);
        out.push(0);
        out
    }

    fn data_block(bytes: Vec<u8>) -> super::DataBlock {
        let content_end = bytes.len();
        super::DataBlock {
            data: bytes,
            content_end,
        }
    }

    /// The borrowing walk sees what the copying iterator sees, and a
    /// truncated entry is an error rather than a silent end.
    #[test]
    fn entries_match_key_values_and_reject_a_truncated_entry() {
        let mut bytes = entry(b"k1", b"v1");
        bytes.extend(entry(b"k2", b"value two"));
        let block = data_block(bytes.clone());
        let walked: Vec<(Vec<u8>, Vec<u8>)> = block
            .entries()
            .map(|e| e.map(|(k, v)| (k.to_vec(), v.to_vec())).unwrap())
            .collect();
        let copied: Vec<(Vec<u8>, Vec<u8>)> = block
            .iter()
            .map(|kv| (kv.key().content().to_vec(), kv.value().to_vec()))
            .collect();
        assert_eq!(walked, copied);
        assert_eq!(walked[1], (b"k2".to_vec(), b"value two".to_vec()));

        // Drop the tail of the second entry.
        let truncated = data_block(bytes[..bytes.len() - 4].to_vec());
        let results: Vec<_> = truncated.entries().collect();
        assert_eq!(results.len(), 2);
        assert!(results[0].is_ok());
        assert!(matches!(
            results[1],
            Err(super::HFileError::InvalidFormat(_))
        ));
    }

    use super::*;

    #[test]
    fn test_block_header_parse_too_small() {
        let bytes = vec![0u8; 10]; // Too small for header
        let err = BlockHeader::parse(&bytes).unwrap_err();
        assert!(matches!(err, HFileError::InvalidFormat(_)));
    }

    #[test]
    fn test_block_header_on_disk_size_with_header() {
        // Create a minimal valid block header with DATA magic
        let mut bytes = vec![];
        bytes.extend_from_slice(b"DATABLK*"); // magic
        bytes.extend_from_slice(&100i32.to_be_bytes()); // on_disk_size_without_header
        bytes.extend_from_slice(&200i32.to_be_bytes()); // uncompressed_size_without_header
        bytes.extend_from_slice(&(-1i64).to_be_bytes()); // prev_block_offset
        bytes.push(1); // checksum_type
        bytes.extend_from_slice(&16384i32.to_be_bytes()); // bytes_per_checksum
        bytes.extend_from_slice(&133i32.to_be_bytes()); // on_disk_data_size_with_header

        let header = BlockHeader::parse(&bytes).unwrap();
        assert_eq!(header.on_disk_size_with_header(), BLOCK_HEADER_SIZE + 100);
        assert_eq!(header.block_type, HFileBlockType::Data);
    }

    #[test]
    fn test_parse_gz_block_at_checksum_chunk_boundary() {
        use flate2::{Compression, write::GzEncoder};
        use std::io::Write;

        let bytes_per_checksum = 16384usize;
        let checksum_size = SIZEOF_INT32;
        let (input, compressed) = (16000..16384)
            .find_map(|len| {
                let input = vec![b'x'; len];
                let mut encoder = GzEncoder::new(Vec::new(), Compression::none());
                encoder.write_all(&input).unwrap();
                let compressed = encoder.finish().unwrap();
                let data_size = BLOCK_HEADER_SIZE + compressed.len();
                (data_size == bytes_per_checksum - 1).then_some((input, compressed))
            })
            .expect("a stored GZIP block one byte below the checksum boundary");
        let data_size = BLOCK_HEADER_SIZE + compressed.len();
        let num_chunks = data_size.div_ceil(bytes_per_checksum);
        let checksum_bytes = num_chunks * checksum_size;
        assert!(num_chunks < (data_size + checksum_bytes).div_ceil(bytes_per_checksum));
        eprintln!(
            "synthetic GZIP boundary: N={data_size}, C={}, chunks={num_chunks}, checksum_bytes={checksum_bytes}",
            compressed.len()
        );

        let mut bytes = Vec::new();
        bytes.extend_from_slice(b"DATABLK*");
        bytes.extend_from_slice(&((compressed.len() + checksum_bytes) as i32).to_be_bytes());
        bytes.extend_from_slice(&(input.len() as i32).to_be_bytes());
        bytes.extend_from_slice(&(-1i64).to_be_bytes());
        bytes.push(1);
        bytes.extend_from_slice(&(bytes_per_checksum as i32).to_be_bytes());
        bytes.extend_from_slice(&(data_size as i32).to_be_bytes());
        assert_eq!(bytes.len(), BLOCK_HEADER_SIZE);
        bytes.extend_from_slice(&compressed);
        bytes.resize(bytes.len() + checksum_bytes, 0);

        let block = HFileBlock::parse(&bytes, CompressionCodec::Gzip).unwrap();
        assert_eq!(block.data, input);
    }

    #[test]
    fn test_block_index_entry_ordering() {
        let key1 = Key::from_bytes(vec![0, 3, b'a', b'a', b'a']);
        let key2 = Key::from_bytes(vec![0, 3, b'b', b'b', b'b']);

        let entry1 = BlockIndexEntry::new(key1.clone(), None, 0, 100);
        let entry2 = BlockIndexEntry::new(key2.clone(), None, 1000, 100);
        let entry3 = BlockIndexEntry::new(key1.clone(), None, 2000, 200);

        assert!(entry1 < entry2);
        assert_eq!(entry1, entry3); // Same key = equal
        assert_eq!(entry1.partial_cmp(&entry2), Some(std::cmp::Ordering::Less));
    }

    #[test]
    fn test_read_var_long_single_byte() {
        // Single byte values (first byte >= -112 as signed, i.e., 0-143 or 144-255 when unsigned)
        let bytes = vec![0u8]; // value 0
        let (value, size) = read_var_long(&bytes, 0);
        assert_eq!(value, 0);
        assert_eq!(size, 1);

        let bytes = vec![100u8]; // value 100
        let (value, size) = read_var_long(&bytes, 0);
        assert_eq!(value, 100);
        assert_eq!(size, 1);

        // 127 as unsigned byte is 127, as signed it's 127 (>= -112), so single byte
        let bytes = vec![127u8];
        let (value, size) = read_var_long(&bytes, 0);
        assert_eq!(value, 127);
        assert_eq!(size, 1);
    }

    #[test]
    fn test_read_var_long_multi_byte() {
        // Multi-byte encoding: first byte < -112 (signed)
        // -113 as signed = 143 as unsigned = 0x8F
        // This means 2 additional bytes (size = -111 - (-113) = 2, total size = 3)
        // Actually let's test with known values

        // For value 1000:
        // In Hadoop VLong, values 128-255 use 2 bytes
        // First byte = -113 (0x8F) means 2 additional bytes
        // But let's use a simpler approach - test with offset

        let bytes = vec![50u8, 0x8F, 0x03, 0xE8]; // offset 0 = 50 (single byte)
        let (value, size) = read_var_long(&bytes, 0);
        assert_eq!(value, 50);
        assert_eq!(size, 1);
    }

    #[test]
    fn test_var_long_size_on_disk() {
        // Single byte values (first byte >= -112 as signed)
        // 0..=127 are positive when signed, so single byte
        assert_eq!(var_long_size_on_disk(&[0], 0), 1);
        assert_eq!(var_long_size_on_disk(&[100], 0), 1);
        assert_eq!(var_long_size_on_disk(&[127], 0), 1);
        // 128..=143 are -128..-113 as signed, still >= -112? No.
        // 143 as u8 = -113 as i8, which is >= -112? -113 >= -112? No, -113 < -112
        // So 143 should be multi-byte. Let's check: -111 - (-113) = 2
        assert_eq!(var_long_size_on_disk(&[143], 0), 2); // -113 as signed, size = 2

        // 144 as u8 = -112 as i8, which is >= -112, so single byte
        assert_eq!(var_long_size_on_disk(&[144], 0), 1);
    }

    #[test]
    fn test_data_block_is_valid_offset() {
        // Create a simple HFileBlock
        let mut header_bytes = vec![];
        header_bytes.extend_from_slice(b"DATABLK*"); // magic
        header_bytes.extend_from_slice(&100i32.to_be_bytes()); // on_disk_size
        header_bytes.extend_from_slice(&50i32.to_be_bytes()); // uncompressed_size
        header_bytes.extend_from_slice(&(-1i64).to_be_bytes()); // prev_block_offset
        header_bytes.push(1); // checksum_type
        header_bytes.extend_from_slice(&16384i32.to_be_bytes()); // bytes_per_checksum
        header_bytes.extend_from_slice(&133i32.to_be_bytes()); // on_disk_data_size

        // Add data after header
        header_bytes.extend_from_slice(&[0u8; 50]); // 50 bytes of data

        let block = HFileBlock::parse(&header_bytes, CompressionCodec::None).unwrap();
        let data_block = DataBlock::from_block(block);

        assert!(data_block.is_valid_offset(0));
        assert!(data_block.is_valid_offset(49));
        assert!(!data_block.is_valid_offset(50));
        assert!(!data_block.is_valid_offset(100));
    }
}