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    clippy::missing_panics_doc,
21    reason = "catalog write functions validate catalog-controlled batches and forward storage errors"
22)]
23
24use nautilus_serialization::arrow::catalog_identifier_from_metadata;
25
26use super::{
27    BTreeMap, CustomData, Data, DataBatch, EncodeToRecordBatch, HasCatalogDataType, HasTsInit,
28    InstrumentAny, NautilusRecordType, ObjectPath, ObjectStoreExt, Params, ParquetDataCatalog,
29    PathBuf, RecordBatch, Serialize, UnixNanos, WRITE_SKIP_DISJOINT_CHECK, are_intervals_disjoint,
30    instrument_any_type, instrument_path_prefix, parquet_data_path_prefix,
31    prepare_custom_data_batch, record_batch_without_identifier_column, record_path_prefix,
32    timestamps_to_filename, to_snake_case, write_batches_to_object_store, write_catalog_batch,
33};
34use crate::{
35    backend::parquet::io::write_batches_to_object_store_create,
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        Ok(())
87    }
88
89    pub(super) fn write_grouped_to_parquet<T>(
90        &self,
91        data: &[T],
92        start: Option<UnixNanos>,
93        end: Option<UnixNanos>,
94        skip_disjoint_check: Option<bool>,
95    ) -> anyhow::Result<()>
96    where
97        T: Clone + HasTsInit + EncodeToRecordBatch + HasCatalogDataType,
98    {
99        let mut groups: BTreeMap<Option<String>, Vec<T>> = BTreeMap::new();
100
101        for item in data {
102            let identifier = catalog_identifier_from_metadata(&item.metadata());
103            groups.entry(identifier).or_default().push(item.clone());
104        }
105
106        for items in groups.into_values() {
107            self.write_to_parquet(&items, start, end, skip_disjoint_check)?;
108        }
109        Ok(())
110    }
111
112    /// Writes Arrow batches into catalog under a record type optional identifier.
113    ///
114    /// # Errors
115    ///
116    /// Returns error if batches do not contain `ts_init` or cannot be persisted.
117    pub fn write_record_batches(
118        &mut self,
119        record_type: &NautilusRecordType,
120        identifier: Option<&str>,
121        batches: &[RecordBatch],
122        params: &Params,
123    ) -> anyhow::Result<()> {
124        if batches.is_empty() || batches.iter().all(|batch| batch.num_rows() == 0) {
125            return Ok(());
126        }
127
128        let (start_ts, end_ts) = record_batch_ts_init_range(batches)?;
129        let record_prefix = record_path_prefix(record_type);
130        let directory = self.make_path(record_prefix.as_ref(), identifier)?;
131        let filename = timestamps_to_filename(UnixNanos::from(start_ts), UnixNanos::from(end_ts));
132        let path = PathBuf::from(directory.clone()).join(&filename);
133        let object_path = self.to_object_path(&path.to_string_lossy())?;
134        let skip_disjoint_check = params.get_bool(WRITE_SKIP_DISJOINT_CHECK).unwrap_or(false);
135
136        if !skip_disjoint_check {
137            let current_intervals = self.get_directory_intervals(&directory)?;
138            let mut intervals = current_intervals.clone();
139            intervals.push((start_ts, end_ts));
140            anyhow::ensure!(
141                are_intervals_disjoint(&intervals),
142                "Writing file {filename} interval ({start_ts}, {end_ts}) would create non-disjoint intervals. Existing intervals: {current_intervals:?}",
143            );
144        }
145
146        self.execute_async(|| async {
147            write_batches_to_object_store(
148                batches,
149                self.object_store.clone(),
150                &object_path,
151                Some(self.compression),
152                Some(self.max_row_group_size),
153                None,
154            )
155            .await
156        })
157    }
158
159    /// Writes typed data to a Parquet file in the catalog.
160    ///
161    /// This is the core method for persisting market data to the catalog. It handles data
162    /// validation, batching, compression, and ensures proper file organization with
163    /// timestamp-based naming.
164    ///
165    /// # Type Parameters
166    ///
167    /// - `T`: The data type to write, must implement required traits for serialization and cataloging.
168    ///
169    /// # Parameters
170    ///
171    /// - `data`: Vector of data records to write (must be in ascending timestamp order).
172    /// - `start`: Optional start timestamp to override the natural data range.
173    /// - `end`: Optional end timestamp to override the natural data range.
174    ///
175    /// # Returns
176    ///
177    /// Returns the [`PathBuf`] of the created file, or an empty path if no data was provided.
178    /// If the target file already exists, returns the path without writing (skips write).
179    ///
180    /// # Errors
181    ///
182    /// Returns an error if:
183    /// - Data serialization to Arrow record batches fails.
184    /// - Object store write operations fail.
185    /// - File path construction fails.
186    /// - Writing would create non-disjoint timestamp intervals.
187    ///
188    /// # Panics
189    ///
190    /// Panics if:
191    /// - Data timestamps are not in ascending order.
192    /// - Record batches are empty after conversion.
193    /// - Required metadata is missing from the schema.
194    ///
195    /// # Examples
196    ///
197    /// ```rust,no_run
198    /// use nautilus_model::data::QuoteTick;
199    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
200    ///
201    /// let mut catalog = ParquetDataCatalog::new(
202    ///     std::path::Path::new("/tmp/nautilus_data"),
203    ///     None,
204    ///     None,
205    ///     None,
206    ///     None,
207    /// );
208    /// let quotes: Vec<QuoteTick> = vec![/* quote data */];
209    ///
210    /// let path = catalog.write_to_parquet(&quotes, None, None, None)?;
211    /// println!("Data written to: {:?}", path);
212    /// # Ok::<(), anyhow::Error>(())
213    /// ```
214    pub fn write_to_parquet<T>(
215        &self,
216        data: &[T],
217        start: Option<UnixNanos>,
218        end: Option<UnixNanos>,
219        skip_disjoint_check: Option<bool>,
220    ) -> anyhow::Result<PathBuf>
221    where
222        T: HasTsInit + EncodeToRecordBatch + HasCatalogDataType,
223    {
224        if data.is_empty() {
225            return Ok(PathBuf::new());
226        }
227
228        let type_name = to_snake_case(std::any::type_name::<T>());
229        Self::check_ascending_timestamps(data, &type_name)?;
230
231        let chunk_metadata = T::chunk_metadata(data);
232        if let Some(position) = data
233            .iter()
234            .position(|item| !item.matches_chunk_metadata(&chunk_metadata))
235        {
236            anyhow::bail!(
237                "Cannot write {type_name} data with mixed identities: element {position} has \
238                 metadata {:?} but the chunk has {chunk_metadata:?}; write each \
239                 instrument or bar type separately",
240                data[position].metadata(),
241            );
242        }
243
244        let start_ts = start.unwrap_or(data.first().unwrap().ts_init());
245        let end_ts = end.unwrap_or(data.last().unwrap().ts_init());
246
247        let batches = self.data_to_record_batches(data)?;
248        let schema = batches.first().expect("Batches are empty.").schema();
249
250        let data_type = T::catalog_data_type();
251        let path_prefix = parquet_data_path_prefix(&data_type);
252        let identifier = if matches!(data_type, super::NautilusDataType::Bar) {
253            schema.metadata.get("bar_type").cloned()
254        } else {
255            schema.metadata.get("instrument_id").cloned()
256        };
257
258        let directory = self.make_path(path_prefix.as_ref(), identifier.as_deref())?;
259        self.write_parquet_file_checked(
260            &directory,
261            start_ts,
262            end_ts,
263            &batches,
264            skip_disjoint_check.unwrap_or(false),
265            "File",
266            Some(&format!("{type_name} data")),
267            None,
268        )
269    }
270
271    /// Writes custom data to a Parquet file in the catalog.
272    ///
273    /// This method handles writing custom data types that implement `CustomDataTrait`.
274    /// Custom data is organized by type name in a `custom/{type_name}/` directory structure.
275    ///
276    /// # Parameters
277    ///
278    /// - `data`: Vector of custom data items to write (must be in ascending timestamp order).
279    /// - `start`: Optional start timestamp to override the natural data range.
280    /// - `end`: Optional end timestamp to override the natural data range.
281    /// - `skip_disjoint_check`: Whether to skip interval disjointness validation.
282    ///
283    /// # Returns
284    ///
285    /// Returns the [`PathBuf`] of the created file, or an empty path if no data was provided.
286    ///
287    /// # Errors
288    ///
289    /// Returns an error if:
290    /// - The registered Arrow schema omits `ts_init` or uses incompatible timestamp types.
291    /// - Data serialization to Arrow record batches fails.
292    /// - Object store write operations fail.
293    /// - File path construction fails.
294    /// - Writing would create non-disjoint timestamp intervals (unless skipped).
295    pub fn write_custom_data_batch<D>(
296        &self,
297        data: D,
298        start: Option<UnixNanos>,
299        end: Option<UnixNanos>,
300        skip_disjoint_check: Option<bool>,
301    ) -> anyhow::Result<PathBuf>
302    where
303        D: AsRef<[CustomData]>,
304    {
305        let data = data.as_ref();
306        let data = data.iter().collect::<Vec<_>>();
307        self.write_custom_data_refs_batch(&data, start, end, skip_disjoint_check)
308    }
309
310    pub(crate) fn write_custom_data_refs_batch(
311        &self,
312        data: &[&CustomData],
313        start: Option<UnixNanos>,
314        end: Option<UnixNanos>,
315        skip_disjoint_check: Option<bool>,
316    ) -> anyhow::Result<PathBuf> {
317        if data.is_empty() {
318            return Ok(PathBuf::new());
319        }
320
321        let (batch, type_name, identifier, start_ts, end_ts) = prepare_custom_data_batch(data)?;
322        let start_ts = start.unwrap_or(start_ts);
323        let end_ts = end.unwrap_or(end_ts);
324        let batches = vec![record_batch_without_identifier_column(batch)?];
325
326        let directory = self.make_path_custom_data(&type_name, identifier.as_deref())?;
327        self.write_parquet_file_checked(
328            &directory,
329            start_ts,
330            end_ts,
331            &batches,
332            skip_disjoint_check.unwrap_or(false),
333            "File",
334            None,
335            None,
336        )
337    }
338
339    /// Writes instruments to Parquet files in the catalog.
340    ///
341    /// Instruments are stored under their instrument ID directory using timestamp-ranged
342    /// file names, allowing multiple historical versions of the same instrument to be
343    /// appended over time:
344    /// `data/instruments/{instrument_id}/{start_ts}-{end_ts}.parquet`
345    ///
346    /// # Parameters
347    ///
348    /// - `instruments`: Vector of instruments to write.
349    ///
350    /// # Returns
351    ///
352    /// Returns a vector of paths to the created files.
353    ///
354    /// # Errors
355    ///
356    /// Returns an error if:
357    /// - Data serialization fails.
358    /// - Object store write operations fail.
359    /// - File path construction fails.
360    ///
361    /// # Examples
362    ///
363    /// ```rust,no_run
364    /// use nautilus_model::instruments::InstrumentAny;
365    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
366    ///
367    /// let mut catalog = ParquetDataCatalog::new(
368    ///     std::path::Path::new("/tmp/nautilus_data"),
369    ///     None,
370    ///     None,
371    ///     None,
372    ///     None,
373    /// );
374    /// let instruments: Vec<InstrumentAny> = vec![/* instruments */];
375    ///
376    /// let paths = catalog.write_instruments(instruments)?;
377    /// # Ok::<(), anyhow::Error>(())
378    /// ```
379    pub fn write_instruments(
380        &self,
381        instruments: Vec<InstrumentAny>,
382    ) -> anyhow::Result<Vec<PathBuf>> {
383        use nautilus_model::instruments::Instrument;
384
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 start_ts = HasTsInit::ts_init(instrument_group.first().unwrap());
409            let end_ts = HasTsInit::ts_init(instrument_group.last().unwrap());
410            let batches = self.data_to_record_batches(&instrument_group)?;
411            if batches.is_empty() {
412                continue;
413            }
414
415            let directory = self.make_path(&instrument_prefix, Some(instrument_id.as_str()))?;
416
417            // ArrowWriter stores the full schema (including "class" metadata) in ARROW:schema.
418            // When reading, use the builder's schema for metadata (see query_instruments).
419            let path = self.write_parquet_file_checked(
420                &directory,
421                start_ts,
422                end_ts,
423                &batches,
424                false,
425                "Instrument file",
426                Some(&format!("instrument data for {instrument_id}")),
427                None,
428            )?;
429
430            paths.push(path);
431        }
432
433        Ok(paths)
434    }
435
436    /// Writes `batches` to a timestamp-named parquet file in `directory`, skipping when the
437    /// target file already exists and validating interval disjointness unless `skip_disjoint_check`.
438    ///
439    /// `file_label` and `data_description` parameterize the log messages so callers keep their
440    /// site-specific wording (`data_description = None` suppresses the pre-write log).
441    #[expect(clippy::too_many_arguments)]
442    pub(crate) fn write_parquet_file_checked(
443        &self,
444        directory: &str,
445        start_ts: UnixNanos,
446        end_ts: UnixNanos,
447        batches: &[RecordBatch],
448        skip_disjoint_check: bool,
449        file_label: &str,
450        data_description: Option<&str>,
451        replay_identity: Option<&str>,
452    ) -> anyhow::Result<PathBuf> {
453        let filename = timestamps_to_filename(start_ts, end_ts);
454        let filename = replay_identity.map_or(filename.clone(), |identity| {
455            let stem = filename.strip_suffix(".parquet").unwrap_or(&filename);
456            let digest = blake3::hash(identity.as_bytes()).to_hex();
457            format!("{stem}_{}.parquet", &digest[..16])
458        });
459        let path = PathBuf::from(directory).join(&filename);
460        let object_path = self.to_object_path(&path.to_string_lossy())?;
461
462        let file_exists = self.execute_async(|| async {
463            let exists: bool = self.object_store.head(&object_path).await.is_ok();
464            Ok(exists)
465        })?;
466
467        if file_exists {
468            log::info!(
469                "{file_label} {} already exists, skipping write",
470                path.display()
471            );
472            return Ok(path);
473        }
474
475        if !skip_disjoint_check {
476            let current_intervals = self.get_directory_intervals(directory)?;
477            let new_interval = (start_ts.as_u64(), end_ts.as_u64());
478            let mut new_intervals = current_intervals.clone();
479            new_intervals.push(new_interval);
480
481            if !are_intervals_disjoint(&new_intervals) {
482                anyhow::bail!(
483                    "Writing file {filename} with interval ({start_ts}, {end_ts}) would create \
484                    non-disjoint intervals. Existing intervals: {current_intervals:?}"
485                );
486            }
487        }
488
489        if let Some(data_description) = data_description {
490            log::info!(
491                "Writing {} batches of {data_description} to {}",
492                batches.len(),
493                path.display(),
494            );
495        }
496
497        self.execute_async(|| async {
498            let result = if replay_identity.is_some() {
499                write_batches_to_object_store_create(
500                    batches,
501                    self.object_store.clone(),
502                    &object_path,
503                    Some(self.compression),
504                    Some(self.max_row_group_size),
505                    None,
506                )
507                .await
508            } else {
509                write_batches_to_object_store(
510                    batches,
511                    self.object_store.clone(),
512                    &object_path,
513                    Some(self.compression),
514                    Some(self.max_row_group_size),
515                    None,
516                )
517                .await
518            };
519
520            if let Err(e) = result {
521                if replay_identity.is_some()
522                    && matches!(
523                        e.downcast_ref::<object_store::Error>(),
524                        Some(object_store::Error::AlreadyExists { .. })
525                    )
526                {
527                    return Ok(());
528                }
529                return Err(e);
530            }
531            Ok(())
532        })?;
533
534        Ok(path)
535    }
536
537    /// Writes typed data to a JSON file in the catalog.
538    ///
539    /// This method provides an alternative to Parquet format for data export and debugging.
540    /// JSON files are human-readable but less efficient for large datasets.
541    ///
542    /// # Type Parameters
543    ///
544    /// - `T`: The data type to write, must implement serialization and cataloging traits.
545    ///
546    /// # Parameters
547    ///
548    /// - `data`: Vector of data records to write (must be in ascending timestamp order).
549    /// - `path`: Optional custom directory path (defaults to catalog's standard structure).
550    /// - `write_metadata`: Whether to write a separate metadata file alongside the data.
551    ///
552    /// # Returns
553    ///
554    /// Returns the [`PathBuf`] of the created JSON file.
555    ///
556    /// # Errors
557    ///
558    /// Returns an error if:
559    /// - JSON serialization fails.
560    /// - Object store write operations fail.
561    /// - File path construction fails.
562    ///
563    /// # Panics
564    ///
565    /// Panics if data timestamps are not in ascending order.
566    ///
567    /// # Examples
568    ///
569    /// ```rust,no_run
570    /// use std::path::PathBuf;
571    /// use nautilus_model::data::TradeTick;
572    /// use nautilus_persistence::backend::parquet::catalog::ParquetDataCatalog;
573    ///
574    /// let mut catalog = ParquetDataCatalog::new(
575    ///     std::path::Path::new("/tmp/nautilus_data"),
576    ///     None,
577    ///     None,
578    ///     None,
579    ///     None,
580    /// );
581    /// let trades: Vec<TradeTick> = vec![/* trade data */];
582    ///
583    /// let path = catalog.write_to_json(
584    ///     trades,
585    ///     Some(PathBuf::from("/custom/path")),
586    ///     true  // write metadata
587    /// )?;
588    /// # Ok::<(), anyhow::Error>(())
589    /// ```
590    pub fn write_to_json<T>(
591        &self,
592        data: Vec<T>,
593        path: Option<PathBuf>,
594        write_metadata: bool,
595    ) -> anyhow::Result<PathBuf>
596    where
597        T: HasTsInit + Serialize + HasCatalogDataType + EncodeToRecordBatch,
598    {
599        if data.is_empty() {
600            return Ok(PathBuf::new());
601        }
602
603        let type_name = to_snake_case(std::any::type_name::<T>());
604        Self::check_ascending_timestamps(&data, &type_name)?;
605
606        let start_ts = data.first().unwrap().ts_init();
607        let end_ts = data.last().unwrap().ts_init();
608
609        let data_type = T::catalog_data_type();
610        let path_prefix = parquet_data_path_prefix(&data_type);
611        let directory = path
612            .unwrap_or_else(|| PathBuf::from(self.make_path(path_prefix.as_ref(), None).unwrap()));
613        let filename = timestamps_to_filename(start_ts, end_ts).replace(".parquet", ".json");
614        let json_path = directory.join(&filename);
615
616        log::info!(
617            "Writing {} records of {type_name} data to {}",
618            data.len(),
619            json_path.display(),
620        );
621
622        if write_metadata {
623            let metadata = T::chunk_metadata(&data);
624            let metadata_path = json_path.with_extension("metadata.json");
625            log::info!("Writing metadata to {}", metadata_path.display());
626
627            // Use object store for metadata file
628            let metadata_object_path = ObjectPath::from(metadata_path.to_string_lossy().as_ref());
629            let metadata_json = serde_json::to_vec_pretty(&metadata)?;
630            self.execute_async(|| async {
631                let _: object_store::PutResult = self
632                    .object_store
633                    .put(&metadata_object_path, metadata_json.into())
634                    .await?;
635                Ok(())
636            })?;
637        }
638
639        // Use object store for main JSON file
640        let json_object_path = ObjectPath::from(json_path.to_string_lossy().as_ref());
641        let json_data = serde_json::to_vec_pretty(&serde_json::to_value(data)?)?;
642        self.execute_async(|| async {
643            let _: object_store::PutResult = self
644                .object_store
645                .put(&json_object_path, json_data.into())
646                .await?;
647            Ok(())
648        })?;
649
650        Ok(json_path)
651    }
652
653    /// Validates that data timestamps are in ascending order.
654    ///
655    /// # Parameters
656    ///
657    /// - `data`: Slice of data records to validate.
658    /// - `type_name`: Name of the data type for error messages.
659    pub fn check_ascending_timestamps<T: HasTsInit>(
660        data: &[T],
661        type_name: &str,
662    ) -> anyhow::Result<()> {
663        if !data
664            .array_windows()
665            .all(|[a, b]| a.ts_init() <= b.ts_init())
666        {
667            anyhow::bail!("{type_name} timestamps must be in ascending order");
668        }
669
670        Ok(())
671    }
672
673    /// Converts data into Arrow record batches for Parquet serialization.
674    ///
675    /// This method chunks the data according to the configured batch size and converts
676    /// each chunk into an Arrow record batch with appropriate metadata.
677    ///
678    /// # Type Parameters
679    ///
680    /// - `T`: The data type to convert, must implement required encoding traits.
681    ///
682    /// # Parameters
683    ///
684    /// - `data`: Vector of data records to convert.
685    ///
686    /// # Returns
687    ///
688    /// Returns a vector of Arrow [`RecordBatch`] instances ready for Parquet serialization.
689    ///
690    /// # Errors
691    ///
692    /// Returns an error if record batch encoding fails for any chunk.
693    pub fn data_to_record_batches<T>(&self, data: &[T]) -> anyhow::Result<Vec<RecordBatch>>
694    where
695        T: HasTsInit + EncodeToRecordBatch,
696    {
697        if data.is_empty() {
698            return Ok(Vec::new());
699        }
700
701        let mut batches = Vec::new();
702        let metadata = EncodeToRecordBatch::chunk_metadata(data);
703
704        for chunk in data.chunks(self.batch_size) {
705            let record_batch = T::encode_batch(&metadata, chunk)?;
706            let record_batch = record_batch_without_identifier_column(record_batch)?;
707            batches.push(record_batch);
708        }
709
710        Ok(batches)
711    }
712}
713
714#[cfg(test)]
715mod tests {
716    use std::{
717        fmt::Display,
718        fs::File,
719        sync::{
720            Arc,
721            atomic::{AtomicUsize, Ordering},
722        },
723    };
724
725    use arrow::{
726        array::Int64Array,
727        datatypes::{DataType, Field, Schema},
728        record_batch::RecordBatch,
729    };
730    use futures::stream::BoxStream;
731    use nautilus_core::UnixNanos;
732    use nautilus_model::data::{
733        OrderBookDelta, OrderBookDepth,
734        stubs::{stub_delta, stub_depth10},
735    };
736    use nautilus_serialization::arrow::{
737        DecodeFromRecordBatch, KEY_PRICE_PRECISION, KEY_SIZE_PRECISION,
738    };
739    use object_store::{
740        CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
741        ObjectStoreExt, PutMode, PutMultipartOptions, PutOptions, PutPayload, PutResult,
742        Result as ObjectStoreResult, memory::InMemory, path::Path as ObjectPath,
743    };
744    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
745    use rstest::rstest;
746    use tempfile::TempDir;
747
748    use super::ParquetDataCatalog;
749    use crate::common::datafusion::DataBackendSession;
750
751    #[rstest]
752    fn depth_write_shares_file_metadata_across_chunks(stub_depth10: OrderBookDepth) {
753        let mut empty = stub_depth10.clone();
754        empty.bids.clear();
755        empty.asks.clear();
756        empty.bid_counts.clear();
757        empty.ask_counts.clear();
758        let directory = TempDir::new().unwrap();
759        let catalog = ParquetDataCatalog::from_uri(
760            directory.path().to_str().unwrap(),
761            None,
762            Some(2),
763            None,
764            Some(2),
765        )
766        .unwrap();
767        let data = vec![empty.clone(), empty, stub_depth10];
768
769        let path = catalog.write_to_parquet(&data, None, None, None).unwrap();
770        let builder = ParquetRecordBatchReaderBuilder::try_new(
771            File::open(directory.path().join(path)).unwrap(),
772        )
773        .unwrap();
774        let metadata = builder.schema().metadata().clone();
775        let batches = builder
776            .build()
777            .unwrap()
778            .collect::<Result<Vec<_>, _>>()
779            .unwrap();
780        let decoded = batches
781            .iter()
782            .cloned()
783            .flat_map(|batch| OrderBookDepth::decode_batch(&metadata, batch).unwrap())
784            .collect::<Vec<_>>();
785
786        assert_eq!(metadata[KEY_PRICE_PRECISION], "2");
787        assert_eq!(metadata[KEY_SIZE_PRECISION], "0");
788        assert_eq!(decoded.len(), 3);
789        assert_eq!(decoded, data);
790        assert_eq!(decoded[2].bids[0].price.precision, 2);
791    }
792
793    #[rstest]
794    fn leading_clear_delta_writes_with_following_precision(stub_delta: OrderBookDelta) {
795        let directory = TempDir::new().unwrap();
796        let catalog = ParquetDataCatalog::from_uri(
797            directory.path().to_str().unwrap(),
798            None,
799            Some(2),
800            None,
801            None,
802        )
803        .unwrap();
804        let clear = OrderBookDelta::clear(
805            stub_delta.instrument_id,
806            0,
807            UnixNanos::from(1),
808            UnixNanos::from(1),
809        );
810        let second_clear = OrderBookDelta::clear(
811            stub_delta.instrument_id,
812            0,
813            UnixNanos::from(2),
814            UnixNanos::from(2),
815        );
816
817        let path = catalog
818            .write_to_parquet(&[clear, second_clear, stub_delta], None, None, None)
819            .unwrap();
820        let builder = ParquetRecordBatchReaderBuilder::try_new(
821            File::open(directory.path().join(&path)).unwrap(),
822        )
823        .unwrap();
824        let metadata = builder.schema().metadata().clone();
825        let decoded = builder
826            .build()
827            .unwrap()
828            .map(|batch| OrderBookDelta::decode_batch(&metadata, batch.unwrap()).unwrap())
829            .collect::<Vec<_>>()
830            .concat();
831
832        assert!(directory.path().join(path).exists());
833        assert_eq!(metadata[KEY_PRICE_PRECISION], "2");
834        assert_eq!(decoded[2].order.price.precision, 2);
835    }
836
837    #[derive(Debug)]
838    struct CreateRaceStore {
839        inner: InMemory,
840        create_calls: AtomicUsize,
841    }
842
843    impl Display for CreateRaceStore {
844        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
845            f.write_str("create-race")
846        }
847    }
848
849    #[async_trait::async_trait]
850    impl ObjectStore for CreateRaceStore {
851        async fn put_opts(
852            &self,
853            location: &ObjectPath,
854            payload: PutPayload,
855            opts: PutOptions,
856        ) -> ObjectStoreResult<PutResult> {
857            if opts.mode == PutMode::Create {
858                self.create_calls.fetch_add(1, Ordering::Relaxed);
859                return Err(object_store::Error::AlreadyExists {
860                    path: location.to_string(),
861                    source: Box::new(std::io::Error::new(
862                        std::io::ErrorKind::AlreadyExists,
863                        "injected create race",
864                    )),
865                });
866            }
867            self.inner.put_opts(location, payload, opts).await
868        }
869
870        async fn put_multipart_opts(
871            &self,
872            location: &ObjectPath,
873            opts: PutMultipartOptions,
874        ) -> ObjectStoreResult<Box<dyn MultipartUpload>> {
875            self.inner.put_multipart_opts(location, opts).await
876        }
877
878        async fn get_opts(
879            &self,
880            location: &ObjectPath,
881            options: GetOptions,
882        ) -> ObjectStoreResult<GetResult> {
883            self.inner.get_opts(location, options).await
884        }
885
886        fn list(
887            &self,
888            prefix: Option<&ObjectPath>,
889        ) -> BoxStream<'static, ObjectStoreResult<ObjectMeta>> {
890            self.inner.list(prefix)
891        }
892
893        async fn list_with_delimiter(
894            &self,
895            prefix: Option<&ObjectPath>,
896        ) -> ObjectStoreResult<ListResult> {
897            self.inner.list_with_delimiter(prefix).await
898        }
899
900        fn delete_stream(
901            &self,
902            locations: BoxStream<'static, ObjectStoreResult<ObjectPath>>,
903        ) -> BoxStream<'static, ObjectStoreResult<ObjectPath>> {
904            self.inner.delete_stream(locations)
905        }
906
907        async fn copy_opts(
908            &self,
909            from: &ObjectPath,
910            to: &ObjectPath,
911            opts: CopyOptions,
912        ) -> ObjectStoreResult<()> {
913            self.inner.copy_opts(from, to, opts).await
914        }
915    }
916
917    #[rstest]
918    fn promotion_writes_to_memory_store_without_copy_support() {
919        let catalog = ParquetDataCatalog {
920            base_path: "catalog".to_string(),
921            original_uri: "memory://".to_string(),
922            object_store: Arc::new(InMemory::new()),
923            session: DataBackendSession::new(5_000),
924            batch_size: 5_000,
925            compression: parquet::basic::Compression::SNAPPY,
926            max_row_group_size: 5_000,
927        };
928        let batch = RecordBatch::try_new(
929            Arc::new(Schema::new(vec![Field::new(
930                "ts_init",
931                DataType::Int64,
932                false,
933            )])),
934            vec![Arc::new(Int64Array::from(vec![1]))],
935        )
936        .unwrap();
937
938        let path = catalog
939            .write_parquet_file_checked(
940                "quotes/TEST",
941                UnixNanos::from(1),
942                UnixNanos::from(1),
943                &[batch],
944                false,
945                "Promoted file",
946                None,
947                Some("memory-replay"),
948            )
949            .unwrap();
950        let object_path = catalog.to_object_path(&path.to_string_lossy()).unwrap();
951        catalog
952            .execute_async(|| async {
953                catalog.object_store.head(&object_path).await?;
954                Ok(())
955            })
956            .unwrap();
957    }
958
959    #[rstest]
960    fn promotion_accepts_already_exists_after_head_miss() {
961        let object_store = Arc::new(CreateRaceStore {
962            inner: InMemory::new(),
963            create_calls: AtomicUsize::new(0),
964        });
965        let catalog = ParquetDataCatalog {
966            base_path: "catalog".to_string(),
967            original_uri: "memory://".to_string(),
968            object_store: object_store.clone(),
969            session: DataBackendSession::new(5_000),
970            batch_size: 5_000,
971            compression: parquet::basic::Compression::SNAPPY,
972            max_row_group_size: 5_000,
973        };
974        let batch = RecordBatch::try_new(
975            Arc::new(Schema::new(vec![Field::new(
976                "ts_init",
977                DataType::Int64,
978                false,
979            )])),
980            vec![Arc::new(Int64Array::from(vec![1]))],
981        )
982        .unwrap();
983
984        catalog
985            .write_parquet_file_checked(
986                "quotes/TEST",
987                UnixNanos::from(1),
988                UnixNanos::from(1),
989                &[batch],
990                true,
991                "Promoted file",
992                None,
993                Some("racing-replay"),
994            )
995            .unwrap();
996
997        assert_eq!(object_store.create_calls.load(Ordering::Relaxed), 1);
998    }
999}
nautilus_persistence/backend/parquet/catalog/write.rs

nautilus_persistence/backend/parquet/catalog/
write.rs
