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+).
20fn strip_column_index_info(metadata: ParquetMetaData) -> Result<ParquetMetaData> {
21    let mut builder = metadata.into_builder();
22    let row_groups = builder
23        .take_row_groups()
24        .into_iter()
25        .map(|row_group| {
26            let mut row_group_builder = row_group.into_builder();
27            let columns = row_group_builder
28                .take_columns()
29                .into_iter()
30                .map(|column| {
31                    column
32                        .into_builder()
33                        .set_column_index_offset(None)
34                        .set_column_index_length(None)
35                        .set_offset_index_offset(None)
36                        .set_offset_index_length(None)
37                        .build()
38                        .context("rebuilding column chunk metadata without column index pointers")
39                })
40                .collect::<Result<Vec<_>>>()?;
41            row_group_builder
42                .set_column_metadata(columns)
43                .build()
44                .context("rebuilding row group metadata without column index pointers")
45        })
46        .collect::<Result<Vec<_>>>()?;
47    Ok(builder.set_row_groups(row_groups).build())
48}
49
50/// Adapts `ObjectStore::get_range` to the `MetadataFetch` interface expected by
51/// `ParquetMetaDataReader`, so the footer read benefits from the same
52/// object-store-backed byte caching as the rest of the file (the object store
53/// itself may be L1-cache-backed, see `object_cache::l1_wrap`).
54///
55/// Also tallies the number of footer bytes fetched via `bytes_read`, so callers can use it
56/// as a cheap proxy for the parsed metadata's weight in `MetadataCache` (see
57/// `load_partition_metadata`).
58struct ObjectStoreFetch<'a> {
59    object_store: &'a Arc<dyn ObjectStore>,
60    path: &'a Path,
61    bytes_read: std::sync::Arc<std::sync::atomic::AtomicU64>,
62}
63
64impl MetadataFetch for ObjectStoreFetch<'_> {
65    fn fetch(
66        &mut self,
67        range: Range<u64>,
68    ) -> BoxFuture<'_, datafusion::parquet::errors::Result<Bytes>> {
69        self.bytes_read.fetch_add(
70            range.end - range.start,
71            std::sync::atomic::Ordering::Relaxed,
72        );
73        let object_store = self.object_store;
74        let path = self.path;
75        async move {
76            object_store
77                .get_range(path, range)
78                .await
79                .map_err(|e| ParquetError::External(Box::new(e)))
80        }
81        .boxed()
82    }
83}
84
85/// Read and parse a partition's parquet footer directly from `object_store` — the sole
86/// partition-metadata read path. Checks and backfills the shared `MetadataCache`
87/// parsed-lookaside: on a cache hit this returns immediately, on a miss it reads the
88/// footer from `object_store` (which may itself be L1-cache-backed) and stores the parsed
89/// result before returning it.
90#[span_fn]
91pub async fn load_partition_metadata(
92    object_store: &Arc<dyn ObjectStore>,
93    path: &Path,
94    file_size: u64,
95    cache: Option<&MetadataCache>,
96) -> datafusion::parquet::errors::Result<Arc<ParquetMetaData>> {
97    let file_path = path.as_ref();
98
99    // Check cache first
100    if let Some(cache) = cache
101        && let Some(metadata) = cache.get(file_path).await
102    {
103        debug!("cache hit for partition metadata path={file_path}");
104        return Ok(metadata);
105    }
106
107    let start = std::time::Instant::now();
108    let bytes_read = std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0));
109    let raw = ParquetMetaDataReader::new()
110        .load_and_finish(
111            ObjectStoreFetch {
112                object_store,
113                path,
114                bytes_read: bytes_read.clone(),
115            },
116            file_size,
117        )
118        .await?;
119    let stripped = strip_column_index_info(raw).map_err(|e| ParquetError::External(e.into()))?;
120    let duration_ms = start.elapsed().as_millis();
121    debug!(
122        "partition_metadata_footer_read file={file_path} file_size={file_size} duration_ms={duration_ms}"
123    );
124    let result = Arc::new(stripped);
125
126    // Store in cache. There's no pre-serialized footer blob to measure the size of here, so
127    // use the number of footer bytes actually read from object storage as the weight: it's a
128    // natural, cheap proxy for the parsed metadata's size, without paying for an extra
129    // re-serialization pass just to compute a weight.
130    if let Some(cache) = cache {
131        let weight = bytes_read.load(std::sync::atomic::Ordering::Relaxed);
132        let weight = u32::try_from(weight).unwrap_or(u32::MAX);
133        cache
134            .insert(file_path.to_string(), result.clone(), weight)
135            .await;
136    }
137
138    Ok(result)
139}
micromegas_analytics/lakehouse/partition_metadata.rs

micromegas_analytics/lakehouse/
partition_metadata.rs
