parx_rs/

reader.rs

/*
 * Copyright 2026 PARX Authors
 *
 * Licensed 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.
 */
//! PARX file reader.

use crate::compression;
use crate::error::{ParxError, Result};
use crate::format::{
    Compression, Header, Trailer, HEADER_SIZE, MAGIC, MIN_FILE_SIZE, TRAILER_SIZE,
};
use crate::proto::ParxManifest;
use bytes::Bytes;
use prost::Message;
use std::ops::Range;

#[derive(Debug, Clone)]
enum Payload {
    Borrowed(Bytes),
    Owned(Bytes),
}

impl Payload {
    fn as_slice(&self) -> &[u8] {
        match self {
            Self::Borrowed(bytes) | Self::Owned(bytes) => bytes.as_ref(),
        }
    }
}

/// Reader for PARX sidecar files.
///
/// Parses and validates PARX files, providing access to the cached Parquet footer bytes
/// and optional page index data.
#[derive(Debug, Clone)]
pub struct ParxReader {
    header: Header,
    manifest: ParxManifest,
    footer_bytes: Payload,
    page_index_bytes: Option<Payload>,
}

impl ParxReader {
    /// Open a PARX file from bytes.
    ///
    /// This validates the file structure, checksums, and extracts the manifest,
    /// footer payload, and optional page index data.
    ///
    /// # Errors
    /// Returns error if file is invalid or corrupted.
    ///
    /// # Panics
    /// Panics if file size is less than minimum header size.
    pub fn open(bytes: &[u8]) -> Result<Self> {
        Self::open_with_payload(bytes, |range| {
            Payload::Owned(Bytes::copy_from_slice(&bytes[range]))
        })
    }

