nautilus_persistence/backend/parquet/
metadata.rs1use 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 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}