nautilus_persistence/backend/parquet/
writer.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautilustrader.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
10//! Feather-staged promoting writer for the Parquet catalog.
11
12use std::{
13    fmt::Debug,
14    sync::{
15        Arc,
16        atomic::{AtomicBool, Ordering},
17    },
18};
19
20use nautilus_common::live::block_on_nautilus_with;
21use object_store::{
22    Error as ObjectStoreError, ObjectStoreExt, PutMode, PutOptions, path::Path as ObjectPath,
23};
24use serde::{Deserialize, Serialize};
25
26use super::catalog::ParquetDataCatalog;
27use crate::{
28    backend::migration::feather_replay_identity,
29    common::{
30        conversion::FeatherConversionSummary, datafusion::identifiers_from_record_batches,
31        storage::create_storage_backend_from_path,
32    },
33    writer::{
34        factory::{PARQUET_WRITER_FACTORY_NAME, WriterConnectConfig, WriterFactoryRegistry},
35        feather::WriterClock,
36        materializer::read_feather_record_batches_with_identity,
37        promotion::{
38            PromotionBackend, PromotionResult, PromotionSession, PromotionWork,
39            StagedPromotionBackend,
40        },
41        run::{FeatherSessionSource, RunStatus},
42        staged::{StagedFeatherWriter, StagedWriter},
43        traits::StreamingDataSink,
44    },
45};
46
47pub(crate) fn register_factory(registry: &mut WriterFactoryRegistry) {
48    registry.insert(
49        PARQUET_WRITER_FACTORY_NAME.to_string(),
50        Arc::new(parquet_writer_factory),
51    );
52}
53
54fn parquet_writer_factory(
55    config: &WriterConnectConfig,
56    clock: WriterClock,
57) -> anyhow::Result<StreamingDataSink> {
58    let params = config.params.as_ref();
59    let interval_ms = params.and_then(|params| params.get_u64("parquet_commit_interval_ms"));
60    let promote_on_close = params
61        .and_then(|params| params.get_bool("promote_on_close"))
62        .unwrap_or(true);
63    let delete_feather_after_commit = params
64        .and_then(|params| params.get_bool("delete_feather_after_commit"))
65        .unwrap_or(false);
66    let use_ts_event_for_ts_init = params
67        .and_then(|params| params.get_bool("use_ts_event_for_ts_init"))
68        .unwrap_or(false);
69    Ok(Box::new(ParquetWriter::new(
70        config,
71        clock,
72        interval_ms,
73        promote_on_close,
74        delete_feather_after_commit,
75        use_ts_event_for_ts_init,
76    )?))
77}
78
79#[expect(
80    clippy::struct_excessive_bools,
81    reason = "the booleans are independent writer policy and lifecycle flags"
82)]
83struct ParquetWriter {
84    core: StagedFeatherWriter<ParquetPromotionBackend>,
85    session: PromotionSession,
86    source: FeatherSessionSource,
87    catalog_uri: String,
88    storage_options: Option<ahash::AHashMap<String, String>>,
89    interval_ms: Option<u64>,
90    promote_on_close: bool,
91    delete_feather_after_commit: bool,
92    use_ts_event_for_ts_init: bool,
93    legacy_manifest_missing: Arc<AtomicBool>,
94    has_data: bool,
95    run_status: RunStatus,
96}
97
98impl Debug for ParquetWriter {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        f.debug_struct(stringify!(ParquetWriter))
101            .field("staging_uri", &self.core.storage.original_uri)
102            .field("catalog_uri", &self.catalog_uri)
103            .field("interval_ms", &self.interval_ms)
104            .finish_non_exhaustive()
105    }
106}
107
108impl ParquetWriter {
109    fn new(
110        config: &WriterConnectConfig,
111        clock: WriterClock,
112        interval_ms: Option<u64>,
113        promote_on_close: bool,
114        delete_feather_after_commit: bool,
115        use_ts_event_for_ts_init: bool,
116    ) -> anyhow::Result<Self> {
117        let storage =
118            create_storage_backend_from_path(&config.uri, config.storage_options.clone())?;
119        let session = StagedFeatherWriter::<ParquetPromotionBackend>::required_session(
120            &storage.original_uri,
121            "Parquet writer URI",
122        )?;
123        let source_storage =
124            create_storage_backend_from_path(&session.catalog_uri, config.storage_options.clone())?;
125        let source = FeatherSessionSource::new(
126            source_storage,
127            session.kind.clone(),
128            session.instance_id.clone(),
129        );
130        let mut core = StagedFeatherWriter::new(
131            storage.clone(),
132            clock,
133            config.rotation_config.clone(),
134            None,
135            None,
136            config.flush_interval_ms,
137            config.record_filter.clone(),
138        )?;
139        block_on_nautilus_with(|| {
140            storage.write_current_run_manifest(
141                &session.kind,
142                &session.instance_id,
143                "in_progress",
144                true,
145            )
146        })?;
147        let timer_catalog_uri = session.catalog_uri.clone();
148        let timer_storage_options = config.storage_options.clone();
149        core.warn_if_orphan_feather_present("Parquet");
150        let legacy_manifest_missing = Arc::new(AtomicBool::new(false));
151        let timer_legacy_manifest_missing = Arc::clone(&legacy_manifest_missing);
152        core.start_promotion_timer(
153            interval_ms,
154            source.clone(),
155            "parquet-promotion",
156            move || {
157                Ok(ParquetPromotionBackend::new(
158                    ParquetDataCatalog::from_uri(
159                        &timer_catalog_uri,
160                        timer_storage_options.clone(),
161                        None,
162                        None,
163                        None,
164                    )?,
165                    Arc::clone(&timer_legacy_manifest_missing),
166                ))
167            },
168            use_ts_event_for_ts_init,
169            delete_feather_after_commit,
170        )?;
171        Ok(Self {
172            core,
173            catalog_uri: session.catalog_uri.clone(),
174            storage_options: config.storage_options.clone(),
175            session,
176            source,
177            interval_ms,
178            promote_on_close,
179            delete_feather_after_commit,
180            use_ts_event_for_ts_init,
181            legacy_manifest_missing,
182            has_data: false,
183            run_status: RunStatus::InProgress,
184        })
185    }
186
187    fn record_status(&mut self, status: RunStatus) -> anyhow::Result<()> {
188        if self.run_status == status {
189            return Ok(());
190        }
191        block_on_nautilus_with(|| {
192            self.core.storage.write_current_run_manifest(
193                &self.session.kind,
194                &self.session.instance_id,
195                status.as_str(),
196                !self.has_data,
197            )
198        })?;
199        self.run_status = status;
200        Ok(())
201    }
202
203    fn mark_non_empty(&mut self) -> anyhow::Result<()> {
204        if self.has_data {
205            return Ok(());
206        }
207        block_on_nautilus_with(|| {
208            self.core.storage.write_current_run_manifest(
209                &self.session.kind,
210                &self.session.instance_id,
211                "in_progress",
212                false,
213            )
214        })?;
215        self.has_data = true;
216        Ok(())
217    }
218
219    fn interval_ns(&self) -> Option<u64> {
220        self.interval_ms
221            .filter(|interval_ms| *interval_ms > 0)
222            .map(|interval_ms| interval_ms.saturating_mul(1_000_000))
223    }
224
225    fn prepare_promotion(&self) -> anyhow::Result<Option<PromotionWork<ParquetPromotionBackend>>> {
226        let catalog = ParquetDataCatalog::from_uri(
227            &self.catalog_uri,
228            self.storage_options.clone(),
229            None,
230            None,
231            None,
232        )?;
233        self.core.prepare_promotion(
234            ParquetPromotionBackend::new(catalog, Arc::clone(&self.legacy_manifest_missing)),
235            self.source.clone(),
236            self.use_ts_event_for_ts_init,
237            self.delete_feather_after_commit,
238        )
239    }
240
241    fn finalize(
242        &mut self,
243        result: PromotionResult,
244    ) -> anyhow::Result<Vec<FeatherConversionSummary>> {
245        if result.converted.is_err() {
246            if result.run_state_recorded {
247                self.run_status = RunStatus::Failed;
248            } else if let Err(e) = self.record_status(RunStatus::Failed) {
249                log::warn!("Failed to record Parquet run Failed state: {e}");
250            }
251        } else if result.run_state_recorded {
252            self.run_status = RunStatus::Promoted;
253        }
254        self.core.finalize_promotion(result)
255    }
256
257    fn drain_completed(&mut self) -> anyhow::Result<Vec<FeatherConversionSummary>> {
258        let mut converted = Vec::new();
259        for result in self.core.promotion_driver.drain_completed()? {
260            converted.extend(self.finalize(result)?);
261        }
262        Ok(converted)
263    }
264
265    fn wait(&mut self) -> anyhow::Result<Vec<FeatherConversionSummary>> {
266        let mut converted = Vec::new();
267        for result in self.core.promotion_driver.wait()? {
268            converted.extend(self.finalize(result)?);
269        }
270        Ok(converted)
271    }
272
273    fn promote_now(&mut self) -> anyhow::Result<Vec<FeatherConversionSummary>> {
274        let mut converted = self.wait()?;
275        if let Some(work) = self.prepare_promotion()? {
276            converted.extend(self.finalize(work.execute())?);
277        }
278        Ok(converted)
279    }
280}
281
282impl StagedWriter for ParquetWriter {
283    type Backend = ParquetPromotionBackend;
284
285    fn staged(&mut self) -> &mut StagedFeatherWriter<Self::Backend> {
286        &mut self.core
287    }
288
289    fn mark_non_empty(&mut self) -> anyhow::Result<()> {
290        Self::mark_non_empty(self)
291    }
292
293    fn maybe_promote(&mut self) -> anyhow::Result<()> {
294        self.drain_completed()?;
295        let Some(interval_ns) = self.interval_ns() else {
296            return Ok(());
297        };
298
299        if self
300            .core
301            .promotion_driver
302            .is_due(self.core.clock.timestamp_ns(), interval_ns)
303            && let Some(work) = self.prepare_promotion()?
304        {
305            self.core
306                .promotion_driver
307                .submit(work, "parquet-promotion")?;
308            self.core
309                .promotion_driver
310                .mark_committed_at(self.core.clock.timestamp_ns());
311        }
312        Ok(())
313    }
314
315    fn wait_for_promotions(&mut self) -> anyhow::Result<Vec<FeatherConversionSummary>> {
316        self.wait()
317    }
318
319    fn promote_on_flush(&self) -> bool {
320        self.interval_ns().is_some_and(|interval_ns| {
321            self.core
322                .promotion_driver
323                .is_due(self.core.clock.timestamp_ns(), interval_ns)
324        })
325    }
326
327    fn promote_on_close(&self) -> bool {
328        self.promote_on_close
329    }
330
331    fn promote_now(&mut self) -> anyhow::Result<Vec<FeatherConversionSummary>> {
332        Self::promote_now(self)
333    }
334
335    fn record_completed(&mut self) {
336        if self.run_status == RunStatus::InProgress
337            && let Err(e) = self.record_status(RunStatus::Completed)
338        {
339            log::warn!("Failed to complete Parquet run manifest: {e}");
340        }
341    }
342}
343
344impl Drop for ParquetWriter {
345    fn drop(&mut self) {
346        self.core.stop_promotion_timer();
347        let promotion_failed = if let Err(e) = self.wait() {
348            log::warn!("ParquetWriter dropped with pending promotion error: {e}");
349            true
350        } else {
351            false
352        };
353
354        if self.run_status == RunStatus::InProgress {
355            let status = if promotion_failed || !self.core.is_closed().unwrap_or(false) {
356                RunStatus::Failed
357            } else {
358                RunStatus::Completed
359            };
360
361            if let Err(e) = self.record_status(status) {
362                log::warn!(
363                    "Failed to record Parquet run {} state during drop: {e}",
364                    status.as_str()
365                );
366            }
367        }
368        self.core.flush_on_drop("ParquetWriter");
369    }
370}
371
372struct ParquetPromotionBackend {
373    catalog: ParquetDataCatalog,
374    legacy_manifest_missing: Arc<AtomicBool>,
375}
376
377const PROMOTION_MANIFEST: &str = "_nautilus_promotions.json";
378const PROMOTION_MARKERS: &str = "_nautilus_promotions";
379
380#[derive(Debug, Default, Deserialize, Serialize)]
381struct ParquetPromotionManifest {
382    identities: Vec<String>,
383}
384
385impl ParquetPromotionBackend {
386    fn new(catalog: ParquetDataCatalog, legacy_manifest_missing: Arc<AtomicBool>) -> Self {
387        Self {
388            catalog,
389            legacy_manifest_missing,
390        }
391    }
392
393    fn manifest_path(&self) -> ObjectPath {
394        let base = self.catalog.base_path.trim_matches('/');
395        if base.is_empty() {
396            ObjectPath::from(PROMOTION_MANIFEST)
397        } else {
398            ObjectPath::from(format!("{base}/{PROMOTION_MANIFEST}"))
399        }
400    }
401
402    fn manifest(&self) -> anyhow::Result<ParquetPromotionManifest> {
403        if self.legacy_manifest_missing.load(Ordering::Relaxed) {
404            return Ok(ParquetPromotionManifest::default());
405        }
406        let path = self.manifest_path();
407        block_on_nautilus_with(|| async {
408            let result = match self.catalog.object_store.get(&path).await {
409                Ok(result) => result,
410                Err(ObjectStoreError::NotFound { .. }) => {
411                    self.legacy_manifest_missing.store(true, Ordering::Relaxed);
412                    return Ok(ParquetPromotionManifest::default());
413                }
414                Err(e) => return Err(e.into()),
415            };
416            Ok(serde_json::from_slice(&result.bytes().await?)?)
417        })
418    }
419
420    fn marker_path(&self, identity: &str) -> ObjectPath {
421        let base = self.catalog.base_path.trim_matches('/');
422        let digest = blake3::hash(identity.as_bytes()).to_hex();
423        let path = format!("{PROMOTION_MARKERS}/{digest}.json");
424        if base.is_empty() {
425            ObjectPath::from(path)
426        } else {
427            ObjectPath::from(format!("{base}/{path}"))
428        }
429    }
430
431    fn identity_recorded(&self, identity: &str) -> anyhow::Result<bool> {
432        let path = self.marker_path(identity);
433        let marker_exists = block_on_nautilus_with(|| async {
434            match self.catalog.object_store.head(&path).await {
435                Ok(_) => Ok::<bool, anyhow::Error>(true),
436                Err(ObjectStoreError::NotFound { .. }) => Ok(false),
437                Err(e) => Err(anyhow::Error::from(e)),
438            }
439        })?;
440        Ok(marker_exists
441            || self
442                .manifest()?
443                .identities
444                .iter()
445                .any(|item| item == identity))
446    }
447
448    fn record_identity(&self, identity: &str) -> anyhow::Result<()> {
449        let bytes = serde_json::to_vec(identity)?;
450        let path = self.marker_path(identity);
451        block_on_nautilus_with(|| async {
452            match self
453                .catalog
454                .object_store
455                .put_opts(
456                    &path,
457                    bytes.into(),
458                    PutOptions {
459                        mode: PutMode::Create,
460                        ..Default::default()
461                    },
462                )
463                .await
464            {
465                Ok(_) | Err(ObjectStoreError::AlreadyExists { .. }) => Ok(()),
466                Err(e) => Err(e.into()),
467            }
468        })
469    }
470}
471
472impl PromotionBackend for ParquetPromotionBackend {
473    type Source = FeatherSessionSource;
474
475    const NAME: &'static str = "Parquet";
476
477    fn convert_file(
478        &mut self,
479        source: &Self::Source,
480        file: &str,
481        use_ts_event_for_ts_init: bool,
482        record_promoted: bool,
483    ) -> anyhow::Result<Option<FeatherConversionSummary>> {
484        let object_path = ObjectPath::from(file);
485        let read = block_on_nautilus_with(|| {
486            read_feather_record_batches_with_identity(
487                source.storage.object_store.clone(),
488                &object_path,
489            )
490        })?;
491        let identifiers = identifiers_from_record_batches(&read.batches)
492            .ok()
493            .filter(|identifiers| !identifiers.is_empty());
494        let identity = feather_replay_identity(
495            &source.storage.original_uri,
496            file,
497            &read.content_hash,
498            identifiers.as_deref(),
499        );
500
501        if self.identity_recorded(&identity)? {
502            return Ok(None);
503        }
504        let summary = self.catalog.promote_feather_file(
505            source,
506            file,
507            read.batches,
508            use_ts_event_for_ts_init,
509            &identity,
510        )?;
511
512        if summary.is_some() && record_promoted {
513            self.record_identity(&identity)?;
514        }
515        Ok(summary)
516    }
517
518    fn delete_file(&mut self, source: &Self::Source, file: &str) -> anyhow::Result<()> {
519        block_on_nautilus_with(|| async {
520            source
521                .storage
522                .object_store
523                .delete(&ObjectPath::from(file))
524                .await?;
525            Ok::<(), anyhow::Error>(())
526        })
527    }
528}
529
530impl StagedPromotionBackend for ParquetPromotionBackend {
531    const REQUIRES_SESSION: bool = true;
532
533    fn record_run_state(
534        &mut self,
535        source: &FeatherSessionSource,
536        _staging_uri: &str,
537        status: crate::writer::run::RunStatus,
538        empty: bool,
539        _error: Option<&str>,
540    ) -> anyhow::Result<()> {
541        block_on_nautilus_with(|| {
542            source.storage.write_run_manifest(
543                &source.kind,
544                &source.instance_id,
545                status.as_str(),
546                empty,
547            )
548        })
549    }
550}
551
552#[cfg(test)]
553mod tests {
554    use std::sync::atomic::AtomicU64;
555
556    use nautilus_core::UnixNanos;
557    use nautilus_model::{
558        data::{Data, DataBatch, NautilusDataType, NautilusRecordType, QuoteTick},
559        identifiers::InstrumentId,
560        types::{ERROR_PRICE, Price, Quantity},
561    };
562    use rstest::rstest;
563    use serde_json::json;
564    use tempfile::TempDir;
565
566    use super::*;
567    use crate::{
568        catalog::traits::{CatalogQuery, CatalogReader},
569        test_data::RustTestHashMapCustomData,
570    };
571
572    #[rstest]
573    fn parquet_default_close_promotes_and_honors_source_retention(
574        #[values(false, true)] delete_source: bool,
575    ) {
576        let directory = TempDir::new().unwrap();
577        let staging = directory.path().join("backtest").join("run-1");
578        let mut config = WriterConnectConfig::new(staging.to_string_lossy(), None);
579        config.params = Some(
580            serde_json::from_value(json!({"delete_feather_after_commit": delete_source})).unwrap(),
581        );
582        let mut sink =
583            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
584                .unwrap();
585        let quote = sample_quote();
586        let mut second = quote;
587        second.instrument_id = InstrumentId::from("BTC/USD.SIM");
588        second.bid_price = Price::from("100.1234");
589        second.ask_price = Price::from("100.5678");
590        second.ts_init = UnixNanos::from(24);
591        let mut catalog = ParquetDataCatalog::from_uri(
592            directory.path().to_str().unwrap(),
593            None,
594            None,
595            None,
596            None,
597        )
598        .unwrap();
599        sink.write_data(Data::Quote(quote)).unwrap();
600        sink.write_data(Data::Quote(second)).unwrap();
601        sink.flush().unwrap();
602        let before = catalog
603            .query_batch(&CatalogQuery::new(NautilusDataType::QuoteTick))
604            .unwrap();
605        sink.close().unwrap();
606        let after = catalog
607            .query_batch(&CatalogQuery::new(NautilusDataType::QuoteTick))
608            .unwrap();
609        let storage = create_storage_backend_from_path(staging.to_str().unwrap(), None).unwrap();
610        let staged = block_on_nautilus_with(|| storage.list_files("", None))
611            .unwrap()
612            .into_iter()
613            .filter(|path| path.ends_with(".feather"))
614            .count();
615        assert!(before.is_empty());
616        let DataBatch::Quote(rows) = after else {
617            panic!("expected quotes")
618        };
619
620        assert_eq!(rows.as_ref(), &[quote, second]);
621        assert_eq!(staged, usize::from(!delete_source));
622    }
623
624    #[rstest]
625    #[case(None)]
626    #[case(Some(0))]
627    #[case(Some(1))]
628    fn parquet_interval_uses_test_clock_and_close_can_leave_staged_data(
629        #[case] interval: Option<u64>,
630    ) {
631        let directory = TempDir::new().unwrap();
632        let staging = directory.path().join("backtest").join("run-2");
633        let mut config = WriterConnectConfig::new(staging.to_string_lossy(), None);
634        config.params = Some(
635            serde_json::from_value(
636                json!({"parquet_commit_interval_ms": interval, "promote_on_close": false}),
637            )
638            .unwrap(),
639        );
640        let clock = Arc::new(AtomicU64::new(0));
641        let mut sink =
642            parquet_writer_factory(&config, WriterClock::Test(Arc::clone(&clock))).unwrap();
643        let quote = sample_quote();
644        sink.write_data(Data::Quote(quote)).unwrap();
645        clock.store(1_000_000, Ordering::Relaxed);
646        sink.flush().unwrap();
647        sink.close().unwrap();
648        let mut catalog = ParquetDataCatalog::from_uri(
649            directory.path().to_str().unwrap(),
650            None,
651            None,
652            None,
653            None,
654        )
655        .unwrap();
656        let before_manual = catalog
657            .query_batch(&CatalogQuery::new(NautilusDataType::QuoteTick))
658            .unwrap();
659
660        if interval != Some(1) {
661            catalog
662                .convert_stream_to_data(
663                    "run-2",
664                    &NautilusDataType::QuoteTick.into(),
665                    Some("backtest"),
666                    None,
667                    false,
668                )
669                .unwrap();
670        }
671        let after = catalog
672            .query_batch(&CatalogQuery::new(NautilusDataType::QuoteTick))
673            .unwrap();
674        assert_eq!(before_manual.len(), usize::from(interval == Some(1)));
675        let DataBatch::Quote(rows) = after else {
676            panic!("expected quotes")
677        };
678        assert_eq!(rows.as_ref(), &[quote]);
679    }
680
681    #[rstest]
682    fn parquet_streaming_funding_round_trip() {
683        use nautilus_model::data::FundingRateUpdate;
684        let directory = TempDir::new().unwrap();
685
686        let config = WriterConnectConfig::new(
687            directory
688                .path()
689                .join("backtest/run-funding")
690                .to_string_lossy(),
691            None,
692        );
693        let mut sink =
694            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
695                .unwrap();
696
697        let funding = FundingRateUpdate::new(
698            InstrumentId::from("AUD/USD.SIM"),
699            "0.0012".parse().unwrap(),
700            Some(480),
701            Some(789.into()),
702            123.into(),
703            456.into(),
704        );
705        assert!(sink.write_any(&funding).unwrap());
706        sink.close().unwrap();
707        let mut catalog = ParquetDataCatalog::from_uri(
708            directory.path().to_str().unwrap(),
709            None,
710            None,
711            None,
712            None,
713        )
714        .unwrap();
715        let rows = catalog
716            .query_batch(&CatalogQuery::new(NautilusDataType::FundingRateUpdate))
717            .unwrap();
718
719        let DataBatch::FundingRate(rows) = rows else {
720            panic!("expected funding rates")
721        };
722
723        assert_eq!(rows.as_ref(), &[funding]);
724    }
725
726    #[rstest]
727    fn parquet_streaming_write_error_reaches_flush() {
728        let directory = TempDir::new().unwrap();
729
730        let config = WriterConnectConfig::new(
731            directory
732                .path()
733                .join("backtest/run-error")
734                .to_string_lossy(),
735            None,
736        );
737        let mut sink =
738            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
739                .unwrap();
740        let mut quote = sample_quote();
741        quote.bid_price = ERROR_PRICE;
742        let write_error = sink.write_any(&quote).unwrap_err().to_string();
743        let flush_error = sink.flush().unwrap_err().to_string();
744        assert_eq!(flush_error, write_error);
745    }
746
747    #[rstest]
748    fn parquet_streaming_instrument_promotes() {
749        use nautilus_model::instruments::{InstrumentAny, stubs::audusd_sim};
750        let directory = TempDir::new().unwrap();
751
752        let config = WriterConnectConfig::new(
753            directory
754                .path()
755                .join("backtest/run-instrument")
756                .to_string_lossy(),
757            None,
758        );
759        let mut sink =
760            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
761                .unwrap();
762        let instrument = InstrumentAny::CurrencyPair(audusd_sim());
763        assert!(sink.write_any(&instrument).unwrap());
764        sink.close().unwrap();
765        let mut catalog = ParquetDataCatalog::from_uri(
766            directory.path().to_str().unwrap(),
767            None,
768            None,
769            None,
770            None,
771        )
772        .unwrap();
773        let rows = catalog
774            .query_batch(&CatalogQuery::new(NautilusDataType::Instrument))
775            .unwrap();
776
777        let DataBatch::Instrument(rows) = rows else {
778            panic!("expected instruments")
779        };
780
781        assert_eq!(rows.as_ref(), &[instrument]);
782    }
783
784    #[rstest]
785    fn parquet_streaming_voided_fill_promotes() {
786        use nautilus_model::events::{OrderFillVoided, order::spec::OrderFillVoidedSpec};
787        use nautilus_serialization::arrow::DecodeTypedFromRecordBatch;
788        let directory = TempDir::new().unwrap();
789
790        let config = WriterConnectConfig::new(
791            directory.path().join("backtest/run-void").to_string_lossy(),
792            None,
793        );
794        let mut sink =
795            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
796                .unwrap();
797        let event = OrderFillVoidedSpec::builder().is_reopened(true).build();
798        assert!(sink.write_any(&event).unwrap());
799        sink.close().unwrap();
800        let mut catalog = ParquetDataCatalog::from_uri(
801            directory.path().to_str().unwrap(),
802            None,
803            None,
804            None,
805            None,
806        )
807        .unwrap();
808        let batches = catalog
809            .query_record_batches(
810                &NautilusRecordType::OrderFillVoided.into(),
811                None,
812                None,
813                None,
814                None,
815                true,
816            )
817            .unwrap();
818        let mut rows = Vec::new();
819        for batch in batches {
820            rows.extend(
821                OrderFillVoided::decode_typed_batch(batch.schema().metadata(), batch.clone())
822                    .unwrap(),
823            );
824        }
825
826        assert_eq!(rows, vec![event]);
827    }
828
829    #[rstest]
830    #[case(true)]
831    #[case(false)]
832    fn parquet_promotion_retains_conflicting_schema_groups(#[case] automatic: bool) {
833        use nautilus_model::{
834            data::{BookOrder, OrderBookDepth},
835            enums::OrderSide,
836        };
837        let directory = TempDir::new().unwrap();
838        let staging = directory.path().join("backtest/run-depth-ties");
839        let mut config = WriterConnectConfig::new(staging.to_string_lossy(), None);
840        config.params = Some(
841            serde_json::from_value(
842                json!({"delete_feather_after_commit": true, "promote_on_close": automatic}),
843            )
844            .unwrap(),
845        );
846        let mut sink =
847            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
848                .unwrap();
849        let id = InstrumentId::from("AUD/USD.SIM");
850        let empty =
851            OrderBookDepth::new(id, vec![], vec![], vec![], vec![], 1, 2, 3.into(), 4.into());
852
853        let order = BookOrder::new(
854            OrderSide::Buy,
855            Price::from("1.23"),
856            Quantity::from("4.5"),
857            6,
858        );
859
860        let populated = OrderBookDepth::new(
861            id,
862            vec![order],
863            vec![],
864            vec![7],
865            vec![],
866            8,
867            9,
868            3.into(),
869            4.into(),
870        );
871        sink.write_any(&empty).unwrap();
872        sink.write_any(&populated).unwrap();
873
874        let error = if automatic {
875            sink.close().unwrap_err().to_string()
876        } else {
877            sink.close().unwrap();
878            let mut catalog = ParquetDataCatalog::from_uri(
879                directory.path().to_str().unwrap(),
880                None,
881                None,
882                None,
883                None,
884            )
885            .unwrap();
886            catalog
887                .convert_stream_to_data(
888                    "run-depth-ties",
889                    &NautilusDataType::OrderBookDepth.into(),
890                    Some("backtest"),
891                    None,
892                    false,
893                )
894                .unwrap_err()
895                .to_string()
896        };
897
898        let storage = create_storage_backend_from_path(staging.to_str().unwrap(), None).unwrap();
899        let staged = block_on_nautilus_with(|| storage.list_files("", Some(".feather"))).unwrap();
900        assert!(error.contains("non-disjoint intervals"), "{error}");
901        assert_eq!(staged.len(), 1);
902    }
903
904    #[rstest]
905    fn parquet_manual_conversion_promotes_custom_data() {
906        use nautilus_model::data::{CustomData, DataType};
907        use nautilus_serialization::ensure_custom_data_registered;
908
909        ensure_custom_data_registered::<RustTestHashMapCustomData>();
910
911        let directory = TempDir::new().unwrap();
912        let staging = directory.path().join("backtest").join("run-custom");
913        let mut config = WriterConnectConfig::new(staging.to_string_lossy(), None);
914        config.params = Some(serde_json::from_value(json!({"promote_on_close": false})).unwrap());
915        let mut sink =
916            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
917                .unwrap();
918
919        let data_type = DataType::new("RustTestHashMapCustomData", None, None);
920        let records = [
921            sample_custom_data("first", "AUD/USD.SIM", "1.23456", 11),
922            sample_custom_data("second", "BTCUSDT.BINANCE", "65432.10", 21),
923        ];
924
925        for record in &records {
926            sink.write_data(Data::Custom(CustomData::new(
927                Arc::new(record.clone()),
928                data_type.clone(),
929            )))
930            .unwrap();
931        }
932
933        sink.close().unwrap();
934
935        let mut catalog = ParquetDataCatalog::from_uri(
936            directory.path().to_str().unwrap(),
937            None,
938            None,
939            None,
940            None,
941        )
942        .unwrap();
943
944        for _ in 0..2 {
945            catalog
946                .convert_stream_to_data(
947                    "run-custom",
948                    &NautilusDataType::Custom {
949                        type_name: "RustTestHashMapCustomData".to_string(),
950                    }
951                    .into(),
952                    Some("backtest"),
953                    None,
954                    false,
955                )
956                .unwrap();
957        }
958
959        let expected: Vec<(Option<String>, RustTestHashMapCustomData)> = records
960            .iter()
961            .map(|record| (None, record.clone()))
962            .collect();
963        assert_eq!(query_custom_records(&mut catalog, None), expected);
964    }
965
966    #[rstest]
967    fn parquet_manual_conversion_filters_custom_data_identifiers() {
968        use nautilus_model::data::{CustomData, DataType};
969        use nautilus_serialization::ensure_custom_data_registered;
970
971        ensure_custom_data_registered::<RustTestHashMapCustomData>();
972
973        let directory = TempDir::new().unwrap();
974        let staging = directory.path().join("backtest").join("run-custom-ids");
975        let mut config = WriterConnectConfig::new(staging.to_string_lossy(), None);
976        config.params = Some(serde_json::from_value(json!({"promote_on_close": false})).unwrap());
977        let mut sink =
978            parquet_writer_factory(&config, WriterClock::Test(Arc::new(AtomicU64::new(0))))
979                .unwrap();
980
981        let audusd = "AUD/USD.SIM";
982        let btcusdt = "BTCUSDT.BINANCE";
983        let records = [
984            (audusd, sample_custom_data("first", audusd, "1.23456", 11)),
985            (audusd, sample_custom_data("second", audusd, "1.23457", 21)),
986            (
987                btcusdt,
988                sample_custom_data("third", btcusdt, "65432.10", 31),
989            ),
990        ];
991
992        for (identifier, record) in &records {
993            let data_type = DataType::new(
994                "RustTestHashMapCustomData",
995                None,
996                Some((*identifier).to_string()),
997            );
998            sink.write_data(Data::Custom(CustomData::new(
999                Arc::new(record.clone()),
1000                data_type,
1001            )))
1002            .unwrap();
1003        }
1004
1005        sink.close().unwrap();
1006
1007        let mut catalog = ParquetDataCatalog::from_uri(
1008            directory.path().to_str().unwrap(),
1009            None,
1010            None,
1011            None,
1012            None,
1013        )
1014        .unwrap();
1015        catalog
1016            .convert_stream_to_data(
1017                "run-custom-ids",
1018                &NautilusDataType::Custom {
1019                    type_name: "RustTestHashMapCustomData".to_string(),
1020                }
1021                .into(),
1022                Some("backtest"),
1023                Some(&[audusd.to_string()]),
1024                false,
1025            )
1026            .unwrap();
1027
1028        let expected_audusd: Vec<(Option<String>, RustTestHashMapCustomData)> = records
1029            .iter()
1030            .filter(|(identifier, _)| *identifier == audusd)
1031            .map(|(identifier, record)| (Some((*identifier).to_string()), record.clone()))
1032            .collect();
1033        assert_eq!(query_custom_records(&mut catalog, None), expected_audusd);
1034
1035        // Catalog identifier directories are urisafe, so a scoped query pins that spelling and
1036        // proves the converted records landed under their own identifier rather than together.
1037        assert_eq!(
1038            query_custom_records(&mut catalog, Some(&["AUDUSD.SIM".to_string()])),
1039            expected_audusd,
1040        );
1041
1042        catalog
1043            .convert_stream_to_data(
1044                "run-custom-ids",
1045                &NautilusDataType::Custom {
1046                    type_name: "RustTestHashMapCustomData".to_string(),
1047                }
1048                .into(),
1049                Some("backtest"),
1050                None,
1051                false,
1052            )
1053            .unwrap();
1054
1055        let expected_all: Vec<(Option<String>, RustTestHashMapCustomData)> = records
1056            .iter()
1057            .map(|(identifier, record)| (Some((*identifier).to_string()), record.clone()))
1058            .collect();
1059        assert_eq!(query_custom_records(&mut catalog, None), expected_all);
1060        let expected_btcusdt: Vec<(Option<String>, RustTestHashMapCustomData)> = records
1061            .iter()
1062            .filter(|(identifier, _)| *identifier == btcusdt)
1063            .map(|(identifier, record)| (Some((*identifier).to_string()), record.clone()))
1064            .collect();
1065        assert_eq!(
1066            query_custom_records(&mut catalog, Some(&[btcusdt.to_string()])),
1067            expected_btcusdt,
1068        );
1069    }
1070
1071    fn sample_custom_data(
1072        name: &str,
1073        instrument_id: &str,
1074        price: &str,
1075        ts: u64,
1076    ) -> RustTestHashMapCustomData {
1077        RustTestHashMapCustomData {
1078            name: name.to_string(),
1079            prices: [(instrument_id.to_string(), Price::from(price))]
1080                .into_iter()
1081                .collect(),
1082            ts_event: UnixNanos::from(ts - 1),
1083            ts_init: UnixNanos::from(ts),
1084        }
1085    }
1086
1087    fn query_custom_records(
1088        catalog: &mut ParquetDataCatalog,
1089        identifiers: Option<&[String]>,
1090    ) -> Vec<(Option<String>, RustTestHashMapCustomData)> {
1091        let loaded = catalog
1092            .query_custom_data_dynamic(
1093                "RustTestHashMapCustomData",
1094                identifiers,
1095                None,
1096                None,
1097                None,
1098                None,
1099                true,
1100            )
1101            .unwrap();
1102        let mut decoded = Vec::new();
1103
1104        for data in &loaded {
1105            let Data::Custom(custom) = data else {
1106                panic!("expected custom data, found {data:?}")
1107            };
1108
1109            assert_eq!(custom.data_type.type_name(), "RustTestHashMapCustomData");
1110            decoded.push((
1111                custom.data_type.identifier().map(ToString::to_string),
1112                custom
1113                    .data
1114                    .as_any()
1115                    .downcast_ref::<RustTestHashMapCustomData>()
1116                    .expect("expected RustTestHashMapCustomData")
1117                    .clone(),
1118            ));
1119        }
1120
1121        decoded.sort_by_key(|(_, record)| record.ts_init);
1122        decoded
1123    }
1124
1125    fn sample_quote() -> QuoteTick {
1126        QuoteTick::new(
1127            InstrumentId::from("AUD/USD.SIM"),
1128            Price::from("0.65"),
1129            Price::from("0.67"),
1130            Quantity::from("11"),
1131            Quantity::from("17"),
1132            UnixNanos::from(19),
1133            UnixNanos::from(23),
1134        )
1135    }
1136}
nautilus_persistence/backend/parquet/writer.rs

nautilus_persistence/backend/parquet/
writer.rs