    fn open_with_payload<F>(bytes: &[u8], make_payload: F) -> Result<Self>
    where
        F: Fn(Range<usize>) -> Payload,
    {
        let file_size = bytes.len();

        // Check minimum size
        if file_size < MIN_FILE_SIZE {
            return Err(ParxError::FileTooSmall {
                size: file_size,
                minimum: MIN_FILE_SIZE,
            });
        }

        // Parse header
        let header_bytes: [u8; HEADER_SIZE] = bytes[..HEADER_SIZE]
            .try_into()
            .expect("header slice length verified above");
        let header = Header::from_bytes(&header_bytes);

        // Validate header magic
        if !header.is_magic_valid(MAGIC) {
            return Err(ParxError::InvalidMagic(header.magic));
        }

        // Validate version
        if !header.is_version_supported() {
            return Err(ParxError::UnsupportedVersion {
                major: header.version_major,
                minor: header.version_minor,
            });
        }

        // Parse trailer
        let trailer_bytes: [u8; TRAILER_SIZE] = bytes[file_size - TRAILER_SIZE..]
            .try_into()
            .expect("trailer slice length verified above");
        let trailer = Trailer::from_bytes(&trailer_bytes);

        // Validate trailer magic
        if !trailer.is_magic_valid(MAGIC) {
            return Err(ParxError::InvalidMagic(trailer.magic));
        }

        // Extract and validate manifest
        let manifest_end = file_size - TRAILER_SIZE;
        let manifest_start = manifest_end
            .checked_sub(trailer.manifest_len as usize)
            .ok_or(ParxError::FileTooSmall {
                size: file_size,
                minimum: MIN_FILE_SIZE + trailer.manifest_len as usize,
            })?;

        let manifest_bytes = &bytes[manifest_start..manifest_end];

        // Verify manifest CRC32C
        let actual_crc = crc32c::crc32c(manifest_bytes);
        if actual_crc != trailer.manifest_crc32c {
            return Err(ParxError::ManifestChecksumMismatch {
                expected: trailer.manifest_crc32c,
                actual: actual_crc,
            });
        }

        // Decode manifest
        let manifest = ParxManifest::decode(manifest_bytes)?;

        // Extract footer bytes
        let footer_offset = usize::try_from(manifest.footer_offset).map_err(|_| {
            ParxError::InvalidPayloadBounds {
                offset: manifest.footer_offset,
                length: manifest.footer_length,
                file_size: file_size as u64,
            }
        })?;
        let footer_length = usize::try_from(manifest.footer_length).map_err(|_| {
            ParxError::InvalidPayloadBounds {
                offset: manifest.footer_offset,
                length: manifest.footer_length,
                file_size: file_size as u64,
            }
        })?;
        let footer_end =
            footer_offset
                .checked_add(footer_length)
                .ok_or(ParxError::InvalidPayloadBounds {
                    offset: manifest.footer_offset,
                    length: manifest.footer_length,
                    file_size: file_size as u64,
                })?;

        // Validate footer bounds
        if footer_offset < HEADER_SIZE || footer_end > manifest_start {
            return Err(ParxError::InvalidPayloadBounds {
                offset: manifest.footer_offset,
                length: manifest.footer_length,
                file_size: file_size as u64,
            });
        }

        let stored_footer_bytes = &bytes[footer_offset..footer_end];

        // Verify footer checksum (on stored/compressed bytes)
        let footer_crc = crc32c::crc32c(stored_footer_bytes);
        if footer_crc.to_le_bytes().as_slice() != manifest.footer_checksum.as_slice() {
            return Err(ParxError::FooterChecksumMismatch);
        }

        // Decompress footer if needed
        let footer_bytes = if let Some(algo) = header.compression_algorithm() {
            let uncompressed_size =
                usize::try_from(manifest.footer_uncompressed_size).map_err(|_| {
                    ParxError::InvalidFormat("footer uncompressed size too large".to_string())
                })?;
            Payload::Owned(Bytes::from(compression::decompress(
                stored_footer_bytes,
                algo,
                uncompressed_size,
            )?))
        } else {
            make_payload(footer_offset..footer_end)
        };

        // Extract page index bytes if present
        let page_index_bytes = if manifest.page_index_length > 0 {
            let page_index_offset = usize::try_from(manifest.page_index_offset).map_err(|_| {
                ParxError::InvalidPayloadBounds {
                    offset: manifest.page_index_offset,
                    length: manifest.page_index_length,
                    file_size: file_size as u64,
                }
            })?;
            let page_index_length = usize::try_from(manifest.page_index_length).map_err(|_| {
                ParxError::InvalidPayloadBounds {
                    offset: manifest.page_index_offset,
                    length: manifest.page_index_length,
                    file_size: file_size as u64,
                }
            })?;
            let page_index_end = page_index_offset.checked_add(page_index_length).ok_or(
                ParxError::InvalidPayloadBounds {
                    offset: manifest.page_index_offset,
                    length: manifest.page_index_length,
                    file_size: file_size as u64,
                },
            )?;

            // Validate page index bounds
            if page_index_offset < footer_end || page_index_end > manifest_start {
                return Err(ParxError::InvalidPayloadBounds {
                    offset: manifest.page_index_offset,
                    length: manifest.page_index_length,
                    file_size: file_size as u64,
                });
            }

            let stored_page_index_bytes = &bytes[page_index_offset..page_index_end];

            // Verify page index checksum
            let page_index_crc = crc32c::crc32c(stored_page_index_bytes);
            if page_index_crc.to_le_bytes().as_slice() != manifest.page_index_checksum.as_slice() {
                return Err(ParxError::PageIndexChecksumMismatch);
            }

            // Decompress page indexes if needed (same compression as footer)
            let page_indexes = if let Some(algo) = header.compression_algorithm() {
                let uncompressed_size = usize::try_from(manifest.page_index_uncompressed_size)
                    .map_err(|_| {
                        ParxError::InvalidFormat(
                            "page index uncompressed size too large".to_string(),
                        )
                    })?;
                Payload::Owned(Bytes::from(compression::decompress(
                    stored_page_index_bytes,
                    algo,
                    uncompressed_size,
                )?))
            } else {
                make_payload(page_index_offset..page_index_end)
            };

            Some(page_indexes)
        } else {
            None
        };

        Ok(Self {
            header,
            manifest,
            footer_bytes,
            page_index_bytes,
        })
    }

