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
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}