nautilus_persistence/backend/parquet/catalog/
write.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 write paths.
17
18#![expect(
19    clippy::missing_errors_doc,
20    reason = "catalog write functions validate catalog-controlled batches and forward storage errors"
21)]
22
23use nautilus_serialization::arrow::catalog_identifier_from_metadata;
24
25use super::{
26    BTreeMap, CustomData, Data, DataBatch, EncodeToRecordBatch, HasCatalogDataType, HasTsInit,
27    Instrument, InstrumentAny, NautilusDataType, NautilusRecordType, ObjectPath, ObjectStoreExt,
28    Params, ParquetDataCatalog, PathBuf, RecordBatch, Serialize, UnixNanos,
29    WRITE_SKIP_DISJOINT_CHECK, are_intervals_disjoint, instrument_any_type, instrument_path_prefix,
30    local_to_object_store_path, parquet_data_path_prefix, prepare_custom_data_batch,
31    record_batch_without_identifier_column, record_path_prefix, timestamps_to_filename,
32    to_snake_case, write_batches_to_object_store, write_catalog_batch,
33};
34use crate::{
35    backend::parquet::{io::write_batches_to_object_store_create, paths::catalog_filename},
36    common::metadata::record_batch_ts_init_range,
37};
38
39impl ParquetDataCatalog {
40    /// Writes mixed data types to the catalog by separating them into type-specific collections.
41    ///
42    /// This method takes a heterogeneous collection of market data and separates it by type,
43    /// then writes each type to its appropriate location in the catalog. This is useful when
44    /// processing mixed data streams or bulk data imports.
45    ///
46    /// # Parameters
47    ///
48    /// - `data`: A vector of mixed [`Data`] enum variants.
49    /// - `start`: Optional start timestamp to override the data's natural range.
50    /// - `end`: Optional end timestamp to override the data's natural range.
51    ///
52    /// # Notes
53    ///
54    /// - Data is automatically sorted by type before writing.
55    /// - Each data type is written to its own directory structure.
56    /// - Instrument data handling is not yet implemented (TODO).
57    ///
58    /// # Examples
59    ///
60    /// ```rust,no_run
61    /// use nautilus_model::data::Data;
62    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
63    ///
64    /// let mut catalog = ParquetDataCatalog::new(
65    ///     std::path::Path::new("/tmp/nautilus_data"),
66    ///     None,
67    ///     None,
68    ///     None,
69    ///     None,
70    /// );
71    /// let mixed_data: Vec<Data> = vec![/* mixed data types */];
72    ///
73    /// catalog.write_data_enum(&mixed_data, None, None, None)?;
74    /// # Ok::<(), anyhow::Error>(())
75    /// ```
76    pub fn write_data_enum(
77        &self,
78        data: &[Data],
79        start: Option<UnixNanos>,
80        end: Option<UnixNanos>,
81        skip_disjoint_check: Option<bool>,
82    ) -> anyhow::Result<()> {
83        for batch in DataBatch::from_data_vec_grouped(data)? {
84            write_catalog_batch(self, &batch, start, end, skip_disjoint_check)?;
85        }
86
87        Ok(())
88    }
89
90    pub(super) fn write_grouped_to_parquet<T>(
91        &self,
92        data: &[T],
93        start: Option<UnixNanos>,
94        end: Option<UnixNanos>,
95        skip_disjoint_check: Option<bool>,
96    ) -> anyhow::Result<()>
97    where
98        T: Clone + HasTsInit + EncodeToRecordBatch + HasCatalogDataType,
99    {
100        let mut groups: BTreeMap<Option<String>, Vec<T>> = BTreeMap::new();
101
102        for item in data {
103            let identifier = catalog_identifier_from_metadata(&item.metadata());
104            groups.entry(identifier).or_default().push(item.clone());
105        }
106
107        for items in groups.into_values() {
108            self.write_to_parquet(&items, start, end, skip_disjoint_check)?;
109        }
110
111        Ok(())
112    }
113
114    /// Writes Arrow batches into catalog under a record type optional identifier.
115    ///
116    /// # Errors
117    ///
118    /// Returns an error if batches do not contain `ts_init` or cannot be persisted.
119    pub fn write_record_batches(
120        &mut self,
121        record_type: &NautilusRecordType,
122        identifier: Option<&str>,
123        batches: &[RecordBatch],
124        params: &Params,
125    ) -> anyhow::Result<()> {
126        if batches.is_empty() || batches.iter().all(|batch| batch.num_rows() == 0) {
127            return Ok(());
128        }
129
130        let (start_ts, end_ts) = record_batch_ts_init_range(batches)?;
131        let record_prefix = record_path_prefix(record_type);
132        let directory = self.make_path(record_prefix.as_ref(), identifier)?;
133        let filename = timestamps_to_filename(UnixNanos::from(start_ts), UnixNanos::from(end_ts));
134        let path = PathBuf::from(directory.clone()).join(&filename);
135        let object_path = self.to_object_path(&path.to_string_lossy())?;
136        let skip_disjoint_check = params.get_bool(WRITE_SKIP_DISJOINT_CHECK).unwrap_or(false);
137
138        if !skip_disjoint_check {
139            let current_intervals = self.get_directory_intervals(&directory)?;
140            let mut intervals = current_intervals.clone();
141            intervals.push((start_ts, end_ts));
142            anyhow::ensure!(
143                are_intervals_disjoint(&intervals),
144                "Writing file {filename} interval ({start_ts}, {end_ts}) would create non-disjoint intervals. Existing intervals: {current_intervals:?}",
145            );
146        }
147
148        self.execute_async(|| async {
149            write_batches_to_object_store(
150                batches,
151                self.object_store.clone(),
152                &object_path,
153                Some(self.compression),
154                Some(self.max_row_group_size),
155                None,
156            )
157            .await
158        })
159    }
160
161    /// Writes typed data to a Parquet file in the catalog.
162    ///
163    /// This is the core method for persisting market data to the catalog. It handles data
164    /// validation, batching, compression, and ensures proper file organization with
165    /// timestamp-based naming.
166    ///
167    /// # Type Parameters
168    ///
169    /// - `T`: The data type to write, must implement required traits for serialization and cataloging.
170    ///
171    /// # Parameters
172    ///
173    /// - `data`: Vector of data records to write (must be in ascending timestamp order).
174    /// - `start`: Optional start timestamp to override the natural data range.
175    /// - `end`: Optional end timestamp to override the natural data range.
176    ///
177    /// # Returns
178    ///
179    /// Returns the [`PathBuf`] of the created file, or an empty path if no data was provided.
180    /// If the target file already exists, returns the path without writing (skips write).
181    ///
182    /// # Errors
183    ///
184    /// Returns an error if:
185    /// - Data serialization to Arrow record batches fails.
186    /// - Object store write operations fail.
187    /// - File path construction fails.
188    /// - Writing would create non-disjoint timestamp intervals.
189    ///
190    /// # Panics
191    ///
192    /// Panics if:
193    /// - Data timestamps are not in ascending order.
194    /// - Record batches are empty after conversion.
195    /// - Required metadata is missing from the schema.
196    ///
197    /// # Examples
198    ///
199    /// ```rust,no_run
200    /// use nautilus_model::data::QuoteTick;
201    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
202    ///
203    /// let mut catalog = ParquetDataCatalog::new(
204    ///     std::path::Path::new("/tmp/nautilus_data"),
205    ///     None,
206    ///     None,
207    ///     None,
208    ///     None,
209    /// );
210    /// let quotes: Vec<QuoteTick> = vec![/* quote data */];
211    ///
212    /// let path = catalog.write_to_parquet(&quotes, None, None, None)?;
213    /// println!("Data written to: {:?}", path);
214    /// # Ok::<(), anyhow::Error>(())
215    /// ```
216    pub fn write_to_parquet<T>(
217        &self,
218        data: &[T],
219        start: Option<UnixNanos>,
220        end: Option<UnixNanos>,
221        skip_disjoint_check: Option<bool>,
222    ) -> anyhow::Result<PathBuf>
223    where
224        T: HasTsInit + EncodeToRecordBatch + HasCatalogDataType,
225    {
226        let (Some(first), Some(last)) = (data.first(), data.last()) else {
227            return Ok(PathBuf::new());
228        };
229
230        let type_name = to_snake_case(std::any::type_name::<T>());
231        Self::check_ascending_timestamps(data, &type_name)?;
232
233        let chunk_metadata = T::chunk_metadata(data);
234        if let Some(position) = data
235            .iter()
236            .position(|item| !item.matches_chunk_metadata(&chunk_metadata))
237        {
238            anyhow::bail!(
239                "Cannot write {type_name} data with mixed identities: element {position} has \
240                 metadata {:?} but the chunk has {chunk_metadata:?}; write each \
241                 instrument or bar type separately",
242                data[position].metadata(),
243            );
244        }
245
246        let start_ts = start.unwrap_or(first.ts_init());
247        let end_ts = end.unwrap_or(last.ts_init());
248
249        let batches = self.data_to_record_batches(data)?;
250        let schema = batches.first().expect("Batches are empty.").schema();
251
252        let data_type = T::catalog_data_type();
253        let path_prefix = parquet_data_path_prefix(&data_type);
254
255        let identifier = if matches!(data_type, NautilusDataType::Bar) {
256            schema.metadata.get("bar_type").cloned()
257        } else {
258            schema.metadata.get("instrument_id").cloned()
259        };
260
261        let directory = self.make_path(path_prefix.as_ref(), identifier.as_deref())?;
262        self.write_parquet_file_checked(
263            &directory,
264            start_ts,
265            end_ts,
266            &batches,
267            skip_disjoint_check.unwrap_or(false),
268            "File",
269            Some(&format!("{type_name} data")),
270            None,
271        )
272    }
273
274    /// Writes custom data to a Parquet file in the catalog.
275    ///
276    /// This method handles writing custom data types that implement `CustomDataTrait`.
277    /// Custom data is organized by type name in a `custom/{type_name}/` directory structure.
278    ///
279    /// # Parameters
280    ///
281    /// - `data`: Vector of custom data items to write (must be in ascending timestamp order).
282    /// - `start`: Optional start timestamp to override the natural data range.
283    /// - `end`: Optional end timestamp to override the natural data range.
284    /// - `skip_disjoint_check`: Whether to skip interval disjointness validation.
285    ///
286    /// # Returns
287    ///
288    /// Returns the [`PathBuf`] of the created file, or an empty path if no data was provided.
289    ///
290    /// # Errors
291    ///
292    /// Returns an error if:
293    /// - The registered Arrow schema omits `ts_init` or uses incompatible timestamp types.
294    /// - Data serialization to Arrow record batches fails.
295    /// - Object store write operations fail.
296    /// - File path construction fails.
297    /// - Writing would create non-disjoint timestamp intervals (unless skipped).
298    pub fn write_custom_data_batch<D>(
299        &self,
300        data: D,
301        start: Option<UnixNanos>,
302        end: Option<UnixNanos>,
303        skip_disjoint_check: Option<bool>,
304    ) -> anyhow::Result<PathBuf>
305    where
306        D: AsRef<[CustomData]>,
307    {
308        let data = data.as_ref().iter().collect::<Vec<_>>();
309        self.write_custom_data_refs_batch(&data, start, end, skip_disjoint_check)
310    }
311
312    pub(crate) fn write_custom_data_refs_batch(
313        &self,
314        data: &[&CustomData],
315        start: Option<UnixNanos>,
316        end: Option<UnixNanos>,
317        skip_disjoint_check: Option<bool>,
318    ) -> anyhow::Result<PathBuf> {
319        if data.is_empty() {
320            return Ok(PathBuf::new());
321        }
322
323        let (batch, type_name, identifier, start_ts, end_ts) = prepare_custom_data_batch(data)?;
324        let start_ts = start.unwrap_or(start_ts);
325        let end_ts = end.unwrap_or(end_ts);
326        let batches = vec![record_batch_without_identifier_column(batch)?];
327
328        let directory = self.make_path_custom_data(&type_name, identifier.as_deref())?;
329        self.write_parquet_file_checked(
330            &directory,
331            start_ts,
332            end_ts,
333            &batches,
334            skip_disjoint_check.unwrap_or(false),
335            "File",
336            None,
337            None,
338        )
339    }
340
341    /// Writes instruments to Parquet files in the catalog.
342    ///
343    /// Instruments are stored under their instrument ID directory using timestamp-ranged
344    /// file names, allowing multiple historical versions of the same instrument to be
345    /// appended over time:
346    /// `data/instruments/{instrument_id}/{start_ts}-{end_ts}.parquet`
347    ///
348    /// # Parameters
349    ///
350    /// - `instruments`: Vector of instruments to write.
351    ///
352    /// # Returns
353    ///
354    /// Returns a vector of paths to the created files.
355    ///
356    /// # Errors
357    ///
358    /// Returns an error if:
359    /// - Data serialization fails.
360    /// - Object store write operations fail.
361    /// - File path construction fails.
362    ///
363    /// # Examples
364    ///
365    /// ```rust,no_run
366    /// use nautilus_model::instruments::InstrumentAny;
367    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
368    ///
369    /// let mut catalog = ParquetDataCatalog::new(
370    ///     std::path::Path::new("/tmp/nautilus_data"),
371    ///     None,
372    ///     None,
373    ///     None,
374    ///     None,
375    /// );
376    /// let instruments: Vec<InstrumentAny> = vec![/* instruments */];
377    ///
378    /// let paths = catalog.write_instruments(instruments)?;
379    /// # Ok::<(), anyhow::Error>(())
380    /// ```
381    pub fn write_instruments(
382        &self,
383        instruments: Vec<InstrumentAny>,
384    ) -> anyhow::Result<Vec<PathBuf>> {
385        if instruments.is_empty() {
386            return Ok(Vec::new());
387        }
388
389        // Group instruments by concrete type and instrument_id so mixed InstrumentAny
390        // inputs are written as separate parquet batches with stable ordering.
391        let mut by_type_and_id: BTreeMap<(String, String), Vec<InstrumentAny>> = BTreeMap::new();
392
393        for instrument in instruments {
394            let instrument_type = instrument_any_type(&instrument);
395            let instrument_prefix = instrument_path_prefix(&instrument_type).to_string();
396            let instrument_id = Instrument::id(&instrument).to_string();
397            by_type_and_id
398                .entry((instrument_prefix, instrument_id))
399                .or_default()
400                .push(instrument);
401        }
402
403        let mut paths = Vec::new();
404
405        for ((instrument_prefix, instrument_id), instrument_group) in by_type_and_id {
406            Self::check_ascending_timestamps(&instrument_group, "instrument")?;
407
408            let (Some(first), Some(last)) = (instrument_group.first(), instrument_group.last())
409            else {
410                continue;
411            };
412
413            let start_ts = HasTsInit::ts_init(first);
414            let end_ts = HasTsInit::ts_init(last);
415            let batches = self.data_to_record_batches(&instrument_group)?;
416
417            let directory = self.make_path(&instrument_prefix, Some(instrument_id.as_str()))?;
418
419            // ArrowWriter stores the full schema (including `type_name` metadata) in ARROW:schema.
420            // When reading, use the builder's schema for metadata (see query_instruments).
421            let path = self.write_parquet_file_checked(
422                &directory,
423                start_ts,
424                end_ts,
425                &batches,
426                false,
427                "Instrument file",
428                Some(&format!("instrument data for {instrument_id}")),
429                None,
430            )?;
431
432            paths.push(path);
433        }
434
435        Ok(paths)
436    }
437
438    /// Writes `batches` to a timestamp-named parquet file in `directory`, skipping when the
439    /// target file already exists and validating interval disjointness unless `skip_disjoint_check`.
440    ///
441    /// `file_label` and `data_description` parameterize the log messages so callers keep their
442    /// site-specific wording (`data_description = None` suppresses the pre-write log).
443    #[expect(clippy::too_many_arguments)]
444    pub(crate) fn write_parquet_file_checked(
445        &self,
446        directory: &str,
447        start_ts: UnixNanos,
448        end_ts: UnixNanos,
449        batches: &[RecordBatch],
450        skip_disjoint_check: bool,
451        file_label: &str,
452        data_description: Option<&str>,
453        replay_identity: Option<&str>,
454    ) -> anyhow::Result<PathBuf> {
455        let filename = catalog_filename(start_ts, end_ts, replay_identity);
456
457        let path = PathBuf::from(directory).join(&filename);
458        let object_path = self.to_object_path(&path.to_string_lossy())?;
459
460        if self.file_exists(&path.to_string_lossy())? {
461            log::info!(
462                "{file_label} {} already exists, skipping write",
463                path.display()
464            );
465            return Ok(path);
466        }
467
468        if !skip_disjoint_check {
469            let current_intervals = self.get_directory_intervals(directory)?;
470            let new_interval = (start_ts.as_u64(), end_ts.as_u64());
471            let mut new_intervals = current_intervals.clone();
472            new_intervals.push(new_interval);
473
474            anyhow::ensure!(
475                are_intervals_disjoint(&new_intervals),
476                "Writing file {filename} with interval ({start_ts}, {end_ts}) would create \
477                non-disjoint intervals. Existing intervals: {current_intervals:?}"
478            );
479        }
480
481        if let Some(data_description) = data_description {
482            log::info!(
483                "Writing {} batches of {data_description} to {}",
484                batches.len(),
485                path.display(),
486            );
487        }
488
489        self.execute_async(|| async {
490            if replay_identity.is_none() {
491                return write_batches_to_object_store(
492                    batches,
493                    self.object_store.clone(),
494                    &object_path,
495                    Some(self.compression),
496                    Some(self.max_row_group_size),
497                    None,
498                )
499                .await;
500            }
501
502            let result = write_batches_to_object_store_create(
503                batches,
504                self.object_store.clone(),
505                &object_path,
506                Some(self.compression),
507                Some(self.max_row_group_size),
508                None,
509            )
510            .await;
511
512            match result {
513                Err(e)
514                    if matches!(
515                        e.downcast_ref::<object_store::Error>(),
516                        Some(object_store::Error::AlreadyExists { .. })
517                    ) =>
518                {
519                    Ok(())
520                }
521                result => result,
522            }
523        })?;
524
525        Ok(path)
526    }
527
528    /// Writes typed data to a JSON file in the catalog.
529    ///
530    /// This method provides an alternative to Parquet format for data export and debugging.
531    /// JSON files are human-readable but less efficient for large datasets.
532    ///
533    /// # Type Parameters
534    ///
535    /// - `T`: The data type to write, must implement serialization and cataloging traits.
536    ///
537    /// # Parameters
538    ///
539    /// - `data`: Vector of data records to write (must be in ascending timestamp order).
540    /// - `path`: Optional custom directory path (defaults to catalog's standard structure).
541    /// - `write_metadata`: Whether to write a separate metadata file alongside the data.
542    ///
543    /// # Returns
544    ///
545    /// Returns the [`PathBuf`] of the created JSON file.
546    ///
547    /// # Errors
548    ///
549    /// Returns an error if:
550    /// - JSON serialization fails.
551    /// - Object store write operations fail.
552    /// - File path construction fails.
553    ///
554    /// # Examples
555    ///
556    /// ```rust,no_run
557    /// use std::path::PathBuf;
558    /// use nautilus_model::data::TradeTick;
559    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
560    ///
561    /// let mut catalog = ParquetDataCatalog::new(
562    ///     std::path::Path::new("/tmp/nautilus_data"),
563    ///     None,
564    ///     None,
565    ///     None,
566    ///     None,
567    /// );
568    /// let trades: Vec<TradeTick> = vec![/* trade data */];
569    ///
570    /// let path = catalog.write_to_json(
571    ///     trades,
572    ///     Some(PathBuf::from("/custom/path")),
573    ///     true  // write metadata
574    /// )?;
575    /// # Ok::<(), anyhow::Error>(())
576    /// ```
577    pub fn write_to_json<T>(
578        &self,
579        data: Vec<T>,
580        path: Option<PathBuf>,
581        write_metadata: bool,
582    ) -> anyhow::Result<PathBuf>
583    where
584        T: HasTsInit + Serialize + HasCatalogDataType + EncodeToRecordBatch,
585    {
586        let (Some(first), Some(last)) = (data.first(), data.last()) else {
587            return Ok(PathBuf::new());
588        };
589
590        let type_name = to_snake_case(std::any::type_name::<T>());
591        Self::check_ascending_timestamps(&data, &type_name)?;
592
593        let start_ts = first.ts_init();
594        let end_ts = last.ts_init();
595
596        let data_type = T::catalog_data_type();
597        let path_prefix = parquet_data_path_prefix(&data_type);
598
599        let directory = match path {
600            Some(path) => path,
601            None => PathBuf::from(self.make_path(path_prefix.as_ref(), None)?),
602        };
603
604        let filename = timestamps_to_filename(start_ts, end_ts).replace(".parquet", ".json");
605        let json_path = directory.join(&filename);
606
607        log::info!(
608            "Writing {} records of {type_name} data to {}",
609            data.len(),
610            json_path.display(),
611        );
612
613        if write_metadata {
614            let metadata = T::chunk_metadata(&data);
615            let metadata_path = json_path.with_extension("metadata.json");
616            log::info!("Writing metadata to {}", metadata_path.display());
617
618            // Use object store for metadata file
619            let metadata_object_path = ObjectPath::from(local_to_object_store_path(&metadata_path));
620            let metadata_json = serde_json::to_vec_pretty(&metadata)?;
621            self.execute_async(|| async {
622                let _: object_store::PutResult = self
623                    .object_store
624                    .put(&metadata_object_path, metadata_json.into())
625                    .await?;
626                Ok(())
627            })?;
628        }
629
630        // Use object store for main JSON file
631        let json_object_path = ObjectPath::from(local_to_object_store_path(&json_path));
632        let json_data = serde_json::to_vec_pretty(&serde_json::to_value(data)?)?;
633        self.execute_async(|| async {
634            let _: object_store::PutResult = self
635                .object_store
636                .put(&json_object_path, json_data.into())
637                .await?;
638            Ok(())
639        })?;
640
641        Ok(json_path)
642    }
643
644    /// Validates that data timestamps are in ascending order.
645    ///
646    /// # Parameters
647    ///
648    /// - `data`: Slice of data records to validate.
649    /// - `type_name`: Name of the data type for error messages.
650    pub fn check_ascending_timestamps<T: HasTsInit>(
651        data: &[T],
652        type_name: &str,
653    ) -> anyhow::Result<()> {
654        anyhow::ensure!(
655            data.array_windows()
656                .all(|[a, b]| a.ts_init() <= b.ts_init()),
657            "{type_name} timestamps must be in ascending order"
658        );
659        Ok(())
660    }
661
662    /// Converts data into Arrow record batches for Parquet serialization.
663    ///
664    /// This method chunks the data according to the configured batch size and converts
665    /// each chunk into an Arrow record batch with appropriate metadata.
666    ///
667    /// # Type Parameters
668    ///
669    /// - `T`: The data type to convert, must implement required encoding traits.
670    ///
671    /// # Parameters
672    ///
673    /// - `data`: Vector of data records to convert.
674    ///
675    /// # Returns
676    ///
677    /// Returns a vector of Arrow [`RecordBatch`] instances ready for Parquet serialization.
678    ///
679    /// # Errors
680    ///
681    /// Returns an error if record batch encoding fails for any chunk.
682    pub fn data_to_record_batches<T>(&self, data: &[T]) -> anyhow::Result<Vec<RecordBatch>>
683    where
684        T: HasTsInit + EncodeToRecordBatch,
685    {
686        if data.is_empty() {
687            return Ok(Vec::new());
688        }
689
690        let mut batches = Vec::new();
691        let metadata = EncodeToRecordBatch::chunk_metadata(data);
692
693        for chunk in data.chunks(self.batch_size) {
694            let record_batch = T::encode_batch(&metadata, chunk)?;
695            let record_batch = record_batch_without_identifier_column(record_batch)?;
696            batches.push(record_batch);
697        }
698
699        Ok(batches)
700    }
701}
702
703#[cfg(test)]
704mod tests {
705    use std::{
706        fmt::Display,
707        fs::File,
708        sync::{
709            Arc,
710            atomic::{AtomicBool, AtomicUsize, Ordering},
711        },
712    };
713
714    use arrow::{
715        array::Int64Array,
716        datatypes::{DataType, Field, Schema},
717        record_batch::RecordBatch,
718    };
719    use futures::{StreamExt, TryStreamExt, stream::BoxStream};
720    use nautilus_core::UnixNanos;
721    use nautilus_model::{
722        data::{
723            NautilusDataType, OrderBookDelta, OrderBookDepth, QuoteTick,
724            stubs::{stub_delta, stub_depth10},
725        },
726        identifiers::InstrumentId,
727        types::{Price, Quantity},
728    };
729    use nautilus_serialization::arrow::{
730        DecodeFromRecordBatch, KEY_PRICE_PRECISION, KEY_SIZE_PRECISION,
731    };
732    use object_store::{
733        CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
734        ObjectStoreExt, PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult,
735        Result as ObjectStoreResult, memory::InMemory, path::Path as ObjectPath,
736    };
737    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
738    use rstest::rstest;
739    use tempfile::TempDir;
740
741    use super::ParquetDataCatalog;
742    use crate::common::datafusion::DataBackendSession;
743
744    #[rstest]
745    fn depth_write_shares_file_metadata_across_chunks(stub_depth10: OrderBookDepth) {
746        let mut empty = stub_depth10.clone();
747        empty.bids.clear();
748        empty.asks.clear();
749        empty.bid_counts.clear();
750        empty.ask_counts.clear();
751        let directory = TempDir::new().unwrap();
752        let catalog = ParquetDataCatalog::from_uri(
753            directory.path().to_str().unwrap(),
754            None,
755            Some(2),
756            None,
757            Some(2),
758        )
759        .unwrap();
760        let data = vec![empty.clone(), empty, stub_depth10];
761
762        let path = catalog.write_to_parquet(&data, None, None, None).unwrap();
763        let builder = ParquetRecordBatchReaderBuilder::try_new(
764            File::open(directory.path().join(path)).unwrap(),
765        )
766        .unwrap();
767        let metadata = builder.schema().metadata().clone();
768        let batches = builder
769            .build()
770            .unwrap()
771            .collect::<Result<Vec<_>, _>>()
772            .unwrap();
773        let decoded = batches
774            .iter()
775            .cloned()
776            .flat_map(|batch| OrderBookDepth::decode_batch(&metadata, batch).unwrap())
777            .collect::<Vec<_>>();
778
779        assert_eq!(metadata[KEY_PRICE_PRECISION], "2");
780        assert_eq!(metadata[KEY_SIZE_PRECISION], "0");
781        assert_eq!(decoded.len(), 3);
782        assert_eq!(decoded, data);
783        assert_eq!(decoded[2].bids[0].price.precision, 2);
784    }
785
786    #[rstest]
787    fn leading_clear_delta_writes_with_following_precision(stub_delta: OrderBookDelta) {
788        let directory = TempDir::new().unwrap();
789        let catalog = ParquetDataCatalog::from_uri(
790            directory.path().to_str().unwrap(),
791            None,
792            Some(2),
793            None,
794            None,
795        )
796        .unwrap();
797        let clear = OrderBookDelta::clear(
798            stub_delta.instrument_id,
799            0,
800            UnixNanos::from(1),
801            UnixNanos::from(1),
802        );
803        let second_clear = OrderBookDelta::clear(
804            stub_delta.instrument_id,
805            0,
806            UnixNanos::from(2),
807            UnixNanos::from(2),
808        );
809
810        let path = catalog
811            .write_to_parquet(&[clear, second_clear, stub_delta], None, None, None)
812            .unwrap();
813        let builder = ParquetRecordBatchReaderBuilder::try_new(
814            File::open(directory.path().join(&path)).unwrap(),
815        )
816        .unwrap();
817        let metadata = builder.schema().metadata().clone();
818        let decoded = builder
819            .build()
820            .unwrap()
821            .map(|batch| OrderBookDelta::decode_batch(&metadata, batch.unwrap()).unwrap())
822            .collect::<Vec<_>>()
823            .concat();
824
825        assert!(directory.path().join(path).exists());
826        assert_eq!(metadata[KEY_PRICE_PRECISION], "2");
827        assert_eq!(decoded[2].order.price.precision, 2);
828    }
829
830    #[derive(Debug)]
831    struct CreateRaceStore {
832        inner: InMemory,
833        create_calls: AtomicUsize,
834        already_exists: bool,
835        fail_head: AtomicBool,
836    }
837
838    impl Display for CreateRaceStore {
839        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
840            f.write_str("create-race")
841        }
842    }
843
844    #[async_trait::async_trait]
845    impl ObjectStore for CreateRaceStore {
846        async fn put_opts(
847            &self,
848            location: &ObjectPath,
849            payload: PutPayload,
850            opts: PutOptions,
851        ) -> ObjectStoreResult<PutResult> {
852            if opts.mode == PutMode::Create {
853                self.create_calls.fetch_add(1, Ordering::Relaxed);
854
855                if !self.already_exists {
856                    return Err(object_store::Error::Generic {
857                        store: "create-race",
858                        source: "injected create failure".into(),
859                    });
860                }
861
862                return Err(object_store::Error::AlreadyExists {
863                    path: location.to_string(),
864                    source: Box::new(std::io::Error::new(
865                        std::io::ErrorKind::AlreadyExists,
866                        "injected create race",
867                    )),
868                });
869            }
870
871            self.inner.put_opts(location, payload, opts).await
872        }
873
874        async fn put_multipart_opts(
875            &self,
876            location: &ObjectPath,
877            opts: PutMultipartOptions,
878        ) -> ObjectStoreResult<Box<dyn MultipartUpload>> {
879            self.inner.put_multipart_opts(location, opts).await
880        }
881
882        async fn get_opts(
883            &self,
884            location: &ObjectPath,
885            options: GetOptions,
886        ) -> ObjectStoreResult<GetResult> {
887            if options.head && self.fail_head.load(Ordering::Relaxed) {
888                return Err(object_store::Error::Generic {
889                    store: "create-race",
890                    source: "injected head failure".into(),
891                });
892            }
893
894            self.inner.get_opts(location, options).await
895        }
896
897        fn list(
898            &self,
899            prefix: Option<&ObjectPath>,
900        ) -> BoxStream<'static, ObjectStoreResult<ObjectMeta>> {
901            self.inner.list(prefix)
902        }
903
904        async fn list_with_delimiter(
905            &self,
906            prefix: Option<&ObjectPath>,
907        ) -> ObjectStoreResult<ListResult> {
908            self.inner.list_with_delimiter(prefix).await
909        }
910
911        fn delete_stream(
912            &self,
913            locations: BoxStream<'static, ObjectStoreResult<ObjectPath>>,
914        ) -> BoxStream<'static, ObjectStoreResult<ObjectPath>> {
915            self.inner.delete_stream(locations)
916        }
917
918        async fn copy_opts(
919            &self,
920            from: &ObjectPath,
921            to: &ObjectPath,
922            opts: CopyOptions,
923        ) -> ObjectStoreResult<()> {
924            self.inner.copy_opts(from, to, opts).await
925        }
926    }
927
928    #[rstest]
929    fn promotion_writes_to_memory_store_without_copy_support() {
930        let catalog = ParquetDataCatalog {
931            base_path: "catalog".to_string(),
932            original_uri: "memory://".to_string(),
933            object_store: Arc::new(InMemory::new()),
934            session: DataBackendSession::new(5_000),
935            batch_size: 5_000,
936            compression: parquet::basic::Compression::SNAPPY,
937            max_row_group_size: 5_000,
938        };
939
940        let batch = RecordBatch::try_new(
941            Arc::new(Schema::new(vec![Field::new(
942                "ts_init",
943                DataType::Int64,
944                false,
945            )])),
946            vec![Arc::new(Int64Array::from(vec![1]))],
947        )
948        .unwrap();
949
950        let path = catalog
951            .write_parquet_file_checked(
952                "quotes/TEST",
953                UnixNanos::from(1),
954                UnixNanos::from(1),
955                &[batch],
956                false,
957                "Promoted file",
958                None,
959                Some("memory-replay"),
960            )
961            .unwrap();
962        let object_path = catalog.to_object_path(&path.to_string_lossy()).unwrap();
963        catalog
964            .execute_async(|| async {
965                catalog.object_store.head(&object_path).await?;
966                Ok(())
967            })
968            .unwrap();
969    }
970
971    #[rstest]
972    fn promotion_accepts_already_exists_after_head_miss() {
973        let object_store = Arc::new(CreateRaceStore {
974            inner: InMemory::new(),
975            create_calls: AtomicUsize::new(0),
976            already_exists: true,
977            fail_head: AtomicBool::new(false),
978        });
979
980        let catalog = ParquetDataCatalog {
981            base_path: "catalog".to_string(),
982            original_uri: "memory://".to_string(),
983            object_store: object_store.clone(),
984            session: DataBackendSession::new(5_000),
985            batch_size: 5_000,
986            compression: parquet::basic::Compression::SNAPPY,
987            max_row_group_size: 5_000,
988        };
989
990        let batch = RecordBatch::try_new(
991            Arc::new(Schema::new(vec![Field::new(
992                "ts_init",
993                DataType::Int64,
994                false,
995            )])),
996            vec![Arc::new(Int64Array::from(vec![1]))],
997        )
998        .unwrap();
999
1000        catalog
1001            .write_parquet_file_checked(
1002                "quotes/TEST",
1003                UnixNanos::from(1),
1004                UnixNanos::from(1),
1005                &[batch],
1006                true,
1007                "Promoted file",
1008                None,
1009                Some("racing-replay"),
1010            )
1011            .unwrap();
1012
1013        assert_eq!(object_store.create_calls.load(Ordering::Relaxed), 1);
1014    }
1015
1016    #[rstest]
1017    fn promotion_propagates_create_errors_other_than_already_exists() {
1018        let object_store = Arc::new(CreateRaceStore {
1019            inner: InMemory::new(),
1020            create_calls: AtomicUsize::new(0),
1021            already_exists: false,
1022            fail_head: AtomicBool::new(false),
1023        });
1024
1025        let catalog = ParquetDataCatalog {
1026            base_path: "catalog".to_string(),
1027            original_uri: "memory://".to_string(),
1028            object_store: object_store.clone(),
1029            session: DataBackendSession::new(5_000),
1030            batch_size: 5_000,
1031            compression: parquet::basic::Compression::SNAPPY,
1032            max_row_group_size: 5_000,
1033        };
1034
1035        let batch = RecordBatch::try_new(
1036            Arc::new(Schema::new(vec![Field::new(
1037                "ts_init",
1038                DataType::Int64,
1039                false,
1040            )])),
1041            vec![Arc::new(Int64Array::from(vec![1]))],
1042        )
1043        .unwrap();
1044
1045        let error = catalog
1046            .write_parquet_file_checked(
1047                "quotes/TEST",
1048                UnixNanos::from(1),
1049                UnixNanos::from(1),
1050                &[batch],
1051                true,
1052                "Promoted file",
1053                None,
1054                Some("failing-replay"),
1055            )
1056            .unwrap_err();
1057
1058        assert!(
1059            format!("{error:#}").contains("injected create failure"),
1060            "unexpected error: {error:#}"
1061        );
1062        assert_eq!(object_store.create_calls.load(Ordering::Relaxed), 1);
1063    }
1064
1065    #[rstest]
1066    #[case::head_failure(true, None)]
1067    #[case::missing_file(false, Some(false))]
1068    fn file_exists_propagates_storage_errors_other_than_not_found(
1069        #[case] fail_head: bool,
1070        #[case] expected: Option<bool>,
1071    ) {
1072        let catalog = ParquetDataCatalog {
1073            base_path: "catalog".to_string(),
1074            original_uri: "memory://".to_string(),
1075            object_store: Arc::new(CreateRaceStore {
1076                inner: InMemory::new(),
1077                create_calls: AtomicUsize::new(0),
1078                already_exists: true,
1079                fail_head: AtomicBool::new(fail_head),
1080            }),
1081            session: DataBackendSession::new(5_000),
1082            batch_size: 5_000,
1083            compression: parquet::basic::Compression::SNAPPY,
1084            max_row_group_size: 5_000,
1085        };
1086
1087        let result = catalog.file_exists("quotes/TEST/missing.parquet");
1088
1089        match expected {
1090            Some(exists) => assert_eq!(result.unwrap(), exists),
1091            None => assert!(
1092                format!("{:#}", result.unwrap_err()).contains("injected head failure"),
1093                "head failure must propagate"
1094            ),
1095        }
1096    }
1097
1098    #[rstest]
1099    fn write_propagates_head_failure_instead_of_writing() {
1100        let object_store = Arc::new(CreateRaceStore {
1101            inner: InMemory::new(),
1102            create_calls: AtomicUsize::new(0),
1103            already_exists: true,
1104            fail_head: AtomicBool::new(true),
1105        });
1106
1107        let catalog = ParquetDataCatalog {
1108            base_path: "catalog".to_string(),
1109            original_uri: "memory://".to_string(),
1110            object_store: object_store.clone(),
1111            session: DataBackendSession::new(5_000),
1112            batch_size: 5_000,
1113            compression: parquet::basic::Compression::SNAPPY,
1114            max_row_group_size: 5_000,
1115        };
1116
1117        let batch = RecordBatch::try_new(
1118            Arc::new(Schema::new(vec![Field::new(
1119                "ts_init",
1120                DataType::Int64,
1121                false,
1122            )])),
1123            vec![Arc::new(Int64Array::from(vec![1]))],
1124        )
1125        .unwrap();
1126
1127        let error = catalog
1128            .write_parquet_file_checked(
1129                "quotes/TEST",
1130                UnixNanos::from(1),
1131                UnixNanos::from(1),
1132                &[batch],
1133                true,
1134                "Data file",
1135                None,
1136                None,
1137            )
1138            .unwrap_err();
1139
1140        let stored = catalog
1141            .execute_async(|| async {
1142                Ok(object_store
1143                    .inner
1144                    .list(None)
1145                    .collect::<Vec<_>>()
1146                    .await
1147                    .len())
1148            })
1149            .unwrap();
1150
1151        assert!(
1152            format!("{error:#}").contains("injected head failure"),
1153            "unexpected error: {error:#}"
1154        );
1155        assert_eq!(stored, 0);
1156    }
1157
1158    #[rstest]
1159    fn period_consolidation_stops_on_head_failure() {
1160        let object_store = Arc::new(CreateRaceStore {
1161            inner: InMemory::new(),
1162            create_calls: AtomicUsize::new(0),
1163            already_exists: true,
1164            fail_head: AtomicBool::new(false),
1165        });
1166
1167        let mut catalog = ParquetDataCatalog {
1168            base_path: String::new(),
1169            original_uri: "memory://".to_string(),
1170            object_store: object_store.clone(),
1171            session: DataBackendSession::new(5_000),
1172            batch_size: 5_000,
1173            compression: parquet::basic::Compression::SNAPPY,
1174            max_row_group_size: 5_000,
1175        };
1176
1177        let quote = |ts: u64| {
1178            QuoteTick::new(
1179                InstrumentId::from("ETH/USDT.BINANCE"),
1180                Price::from("1.0001"),
1181                Price::from("1.0002"),
1182                Quantity::from("100"),
1183                Quantity::from("100"),
1184                UnixNanos::from(ts),
1185                UnixNanos::from(ts),
1186            )
1187        };
1188
1189        catalog
1190            .write_to_parquet(&[quote(1)], None, None, None)
1191            .unwrap();
1192        catalog
1193            .write_to_parquet(&[quote(2)], None, None, None)
1194            .unwrap();
1195        object_store.fail_head.store(true, Ordering::Relaxed);
1196
1197        let error = catalog
1198            .consolidate_data_by_period(
1199                &NautilusDataType::QuoteTick.into(),
1200                Some("ETH/USDT.BINANCE"),
1201                Some(86_400_000_000_000),
1202                None,
1203                None,
1204                Some(true),
1205            )
1206            .unwrap_err();
1207
1208        let files = catalog
1209            .query_files(&NautilusDataType::QuoteTick.into(), None, None, None)
1210            .unwrap();
1211        assert!(
1212            format!("{error:#}").contains("injected head failure"),
1213            "unexpected error: {error:#}"
1214        );
1215        assert_eq!(files.len(), 2);
1216    }
1217
1218    #[rstest]
1219    fn write_to_json_normalizes_windows_separators() {
1220        let catalog = ParquetDataCatalog {
1221            base_path: String::new(),
1222            original_uri: "memory://".to_string(),
1223            object_store: Arc::new(InMemory::new()),
1224            session: DataBackendSession::new(5_000),
1225            batch_size: 5_000,
1226            compression: parquet::basic::Compression::SNAPPY,
1227            max_row_group_size: 5_000,
1228        };
1229
1230        let quote = QuoteTick::new(
1231            InstrumentId::from("ETH/USDT.BINANCE"),
1232            Price::from("1.0001"),
1233            Price::from("1.0002"),
1234            Quantity::from("100"),
1235            Quantity::from("100"),
1236            UnixNanos::from(1),
1237            UnixNanos::from(1),
1238        );
1239
1240        catalog
1241            .write_to_json(vec![quote], Some(r"data\quotes".into()), true)
1242            .unwrap();
1243
1244        let mut keys = catalog
1245            .execute_async(|| async {
1246                Ok(catalog
1247                    .object_store
1248                    .list(None)
1249                    .map_ok(|object| object.location.to_string())
1250                    .try_collect::<Vec<_>>()
1251                    .await?)
1252            })
1253            .unwrap();
1254
1255        keys.sort();
1256
1257        assert_eq!(
1258            keys,
1259            vec![
1260                "data/quotes/1970-01-01T00-00-00-000000001Z_1970-01-01T00-00-00-000000001Z.json",
1261                "data/quotes/1970-01-01T00-00-00-000000001Z_1970-01-01T00-00-00-000000001Z.metadata.json",
1262            ]
1263        );
1264    }
1265}
nautilus_persistence/backend/parquet/catalog/write.rs

nautilus_persistence/backend/parquet/catalog/
write.rs
