Skip to main content

nautilus_persistence/backend/parquet/
metadata.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Parquet catalog metadata queries.
17
18use std::collections::{BTreeMap, HashMap};
19
20use datafusion::arrow::record_batch::RecordBatch;
21use nautilus_core::UnixNanos;
22use nautilus_serialization::arrow::U64ColumnRef;
23
24use crate::{
25    backend::parquet::{
26        catalog::ParquetDataCatalog, io::read_parquet_schema_from_object_store,
27        paths::make_sql_safe_identifier,
28    },
29    catalog::{
30        traits::CatalogMetadata,
31        types::{CatalogDataType, parquet_catalog_data_type_table_stem},
32    },
33    common::{datafusion::build_query, metadata::arrow_metadata_to_params},
34};
35
36impl ParquetDataCatalog {
37    /// Queries Arrow schema metadata and the first queried timestamp where each metadata is used.
38    ///
39    /// # Errors
40    ///
41    /// Returns an error if file discovery, Parquet metadata reading, or query execution fails.
42    pub fn query_metadata(
43        &mut self,
44        data_type: &CatalogDataType,
45        identifiers: Option<Vec<String>>,
46        start: Option<UnixNanos>,
47        end: Option<UnixNanos>,
48        where_clause: Option<&str>,
49    ) -> anyhow::Result<Vec<CatalogMetadata>> {
50        self.clear_session_tables();
51        self.register_remote_object_store()?;
52
53        let files_list = self.query_files(data_type, identifiers, start, end)?;
54        let table_prefix =
55            make_sql_safe_identifier(&parquet_catalog_data_type_table_stem(data_type));
56        let mut metadata_by_key: BTreeMap<String, CatalogMetadata> = BTreeMap::new();
57
58        for (index, file_uri) in files_list.iter().enumerate() {
59            let object_path = self.to_object_path_parsed(file_uri)?;
60
61            let metadata = self.execute_async(|| async {
62                let schema =
63                    read_parquet_schema_from_object_store(self.object_store.clone(), &object_path)
64                        .await?;
65                Ok::<HashMap<String, String>, anyhow::Error>(schema.metadata().clone())
66            })?;
67
68            let table_name = format!("{table_prefix}_{index}");
69            let query = build_query(&table_name, start, end, where_clause);
70            let resolved_path = self.resolve_path_for_datafusion(file_uri);
71            let batches = self.session.collect_parquet_files_batches(
72                &table_name,
73                vec![resolved_path],
74                Some(&query),
75            )?;
76
77            let Some(first_ts_init) = first_ts_init_from_batches(&batches)? else {
78                continue;
79            };
80
81            let key = canonical_metadata_key(&metadata)?;
82            let metadata = arrow_metadata_to_params(&metadata);
83
84            metadata_by_key
85                .entry(key)
86                .and_modify(|existing| {
87                    existing.first_ts_init = existing.first_ts_init.min(first_ts_init);
88                })
89                .or_insert(CatalogMetadata {
90                    first_ts_init,
91                    metadata,
92                });
93        }
94
95        let mut metadata = metadata_by_key.into_values().collect::<Vec<_>>();
96        metadata.sort_by_key(|item| item.first_ts_init);
97        Ok(metadata)
98    }
99}
100
101fn first_ts_init_from_batches(batches: &[RecordBatch]) -> anyhow::Result<Option<UnixNanos>> {
102    let mut first_ts_init: Option<u64> = None;
103
104    for batch in batches {
105        if batch.num_rows() == 0 {
106            continue;
107        }
108
109        let column = batch
110            .column_by_name("ts_init")
111            .ok_or_else(|| anyhow::anyhow!("ts_init column not found"))?;
112        let ts_init = U64ColumnRef::try_from_array(column.as_ref())
113            .ok_or_else(|| anyhow::anyhow!("ts_init column has an unsupported type"))?;
114
115        for row in 0..batch.num_rows() {
116            if ts_init.is_null(row) {
117                continue;
118            }
119
120            let value = ts_init
121                .value(row)
122                .ok_or_else(|| anyhow::anyhow!("ts_init value cannot be negative"))?;
123            first_ts_init = Some(first_ts_init.map_or(value, |current| current.min(value)));
124        }
125    }
126
127    Ok(first_ts_init.map(UnixNanos::from))
128}
129
130fn canonical_metadata_key(metadata: &HashMap<String, String>) -> anyhow::Result<String> {
131    let ordered = metadata.iter().collect::<BTreeMap<_, _>>();
132    Ok(serde_json::to_string(&ordered)?)
133}