nautilus_persistence/backend/parquet/catalog/
mod.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 data catalog for efficient storage and retrieval of financial market data.
17//!
18//! This module provides a data catalog implementation that uses Apache Parquet
19//! format for storing financial market data with object store backends. The catalog supports
20//! various data types including quotes, trades, bars, order book data, and other market events.
21//!
22//! # Key Features
23//!
24//! - **Object Store Integration**: Works with local filesystems, S3, and other object stores.
25//! - **Data Type Support**: Handles all major financial data types (quotes, trades, bars, etc.).
26//! - **Time-based Organization**: Organizes data by timestamp ranges for efficient querying.
27//! - **Consolidation**: Merges multiple files to optimize storage and query performance.
28//! - **Validation**: Ensures data integrity with timestamp ordering and interval validation.
29//!
30//! # Architecture
31//!
32//! The catalog organizes data in a hierarchical structure:
33//! ```text
34//! data/
35//! ├── quotes/
36//! │   └── INSTRUMENT_ID/
37//! │       └── start_ts-end_ts.parquet
38//! ├── trades/
39//! │   └── INSTRUMENT_ID/
40//! │       └── start_ts-end_ts.parquet
41//! └── bars/
42//!     └── INSTRUMENT_ID/
43//!         └── start_ts-end_ts.parquet
44//! ```
45//!
46//! # Usage
47//!
48//! ```rust,no_run
49//! use std::path::Path;
50//!
51//! use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
52//!
53//! // Create a new catalog
54//! let catalog = ParquetDataCatalog::new(
55//!     Path::new("/path/to/data"),
56//!     None,       // storage_options
57//!     Some(5000), // batch_size
58//!     None,       // compression (defaults to ZSTD level 1)
59//!     None,       // max_row_group_size (defaults to 131,072)
60//! );
61//!
62//! // Write data to the catalog
63//! // catalog.write_to_parquet(&data, None, None, None)?;
64//! ```
65
66#![expect(
67    clippy::missing_fields_in_debug,
68    reason = "catalog Debug redacts internal caches"
69)]
70
71use std::{
72    collections::{BTreeMap, HashMap, HashSet},
73    fmt::Debug,
74    path::{Path, PathBuf},
75    sync::Arc,
76};
77
78use ahash::AHashMap;
79use arrow::record_batch::RecordBatch;
80use futures::StreamExt;
81use nautilus_core::{Params, UnixNanos, string::conversions::to_snake_case};
82use nautilus_model::{
83    data::{
84        Bar, CustomData, Data, DataBatch, FundingRateUpdate, HasTsInit, IndexPriceUpdate,
85        InstrumentStatus, MarkPriceUpdate, NautilusDataType, NautilusRecordType, OptionGreeks,
86        OrderBookDelta, OrderBookDepth, QuoteTick, TradeTick, close::InstrumentClose,
87        is_monotonically_increasing_by_init,
88    },
89    instruments::{Instrument, InstrumentAny},
90};
91use nautilus_serialization::arrow::{
92    ArrowSchemaProvider, DecodeDataFromRecordBatch, DecodeTypedFromRecordBatch,
93    EncodeToRecordBatch, catalog_display::catalog_record_batch_to_display,
94    custom::CustomDataDecoder, display::instrument::encode_instruments,
95    record_batch_without_identifier_column,
96};
97use object_store::{ObjectStore, ObjectStoreExt, path::Path as ObjectPath};
98use serde::Serialize;
99
100use crate::{
101    backend::parquet::{
102        intervals::query_interval_diff,
103        io::{
104            append_path_to_file_uri, decode_object_store_segment, is_remote_uri_scheme,
105            read_parquet_from_object_store, read_parquet_schema_from_object_store, remote_full_uri,
106            remote_store_root_url, write_batches_to_object_store,
107        },
108        paths::{extract_bar_type_instrument_id, query_intersects_filename},
109    },
110    catalog::{
111        session::{DEFAULT_DATA_BATCH_CHUNK_SIZE, DataBatchQueryResult, TypedDataBatchSession},
112        traits::{
113            CatalogInstrumentQuery, CatalogMetadata, CatalogQuery, CatalogReader,
114            CatalogRecordQuery, CatalogWriter, DataCatalog, filter_instrument_query_result,
115            filter_instruments_for_request_range,
116        },
117        types::{
118            CatalogAsOf, CatalogDataType, HasCatalogDataType, INSTRUMENT_PATH_PREFIXES,
119            instrument_any_type, instrument_path_prefix, parquet_data_path_prefix,
120            record_path_prefix,
121        },
122    },
123    common::{
124        custom::prepare_custom_data_batch,
125        datafusion::{self as datafusion, DataBackendSession, build_query},
126    },
127    config::validate_catalog_counts,
128};
129
130mod coverage;
131mod query;
132mod session;
133mod store;
134mod write;
135
136// Re-export public items from sibling modules so historical
137// `crate::backend::parquet::catalog::...` imports continue to resolve.
138pub use crate::backend::parquet::{
139    intervals::{are_intervals_contiguous, are_intervals_disjoint},
140    paths::{
141        CatalogPathPrefix, extract_identifier_from_path, extract_path_components,
142        extract_sql_safe_filename, local_to_object_store_path, make_local_path,
143        make_object_store_path, make_sql_safe_identifier, parse_filename_timestamps,
144        safe_directory_identifier, timestamps_to_filename, urisafe_instrument_id,
145    },
146};
147
148/// Optional Parquet query file-loading hint parameter key.
149pub const QUERY_OPTIMIZE_FILE_LOADING: &str = "optimize_file_loading";
150
151/// Optional Parquet write overlap bypass parameter key.
152pub const WRITE_SKIP_DISJOINT_CHECK: &str = "skip_disjoint_check";
153
154macro_rules! define_builtin_data_dispatch {
155    (
156        (Instrument, InstrumentAny, Instrument, Instrument, $instrument_prefix:literal),
157        $(($variant:ident, $type:ident, $data:ident, $batch:ident, $prefix:literal)),+ $(,)?
158    ) => {
159
160
161        fn query_builtin_batch(
162            catalog: &mut ParquetDataCatalog,
163            data_type: &NautilusDataType,
164            identifiers: Option<Vec<String>>,
165            start: Option<UnixNanos>,
166            end: Option<UnixNanos>,
167            where_clause: Option<&str>,
168            optimize_file_loading: bool,
169        ) -> Option<anyhow::Result<DataBatch>> {
170            match data_type {
171                $(
172                    NautilusDataType::$variant => Some(
173                        catalog
174                            .query::<$type>(
175                                identifiers,
176                                start,
177                                end,
178                                where_clause,
179                                None,
180                                optimize_file_loading,
181                            )
182                            .map(|data| DataBatch::$batch(data.into())),
183                    ),
184                )+
185                _ => None,
186            }
187        }
188
189        fn write_catalog_batch(
190            catalog: &ParquetDataCatalog,
191            batch: &DataBatch,
192            start: Option<UnixNanos>,
193            end: Option<UnixNanos>,
194            skip_disjoint_check: Option<bool>,
195        ) -> anyhow::Result<()> {
196            #[allow(unreachable_patterns, reason = "reject unsupported variants introduced by feature unification")]
197            match batch {
198                DataBatch::BookDeltas(data) => {
199                    let deltas = data.iter().flat_map(|batch| batch.deltas.iter().copied()).collect::<Vec<_>>();
200                    catalog.write_grouped_to_parquet(&deltas, start, end, skip_disjoint_check)
201                }
202                DataBatch::Custom(data) => catalog.write_custom_data_batch(data.as_ref(), start, end, skip_disjoint_check).map(|_| ()),
203                DataBatch::Instrument(data) => catalog.write_instruments(data.as_ref().to_vec()).map(|_| ()),
204                $(DataBatch::$batch(data) => catalog.write_grouped_to_parquet(data.as_ref(), start, end, skip_disjoint_check),)+
205                _ => anyhow::bail!("Unsupported catalog data batch: {}", batch.data_type_name()),
206            }
207        }
208    };
209}
210
211nautilus_model::for_each_data_type!(define_builtin_data_dispatch);
212
213/// A high-performance data catalog for storing and retrieving financial market data using Apache Parquet format.
214///
215/// The `ParquetDataCatalog` provides a solution for managing large volumes of financial
216/// market data with efficient storage, querying, and consolidation capabilities. It supports various
217/// object store backends including local filesystems, AWS S3, and other cloud storage providers.
218///
219/// # Features
220///
221/// - **Efficient Storage**: Uses Apache Parquet format with configurable compression.
222/// - **Object Store Backend**: Supports multiple storage backends through the `object_store` crate.
223/// - **Time-based Organization**: Organizes data by timestamp ranges for optimal query performance.
224/// - **Data Validation**: Ensures timestamp ordering and interval consistency.
225/// - **Consolidation**: Merges multiple files to reduce storage overhead and improve query speed.
226/// - **Type Safety**: Strongly typed data handling with compile-time guarantees.
227///
228/// # Data Organization
229///
230/// Data is organized hierarchically by data type and instrument:
231/// - `data/{data_type}/{instrument_id}/{start_ts}-{end_ts}.parquet`.
232/// - Files are named with their timestamp ranges for efficient range queries.
233/// - Intervals are validated to be disjoint to prevent data overlap.
234///
235/// # Performance Considerations
236///
237/// - **Batch Size**: Controls memory usage during data processing.
238/// - **Compression**: Affects file size and read and write speed (default: ZSTD level 1).
239/// - **Row Group Size**: Affects query performance and memory usage.
240/// - **File Consolidation**: Reduces the number of files for better query performance.
241pub struct ParquetDataCatalog {
242    /// The base path for data storage within the object store.
243    pub base_path: String,
244    /// The original URI provided when creating the catalog.
245    pub original_uri: String,
246    /// The object store backend for data persistence.
247    pub object_store: Arc<dyn ObjectStore>,
248    /// The DataFusion session for query execution.
249    pub session: DataBackendSession,
250    /// The number of records to process in each batch.
251    pub batch_size: usize,
252    /// The compression algorithm used for Parquet files.
253    pub compression: parquet::basic::Compression,
254    /// The maximum number of rows in each Parquet row group.
255    pub max_row_group_size: usize,
256}
257
258impl Debug for ParquetDataCatalog {
259    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
260        f.debug_struct(stringify!(ParquetDataCatalog))
261            .field("base_path", &self.base_path)
262            .finish()
263    }
264}
265
266impl ParquetDataCatalog {
267    /// Creates a new [`ParquetDataCatalog`] instance from a local file path.
268    ///
269    /// This is a convenience constructor that converts a local path to a URI format
270    /// and delegates to [`Self::from_uri`].
271    ///
272    /// # Parameters
273    ///
274    /// - `base_path`: The base directory path for data storage.
275    /// - `storage_options`: Optional `HashMap` containing storage-specific configuration options.
276    /// - `batch_size`: Number of records to process in each batch (default: 10,000).
277    /// - `compression`: Parquet compression algorithm (default: ZSTD level 1).
278    /// - `max_row_group_size`: Maximum rows per Parquet row group (default: 131,072).
279    ///
280    /// # Panics
281    ///
282    /// Panics if the path cannot be converted to a valid URI or if the object store
283    /// cannot be created from the path.
284    ///
285    /// # Examples
286    ///
287    /// ```rust,no_run
288    /// use std::path::Path;
289    ///
290    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
291    ///
292    /// let catalog = ParquetDataCatalog::new(
293    ///     Path::new("/tmp/nautilus_data"),
294    ///     None,       // no storage options
295    ///     Some(1000), // smaller batch size
296    ///     None,       // default compression
297    ///     None,       // default row group size
298    /// );
299    /// ```
300    #[must_use]
301    pub fn new(
302        base_path: &Path,
303        storage_options: Option<AHashMap<String, String>>,
304        batch_size: Option<usize>,
305        compression: Option<parquet::basic::Compression>,
306        max_row_group_size: Option<usize>,
307    ) -> Self {
308        let path_str = base_path.to_string_lossy().to_string();
309        Self::from_uri(
310            &path_str,
311            storage_options,
312            batch_size,
313            compression,
314            max_row_group_size,
315        )
316        .expect("Failed to create catalog from path")
317    }
318
319    /// Creates a new [`ParquetDataCatalog`] instance from a URI with optional storage options.
320    ///
321    /// Supports various URI schemes including local file paths and multiple cloud storage backends
322    /// supported by the `object_store` crate.
323    ///
324    /// # Supported URI Schemes
325    ///
326    /// - **AWS S3**: `s3://bucket/path`.
327    /// - **Google Cloud Storage**: `gs://bucket/path` or `gcs://bucket/path`.
328    /// - **Azure Blob Storage**: `az://container/path` or `abfs://container@account.dfs.core.windows.net/path`.
329    /// - **HTTP/WebDAV**: `http://` or `https://`.
330    /// - **Local files**: `file://path` or plain paths.
331    ///
332    /// # Parameters
333    ///
334    /// - `uri`: The URI for the data storage location.
335    /// - `storage_options`: Optional `HashMap` of `object_store` configuration keys for the scheme:
336    ///   - For S3: `endpoint_url`, region, `access_key_id`, `secret_access_key`, `session_token`, etc.
337    ///   - For GCS: `service_account_path`, `service_account_key`, `application_credentials`, etc.
338    ///   - For Azure: `account_name`, `account_key`, `sas_token`, etc.
339    /// - `batch_size`: Number of records to process in each batch (default: 10,000).
340    /// - `compression`: Parquet compression algorithm (default: ZSTD level 1).
341    /// - `max_row_group_size`: Maximum rows per Parquet row group (default: 131,072).
342    ///
343    /// # Errors
344    ///
345    /// Returns an error if:
346    /// - `batch_size` or `max_row_group_size` is zero.
347    /// - The URI format is invalid or unsupported.
348    /// - A storage option key is unknown for the URI scheme.
349    /// - The object store cannot be created or accessed.
350    /// - Authentication fails for cloud storage backends.
351    ///
352    /// # Examples
353    ///
354    /// ```rust,no_run
355    /// use ahash::AHashMap;
356    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
357    ///
358    /// // Local filesystem
359    /// let local_catalog = ParquetDataCatalog::from_uri("/tmp/nautilus_data", None, None, None, None)?;
360    ///
361    /// // S3 bucket
362    /// let s3_catalog =
363    ///     ParquetDataCatalog::from_uri("s3://my-bucket/nautilus-data", None, None, None, None)?;
364    ///
365    /// // Google Cloud Storage
366    /// let gcs_catalog =
367    ///     ParquetDataCatalog::from_uri("gs://my-bucket/nautilus-data", None, None, None, None)?;
368    ///
369    /// // Azure Blob Storage
370    /// let azure_catalog =
371    ///     ParquetDataCatalog::from_uri("az://container/nautilus-data", None, None, None, None)?;
372    ///
373    /// // S3 with custom endpoint and credentials
374    /// let mut storage_options = AHashMap::new();
375    /// storage_options.insert(
376    ///     "endpoint_url".to_string(),
377    ///     "https://my-s3-endpoint.com".to_string(),
378    /// );
379    /// storage_options.insert("access_key_id".to_string(), "my-key".to_string());
380    /// storage_options.insert("secret_access_key".to_string(), "my-secret".to_string());
381    ///
382    /// let custom_s3_catalog = ParquetDataCatalog::from_uri(
383    ///     "s3://my-bucket/nautilus-data",
384    ///     Some(storage_options),
385    ///     None,
386    ///     None,
387    ///     None,
388    /// )?;
389    /// # Ok::<(), anyhow::Error>(())
390    /// ```
391    pub fn from_uri(
392        uri: &str,
393        storage_options: Option<AHashMap<String, String>>,
394        batch_size: Option<usize>,
395        compression: Option<parquet::basic::Compression>,
396        max_row_group_size: Option<usize>,
397    ) -> anyhow::Result<Self> {
398        validate_catalog_counts(batch_size, max_row_group_size)?;
399
400        let batch_size = batch_size.unwrap_or(DEFAULT_DATA_BATCH_CHUNK_SIZE);
401        let compression = compression.unwrap_or(parquet::basic::Compression::ZSTD(
402            parquet::basic::ZstdLevel::default(),
403        ));
404        let max_row_group_size =
405            max_row_group_size.unwrap_or(crate::backend::parquet::DEFAULT_ROW_GROUP_SIZE);
406
407        let location = crate::backend::parquet::io::create_object_store_location_from_path(
408            uri,
409            storage_options,
410        )?;
411
412        Ok(Self {
413            base_path: location.base_path,
414            original_uri: location.original_uri,
415            object_store: location.object_store,
416            session: DataBackendSession::new(batch_size),
417            batch_size,
418            compression,
419            max_row_group_size,
420        })
421    }
422
423    /// Returns the base path of the catalog for testing purposes.
424    #[must_use]
425    pub fn get_base_path(&self) -> String {
426        self.base_path.clone()
427    }
428
429    /// Clears cached table registrations so a later query re-reads files that changed.
430    ///
431    /// Catalog operations that modify files call this before querying again.
432    pub(crate) fn clear_session_tables(&mut self) {
433        self.session.clear_registered_tables();
434    }
435}
436
437impl CatalogReader for ParquetDataCatalog {
438    fn fork_query_catalog(&self) -> anyhow::Result<Option<DataCatalog>> {
439        Ok(Some(Box::new(Self {
440            base_path: self.base_path.clone(),
441            original_uri: self.original_uri.clone(),
442            object_store: self.object_store.clone(),
443            session: DataBackendSession::new(self.batch_size),
444            batch_size: self.batch_size,
445            compression: self.compression,
446            max_row_group_size: self.max_row_group_size,
447        })))
448    }
449
450    fn query_batch_session(
451        &mut self,
452        query: &CatalogQuery,
453        chunk_size: Option<usize>,
454    ) -> anyhow::Result<DataBatchQueryResult> {
455        ensure_latest_query(query.as_of)?;
456        macro_rules! batch_session {
457            ((Instrument, InstrumentAny, Instrument, Instrument, $instrument_prefix:literal), $(($variant:ident, $type:ident, $data:ident, $batch:ident, $prefix:literal)),+ $(,)?) => {
458                match &query.data_type {
459                    $(NautilusDataType::$variant => {
460                        let pages = self.query_typed_pages::<$type>(
461                            query.identifiers.clone(), query.start, query.end, query.where_clause.as_deref(), None,
462                            optimize_file_loading(query.params.as_ref()),
463                        )?;
464                        Ok(Box::new(TypedDataBatchSession::new(pages, chunk_size)) as DataBatchQueryResult)
465                    },)+
466                    _ => match self.query_batch(query)? {
467                        DataBatch::Instrument(data) => Ok(Box::new(TypedDataBatchSession::from_vec(data.as_ref().to_vec(), chunk_size))),
468                        DataBatch::Custom(data) => Ok(Box::new(TypedDataBatchSession::from_vec(data.as_ref().to_vec(), chunk_size))),
469                        _ => anyhow::bail!("Unsupported Parquet query family"),
470                    },
471                }
472            };
473        }
474        nautilus_model::for_each_data_type!(batch_session)
475    }
476
477    fn reset_session(&mut self) {
478        self.clear_session_tables();
479    }
480
481    fn instruments(
482        &mut self,
483        query: &CatalogInstrumentQuery,
484    ) -> anyhow::Result<Vec<InstrumentAny>> {
485        let CatalogInstrumentQuery {
486            instrument_ids,
487            start,
488            end,
489            where_clause,
490            instrument_type,
491        } = query.clone();
492        let instrument_ids = instrument_ids.as_deref();
493        self.query_instruments_filtered_with_where_and_type(
494            instrument_ids,
495            start,
496            end,
497            where_clause.as_deref(),
498            instrument_type.as_ref(),
499        )
500    }
501
502    fn query_batch(&mut self, query: &CatalogQuery) -> anyhow::Result<DataBatch> {
503        ensure_latest_query(query.as_of)?;
504        let CatalogQuery {
505            data_type,
506            identifiers,
507            start,
508            end,
509            where_clause,
510            params,
511            instrument_type,
512            ..
513        } = query.clone();
514        let where_clause = where_clause.as_deref();
515        let optimize_file_loading = optimize_file_loading(params.as_ref());
516
517        match data_type {
518            NautilusDataType::Instrument => {
519                let data = self.query_instruments_filtered_with_where_and_type(
520                    identifiers.as_deref(),
521                    start,
522                    end,
523                    where_clause,
524                    instrument_type.as_ref(),
525                )?;
526                Ok(DataBatch::Instrument(
527                    filter_instrument_query_result(data, start, params.as_ref()).into(),
528                ))
529            }
530            NautilusDataType::Custom { type_name } => {
531                let data = self.query_custom_data_dynamic(
532                    &type_name,
533                    identifiers.as_deref(),
534                    start,
535                    end,
536                    where_clause,
537                    None,
538                    optimize_file_loading,
539                )?;
540                Ok(DataBatch::Custom(
541                    data.into_iter()
542                        .filter_map(|item| match item {
543                            Data::Custom(custom) => Some(custom),
544                            _ => None,
545                        })
546                        .collect::<Vec<_>>()
547                        .into(),
548                ))
549            }
550            #[cfg(feature = "defi")]
551            NautilusDataType::Defi => Err(anyhow::Error::from(
552                crate::errors::PersistenceError::unsupported("Parquet catalog DeFi data"),
553            )),
554            data_type => query_builtin_batch(
555                self,
556                &data_type,
557                identifiers,
558                start,
559                end,
560                where_clause,
561                optimize_file_loading,
562            )
563            .unwrap_or_else(|| {
564                Err(anyhow::Error::from(
565                    crate::errors::PersistenceError::unsupported(format!(
566                        "Parquet catalog {data_type} data"
567                    )),
568                ))
569            }),
570        }
571    }
572
573    fn query_identifiers(&mut self, query: &CatalogQuery) -> anyhow::Result<Vec<String>> {
574        ensure_latest_query(query.as_of)?;
575        let CatalogQuery {
576            data_type,
577            identifiers,
578            start,
579            end,
580            where_clause,
581            params,
582            instrument_type,
583            ..
584        } = query.clone();
585
586        if data_type == NautilusDataType::Instrument {
587            let mut identifiers = self
588                .query_instruments_filtered_with_where_and_type(
589                    identifiers.as_deref(),
590                    start,
591                    end,
592                    where_clause.as_deref(),
593                    instrument_type.as_ref(),
594                )?
595                .into_iter()
596                .map(|instrument| instrument.id().to_string())
597                .collect::<Vec<_>>();
598            identifiers.sort();
599            identifiers.dedup();
600            return Ok(identifiers);
601        }
602
603        Self::query_identifiers(
604            self,
605            &CatalogDataType::Data(data_type),
606            identifiers,
607            start,
608            end,
609            where_clause.as_deref(),
610            optimize_file_loading(params.as_ref()),
611        )
612    }
613
614    fn query_display_record_batches(
615        &mut self,
616        query: &CatalogQuery,
617    ) -> anyhow::Result<Vec<RecordBatch>> {
618        ensure_latest_query(query.as_of)?;
619        let CatalogQuery {
620            data_type,
621            identifiers,
622            start,
623            end,
624            where_clause,
625            params,
626            instrument_type,
627            ..
628        } = query.clone();
629
630        if data_type == NautilusDataType::Instrument {
631            let instruments = self.query_instruments_filtered_with_where_and_type(
632                identifiers.as_deref(),
633                start,
634                end,
635                where_clause.as_deref(),
636                instrument_type.as_ref(),
637            )?;
638            return if instruments.is_empty() {
639                Ok(Vec::new())
640            } else {
641                Ok(vec![encode_instruments(&instruments)?])
642            };
643        }
644
645        Self::query_display_record_batches(
646            self,
647            &data_type,
648            identifiers,
649            start,
650            end,
651            where_clause.as_deref(),
652            optimize_file_loading(params.as_ref()),
653        )
654    }
655
656    fn query_record_batches(
657        &mut self,
658        query: &CatalogRecordQuery,
659    ) -> anyhow::Result<Vec<RecordBatch>> {
660        ensure_latest_query(query.as_of)?;
661        let CatalogRecordQuery {
662            record_type,
663            identifier,
664            start,
665            end,
666            where_clause,
667            params,
668            ..
669        } = query.clone();
670        Self::query_record_batches(
671            self,
672            &CatalogDataType::Record(record_type),
673            identifier,
674            start,
675            end,
676            where_clause.as_deref(),
677            optimize_file_loading(params.as_ref()),
678        )
679    }
680
681    fn query_record_display_batches(
682        &mut self,
683        query: &CatalogRecordQuery,
684    ) -> anyhow::Result<Vec<RecordBatch>> {
685        CatalogReader::query_record_batches(self, query)
686    }
687
688    fn query_metadata(&mut self, query: &CatalogQuery) -> anyhow::Result<Vec<CatalogMetadata>> {
689        ensure_latest_query(query.as_of)?;
690        let CatalogQuery {
691            data_type,
692            identifiers,
693            start,
694            end,
695            where_clause,
696            instrument_type,
697            ..
698        } = query.clone();
699
700        let data_type = match (data_type, instrument_type) {
701            (NautilusDataType::Instrument, Some(instrument_type)) => {
702                CatalogDataType::Instrument(instrument_type)
703            }
704            (data_type, _) => CatalogDataType::Data(data_type),
705        };
706
707        Self::query_metadata(
708            self,
709            &data_type,
710            identifiers,
711            start,
712            end,
713            where_clause.as_deref(),
714        )
715    }
716
717    fn get_missing_intervals_for_request(
718        &mut self,
719        start: UnixNanos,
720        end: UnixNanos,
721        data_type: NautilusDataType,
722        identifier: Option<&str>,
723    ) -> anyhow::Result<Vec<(u64, u64)>> {
724        match data_type {
725            NautilusDataType::Instrument => {
726                let identifiers = identifier.map(|value| vec![value.to_string()]);
727                let data = self.query_instruments_filtered(
728                    identifiers.as_deref(),
729                    Some(start),
730                    Some(end),
731                )?;
732                Ok(if data.is_empty() {
733                    vec![(start.as_u64(), end.as_u64())]
734                } else {
735                    Vec::new()
736                })
737            }
738            NautilusDataType::Custom { type_name } => {
739                if let Some(identifier) = identifier {
740                    let directory = self.make_path_custom_data(&type_name, Some(identifier))?;
741                    let intervals = self.get_directory_intervals(&directory)?;
742
743                    Ok(query_interval_diff(
744                        start.as_u64(),
745                        end.as_u64(),
746                        &intervals,
747                    ))
748                } else {
749                    Self::get_missing_intervals_for_request(
750                        self,
751                        start.as_u64(),
752                        end.as_u64(),
753                        &CatalogDataType::Data(NautilusDataType::Custom { type_name }),
754                        None,
755                    )
756                }
757            }
758            _ => Self::get_missing_intervals_for_request(
759                self,
760                start.as_u64(),
761                end.as_u64(),
762                &CatalogDataType::Data(data_type),
763                identifier,
764            ),
765        }
766    }
767
768    fn query_last_timestamp(
769        &mut self,
770        data_type: NautilusDataType,
771        identifier: Option<&str>,
772    ) -> anyhow::Result<Option<u64>> {
773        match data_type {
774            NautilusDataType::Instrument => {
775                let identifiers = identifier.map(|value| vec![value.to_string()]);
776                Ok(self
777                    .query_instruments(identifiers.as_deref())?
778                    .into_iter()
779                    .map(|instrument| HasTsInit::ts_init(&instrument).as_u64())
780                    .max())
781            }
782            NautilusDataType::Custom { type_name } => {
783                if let Some(identifier) = identifier {
784                    let directory = self.make_path_custom_data(&type_name, Some(identifier))?;
785                    let intervals = self.get_directory_intervals(&directory)?;
786
787                    Ok(intervals.into_iter().map(|(_, end)| end).max())
788                } else {
789                    Self::query_last_timestamp(
790                        self,
791                        &CatalogDataType::Data(NautilusDataType::Custom { type_name }),
792                        None,
793                    )
794                }
795            }
796            _ => Self::query_last_timestamp(self, &CatalogDataType::Data(data_type), identifier),
797        }
798    }
799}
800
801impl CatalogWriter for ParquetDataCatalog {
802    fn write_instruments(&mut self, instruments: &[InstrumentAny]) -> anyhow::Result<()> {
803        Self::write_instruments(self, instruments.to_vec()).map(|_| ())
804    }
805
806    fn write_data(
807        &mut self,
808        data: &[Data],
809        start: Option<UnixNanos>,
810        end: Option<UnixNanos>,
811        params: Option<Params>,
812    ) -> anyhow::Result<()> {
813        if data.is_empty() {
814            return Ok(());
815        }
816
817        let skip_disjoint_check = params
818            .as_ref()
819            .and_then(|params| params.get_bool(WRITE_SKIP_DISJOINT_CHECK));
820        self.write_data_enum(data, start, end, skip_disjoint_check)
821    }
822
823    fn write_data_batch(
824        &mut self,
825        batch: &DataBatch,
826        start: Option<UnixNanos>,
827        end: Option<UnixNanos>,
828        params: Option<Params>,
829    ) -> anyhow::Result<()> {
830        let skip_disjoint_check = params
831            .as_ref()
832            .and_then(|params| params.get_bool(WRITE_SKIP_DISJOINT_CHECK));
833
834        write_catalog_batch(self, batch, start, end, skip_disjoint_check)
835    }
836
837    fn write_records(
838        &mut self,
839        record_type: NautilusRecordType,
840        batches: &[RecordBatch],
841        params: Option<Params>,
842    ) -> anyhow::Result<()> {
843        let params = params.unwrap_or_default();
844        self.write_record_batches(&record_type, params.get_str("identifier"), batches, &params)
845    }
846
847    fn record_empty_coverage(
848        &mut self,
849        data_type: NautilusDataType,
850        identifier: Option<&str>,
851        start: UnixNanos,
852        end: UnixNanos,
853    ) -> anyhow::Result<()> {
854        match &data_type {
855            NautilusDataType::Instrument => {
856                anyhow::bail!(
857                    "Cannot record empty instrument coverage without a concrete instrument type"
858                )
859            }
860            NautilusDataType::Custom { type_name } => {
861                let directory = self.make_path_custom_data(type_name, identifier)?;
862                self.extend_file_name_in_directory(&directory, start, end)
863            }
864            _ => self.extend_file_name(&CatalogDataType::Data(data_type), identifier, start, end),
865        }
866    }
867}
868
869fn ensure_latest_query(as_of: CatalogAsOf) -> anyhow::Result<()> {
870    anyhow::ensure!(
871        as_of == CatalogAsOf::Latest,
872        "Parquet catalog does not support historical queries"
873    );
874    Ok(())
875}
876
877fn optimize_file_loading(params: Option<&Params>) -> bool {
878    params
879        .and_then(|params| params.get_bool(QUERY_OPTIMIZE_FILE_LOADING))
880        .unwrap_or(true)
881}
nautilus_persistence/backend/parquet/catalog/mod.rs

nautilus_persistence/backend/parquet/catalog/
mod.rs