    /// Open a PARX file from owned Bytes (zero-copy where possible).
    /// Use this when you already have `Bytes` (e.g., from object_store) to potentially reduce allocations.
    ///
    /// # Errors
    /// Returns error if file is invalid or corrupted.
    pub fn open_bytes(bytes: &Bytes) -> Result<Self> {
        Self::open_with_payload(bytes, |range| Payload::Borrowed(bytes.slice(range)))
    }

    /// Get the parsed header.
    #[inline]
    pub const fn header(&self) -> &Header {
        &self.header
    }

    /// Get the parsed manifest.
    #[inline]
    pub const fn manifest(&self) -> &ParxManifest {
        &self.manifest
    }

    /// Get the cached Parquet footer bytes (decompressed if necessary).
    #[inline]
    pub fn footer_bytes(&self) -> &[u8] {
        self.footer_bytes.as_slice()
    }

    /// Check if the footer was stored compressed.
    #[inline]
    pub const fn is_compressed(&self) -> bool {
        self.header.is_footer_compressed()
    }

    /// Get the compression algorithm used, if any.
    #[inline]
    pub const fn compression_algorithm(&self) -> Option<Compression> {
        self.header.compression_algorithm()
    }

    /// Get the original uncompressed footer size.
    ///
    /// Returns 0 if footer was not compressed.
    #[inline]
    pub const fn uncompressed_footer_size(&self) -> u64 {
        self.manifest.footer_uncompressed_size
    }

    // === Page Index Methods ===

    /// Check if this PARX file contains page index data.
    #[inline]
    pub const fn has_page_indexes(&self) -> bool {
        self.page_index_bytes.is_some()
    }

    /// Get raw page index payload bytes, if present (decompressed).
    ///
    /// Page indexes include concatenated ColumnIndex and OffsetIndex structures
    /// from Parquet V2 format.
    #[inline]
    pub fn page_index_bytes(&self) -> Option<&[u8]> {
        self.page_index_bytes.as_ref().map(Payload::as_slice)
    }

    /// Get the original uncompressed page index size.
    ///
    /// Returns 0 if page indexes were not compressed.
    #[inline]
    pub const fn uncompressed_page_index_size(&self) -> u64 {
        self.manifest.page_index_uncompressed_size
    }

    // === Validation Methods ===

    /// Validate that this PARX file matches the given Parquet file size.
    ///
    /// This is the fast validation check that should always be performed
    /// before using cached footer bytes.
    #[inline]
    pub const fn validate_source_size(&self, source_size: u64) -> bool {
        self.manifest.source_size == source_size
    }

    /// Validate that this PARX file matches the given Parquet footer hash.
    ///
    /// This is the paranoid validation check that re-hashes the original
    /// Parquet footer bytes. Only use when you have the original footer.
    pub fn validate_source_footer(&self, original_footer: &[u8]) -> bool {
        if self.manifest.source_footer_checksum.len() != 4 {
            return false; // Wrong checksum length
        }
        let footer_crc32c = crc32c::crc32c(original_footer);
        footer_crc32c.to_le_bytes().as_slice() == self.manifest.source_footer_checksum.as_slice()
    }

    /// Get the source URI from the manifest (may be empty).
    #[inline]
    pub fn source_uri(&self) -> &str {
        &self.manifest.source_uri
    }

    /// Get the source file size from the manifest.
    #[inline]
    pub const fn source_size(&self) -> u64 {
        self.manifest.source_size
    }

