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            let metadata = self.execute_async(|| async {
61                let schema =
62                    read_parquet_schema_from_object_store(self.object_store.clone(), &object_path)
63                        .await?;
64                Ok::<HashMap<String, String>, anyhow::Error>(schema.metadata().clone())
65            })?;
66
67            let table_name = format!("{table_prefix}_{index}");
68            let query = build_query(&table_name, start, end, where_clause);
69            let resolved_path = self.resolve_path_for_datafusion(file_uri);
70            let batches = self.session.collect_parquet_files_batches(
71                &table_name,
72                vec![resolved_path],
73                Some(&query),
74            )?;
75
76            let Some(first_ts_init) = first_ts_init_from_batches(&batches)? else {
77                continue;
78            };
79
80            let key = canonical_metadata_key(&metadata)?;
81            let metadata = arrow_metadata_to_params(&metadata);
82
83            match metadata_by_key.get_mut(&key) {
84                Some(existing) => {
85                    if first_ts_init < existing.first_ts_init {
86                        existing.first_ts_init = first_ts_init;
87                    }
88                }
89                None => {
90                    metadata_by_key.insert(
91                        key,
92                        CatalogMetadata {
93                            first_ts_init,
94                            metadata,
95                        },
96                    );
97                }
98            }
99        }
100
101        let mut metadata = metadata_by_key.into_values().collect::<Vec<_>>();
102        metadata.sort_by_key(|item| item.first_ts_init);
103        Ok(metadata)
104    }
105}
106
107fn first_ts_init_from_batches(batches: &[RecordBatch]) -> anyhow::Result<Option<UnixNanos>> {
108    let mut first_ts_init: Option<u64> = None;
109
110    for batch in batches {
111        if batch.num_rows() == 0 {
112            continue;
113        }
114
115        let column = batch
116            .column_by_name("ts_init")
117            .ok_or_else(|| anyhow::anyhow!("ts_init column not found"))?;
118        let ts_init = U64ColumnRef::try_from_array(column.as_ref())
119            .ok_or_else(|| anyhow::anyhow!("ts_init column has an unsupported type"))?;
120
121        for row in 0..batch.num_rows() {
122            if ts_init.is_null(row) {
123                continue;
124            }
125
126            let value = ts_init
127                .value(row)
128                .ok_or_else(|| anyhow::anyhow!("ts_init value cannot be negative"))?;
129            first_ts_init = Some(first_ts_init.map_or(value, |current| current.min(value)));
130        }
131    }
132
133    Ok(first_ts_init.map(UnixNanos::from))
134}
135
136fn canonical_metadata_key(metadata: &HashMap<String, String>) -> anyhow::Result<String> {
137    let ordered = metadata
138        .iter()
139        .map(|(key, value)| (key.clone(), value.clone()))
140        .collect::<BTreeMap<_, _>>();
141    Ok(serde_json::to_string(&ordered)?)
142}