micromegas_analytics/lakehouse/
partition_metadata.rs

1use anyhow::{Context, Result};
2use bytes::Bytes;
3use datafusion::parquet::arrow::async_reader::MetadataFetch;
4use datafusion::parquet::errors::ParquetError;
5use datafusion::parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader};
6use futures::FutureExt;
7use futures::future::BoxFuture;
8use micromegas_tracing::prelude::*;
9use object_store::{ObjectStore, ObjectStoreExt, path::Path};
10use std::ops::Range;
11use std::sync::Arc;
12
13use super::metadata_cache::MetadataCache;
14
15/// Strips column index information from Parquet metadata
16///
17/// This removes column_index_offset and column_index_length from ColumnChunk metadata
18/// to prevent DataFusion from trying to read legacy ColumnIndex structures that may
19/// have incomplete or malformed null_pages fields (required in Arrow 57.0+).
20///
21/// The approach: serialize metadata to thrift, modify it, then re-parse.
22#[allow(deprecated)]
23fn strip_column_index_info(metadata: ParquetMetaData) -> Result<ParquetMetaData> {
24    use datafusion::parquet::file::metadata::ParquetMetaDataWriter;
25    use parquet::format::FileMetaData as ThriftFileMetaData;
26    use parquet::thrift::TSerializable;
27    use thrift::protocol::{TCompactInputProtocol, TCompactOutputProtocol, TOutputProtocol};
28    // Serialize metadata using ParquetMetaDataWriter
29    let mut buffer = Vec::new();
30    let writer = ParquetMetaDataWriter::new(&mut buffer, &metadata);
31    writer.finish()?;
32    // Extract FileMetaData portion: the parquet footer is laid out as
33    // [Page Indexes][FileMetaData][Length][PAR1]
34    let metadata_len = u32::from_le_bytes([
35        buffer[buffer.len() - 8],
36        buffer[buffer.len() - 7],
37        buffer[buffer.len() - 6],
38        buffer[buffer.len() - 5],
39    ]) as usize;
40    let file_metadata_start = buffer.len() - 8 - metadata_len;
41    let file_metadata_bytes = &buffer[file_metadata_start..buffer.len() - 8];
42    // Parse FileMetaData with thrift
43    let mut transport =
44        thrift::transport::TBufferChannel::with_capacity(file_metadata_bytes.len(), 0);
45    transport.set_readable_bytes(file_metadata_bytes);
46    let mut protocol = TCompactInputProtocol::new(transport);
47    let mut thrift_meta = ThriftFileMetaData::read_from_in_protocol(&mut protocol)
48        .context("parsing thrift metadata to strip column index")?;
49    // Remove column index information from all row groups and columns
50    for rg in thrift_meta.row_groups.iter_mut() {
51        for col in rg.columns.iter_mut() {
52            col.column_index_offset = None;
53            col.column_index_length = None;
54            // Also remove offset index for consistency
55            col.offset_index_offset = None;
56            col.offset_index_length = None;
57        }
58    }
59    // Re-serialize - use Vec<u8> which auto-grows as needed
60    let mut modified_bytes: Vec<u8> = Vec::with_capacity(file_metadata_bytes.len() * 2);
61    let mut out_protocol = TCompactOutputProtocol::new(&mut modified_bytes);
62    thrift_meta
63        .write_to_out_protocol(&mut out_protocol)
64        .context("serializing modified thrift metadata")?;
65    out_protocol.flush()?;
66    // Parse back to ParquetMetaData
67    ParquetMetaDataReader::decode_metadata(&Bytes::copy_from_slice(&modified_bytes))
68        .context("re-parsing metadata after stripping column index")
69}
70
71/// Adapts `ObjectStore::get_range` to the `MetadataFetch` interface expected by
72/// `ParquetMetaDataReader`, so the footer read benefits from the same
73/// object-store-backed byte caching as the rest of the file (the object store
74/// itself may be L1-cache-backed, see `object_cache::l1_wrap`).
75///
76/// Also tallies the number of footer bytes fetched via `bytes_read`, so callers can use it
77/// as a cheap proxy for the parsed metadata's weight in `MetadataCache` (see
78/// `load_partition_metadata`).
79struct ObjectStoreFetch<'a> {
80    object_store: &'a Arc<dyn ObjectStore>,
81    path: &'a Path,
82    bytes_read: std::sync::Arc<std::sync::atomic::AtomicU64>,
83}
84
85impl MetadataFetch for ObjectStoreFetch<'_> {
86    fn fetch(
87        &mut self,
88        range: Range<u64>,
89    ) -> BoxFuture<'_, datafusion::parquet::errors::Result<Bytes>> {
90        self.bytes_read.fetch_add(
91            range.end - range.start,
92            std::sync::atomic::Ordering::Relaxed,
93        );
94        let object_store = self.object_store;
95        let path = self.path;
96        async move {
97            object_store
98                .get_range(path, range)
99                .await
100                .map_err(|e| ParquetError::External(Box::new(e)))
101        }
102        .boxed()
103    }
104}
105
106/// Read and parse a partition's parquet footer directly from `object_store` — the sole
107/// partition-metadata read path. Checks and backfills the shared `MetadataCache`
108/// parsed-lookaside: on a cache hit this returns immediately, on a miss it reads the
109/// footer from `object_store` (which may itself be L1-cache-backed) and stores the parsed
110/// result before returning it.
111#[span_fn]
112pub async fn load_partition_metadata(
113    object_store: &Arc<dyn ObjectStore>,
114    path: &Path,
115    file_size: u64,
116    cache: Option<&MetadataCache>,
117) -> datafusion::parquet::errors::Result<Arc<ParquetMetaData>> {
118    let file_path = path.as_ref();
119
120    // Check cache first
121    if let Some(cache) = cache
122        && let Some(metadata) = cache.get(file_path).await
123    {
124        debug!("cache hit for partition metadata path={file_path}");
125        return Ok(metadata);
126    }
127
128    let start = std::time::Instant::now();
129    let bytes_read = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
130    let raw = ParquetMetaDataReader::new()
131        .load_and_finish(
132            ObjectStoreFetch {
133                object_store,
134                path,
135                bytes_read: bytes_read.clone(),
136            },
137            file_size,
138        )
139        .await?;
140    let stripped = strip_column_index_info(raw).map_err(|e| ParquetError::External(e.into()))?;
141    let duration_ms = start.elapsed().as_millis();
142    debug!(
143        "partition_metadata_footer_read file={file_path} file_size={file_size} duration_ms={duration_ms}"
144    );
145    let result = Arc::new(stripped);
146
147    // Store in cache. There's no pre-serialized footer blob to measure the size of here, so
148    // use the number of footer bytes actually read from object storage as the weight: it's a
149    // natural, cheap proxy for the parsed metadata's size, without paying for an extra
150    // re-serialization pass just to compute a weight.
151    if let Some(cache) = cache {
152        let weight = bytes_read.load(std::sync::atomic::Ordering::Relaxed);
153        let weight = u32::try_from(weight).unwrap_or(u32::MAX);
154        cache
155            .insert(file_path.to_string(), result.clone(), weight)
156            .await;
157    }
158
159    Ok(result)
160}
micromegas_analytics/lakehouse/partition_metadata.rs

micromegas_analytics/lakehouse/
partition_metadata.rs