    /// Get the creation timestamp in milliseconds.
    #[inline]
    pub const fn created_at_ms(&self) -> u64 {
        self.manifest.created_at_ms
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::writer::ParxWriter;

    #[test]
    fn test_roundtrip() {
        let footer_bytes = b"fake parquet footer data for testing";
        let source_size = 1024 * 1024; // 1 MB

        let mut writer = ParxWriter::new();
        writer.set_source_uri("s3://bucket/table/part-0000.parquet");
        writer.set_source_size(source_size);
        writer.set_footer(footer_bytes);

        let parx_bytes = writer.finish();

        let reader = ParxReader::open(&parx_bytes).expect("failed to open PARX");

        assert_eq!(reader.footer_bytes(), footer_bytes);
        assert_eq!(reader.source_size(), source_size);
        assert!(reader.validate_source_size(source_size));
        assert!(!reader.validate_source_size(source_size + 1));
        assert_eq!(reader.source_uri(), "s3://bucket/table/part-0000.parquet");
        assert!(!reader.is_compressed());
    }

    #[test]
    fn test_open_bytes() {
        let footer_bytes = b"test footer";
        let source_size = 500;

        let mut writer = ParxWriter::new();
        writer.set_source_size(source_size);
        writer.set_footer(footer_bytes);

        let parx_bytes = Bytes::from(writer.finish());
        let reader = ParxReader::open_bytes(&parx_bytes).expect("failed to open PARX");

        assert_eq!(reader.footer_bytes(), footer_bytes);
        let footer_offset = HEADER_SIZE;
        assert_eq!(
            reader.footer_bytes().as_ptr(),
            parx_bytes[footer_offset..footer_offset + footer_bytes.len()].as_ptr()
        );
    }

    #[test]
    fn test_invalid_checksum_length() {
        // Create a valid PARX file
        let mut writer = ParxWriter::new();
        writer.set_footer(b"test footer");
        writer.set_source_size(1000);
        let parx_bytes = writer.finish();

        // Manually corrupt the manifest to have wrong checksum length
        // (This requires parsing and re-encoding the manifest)
        // For now, let's test that validate_source_footer handles it gracefully

        let reader = ParxReader::open(&parx_bytes).unwrap();

        // Valid checksum (4 bytes)
        assert!(reader.validate_source_footer(b"test footer"));
    }

    #[test]
    fn test_source_footer_validation() {
        let mut writer = ParxWriter::new();
        let footer = b"test footer bytes";
        writer.set_footer(footer);
        writer.set_source_size(1000);
        let parx_bytes = writer.finish();

        let reader = ParxReader::open(&parx_bytes).unwrap();

        // Correct footer validates
        assert!(reader.validate_source_footer(footer));

        // Wrong footer fails
        assert!(!reader.validate_source_footer(b"wrong footer"));

        // Empty footer fails
        assert!(!reader.validate_source_footer(b""));
    }

    #[test]
    fn test_roundtrip_with_compression() {
        let footer_bytes = b"test footer data for compression".repeat(100);
        let source_size = 1000;

        for algo in [Compression::Zstd, Compression::Lz4, Compression::Gzip] {
            let mut writer = ParxWriter::new();
            writer.set_source_size(source_size);
            writer.set_footer(&footer_bytes);
            writer.set_compression(algo);

            let parx_bytes = writer.finish();
            let reader = ParxReader::open(&parx_bytes).expect("failed to open PARX");

            assert_eq!(reader.footer_bytes(), footer_bytes.as_slice());
            assert!(reader.is_compressed());
            assert_eq!(reader.compression_algorithm(), Some(algo));
            assert_eq!(reader.uncompressed_footer_size(), footer_bytes.len() as u64);
        }
    }

    #[test]
    fn test_open_bytes_with_page_indexes_borrows_uncompressed_payloads() {
        let mut writer = ParxWriter::new();
        writer.set_source_size(1000);
        writer.set_footer(b"footer");
        writer.set_page_indexes(b"page-index");

        let parx_bytes = Bytes::from(writer.finish());
        let reader = ParxReader::open_bytes(&parx_bytes).expect("failed to open PARX");

        assert_eq!(reader.footer_bytes(), b"footer");
        assert_eq!(reader.page_index_bytes(), Some(b"page-index".as_slice()));
        assert_eq!(
            reader.footer_bytes().as_ptr(),
            parx_bytes[HEADER_SIZE..HEADER_SIZE + b"footer".len()].as_ptr()
        );
    }
}