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

nautilus_persistence/backend/parquet/catalog/
write.rs
