1#![expect(
17 clippy::match_same_arms,
18 clippy::too_many_arguments,
19 reason = "PyO3 catalog wrapper mirrors Python API dispatch surface"
20)]
21
22use std::collections::HashMap;
23
24use nautilus_core::{
25 UnixNanos,
26 python::{to_pytype_err, to_pyvalue_err},
27};
28use nautilus_model::{
29 data::{
30 Bar, Data, FundingRateUpdate, HasTsInit, IndexPriceUpdate, InstrumentStatus,
31 MarkPriceUpdate, NautilusDataType, OptionGreeks, OrderBookDelta, OrderBookDepth, QuoteTick,
32 TradeTick, close::InstrumentClose,
33 },
34 python::{
35 data::data_to_pyobject,
36 instruments::{
37 PyNautilusInstrumentType, instrument_any_to_pyobject, pyobject_to_instrument_any,
38 },
39 },
40};
41use nautilus_serialization::{
42 arrow::{
43 DecodeTypedFromRecordBatch, EncodeToRecordBatch, display::instrument::encode_instruments,
44 },
45 python::arrow::{arrow_record_batches_to_pyarrow_stream, arrow_record_batches_to_pybytes},
46};
47use pyo3::{
48 exceptions::PyIOError,
49 prelude::*,
50 types::{PyBytes, PyDict, PyList},
51};
52
53use crate::{
54 backend::{migration::build_catalog_migration_plan, parquet::catalog::ParquetDataCatalog},
55 catalog::{
56 traits::{CatalogQuery, CatalogReader, CatalogRecordQuery, CatalogWriter},
57 types::{
58 HasCatalogDataType, custom_data_read_prefixes, custom_type_name,
59 parquet_catalog_data_type_path_prefixes,
60 },
61 },
62 python::backend::{
63 PyCatalogDataType, arrow_ipc_batches, arrow_ipc_data_schema, arrow_ipc_record_schema,
64 arrow_record_batches_from_pybytes, catalog_metadata_to_pydict, catalog_record_type_from_py,
65 nautilus_data_type_from_py, to_pyio_err, write_record_params_from_py,
66 },
67};
68
69#[expect(
70 clippy::needless_pass_by_value,
71 reason = "PyO3 supplies owned Python data at the catalog write boundary"
72)]
73fn write_parquet_data<T>(
74 catalog: &ParquetDataCatalog,
75 data: Vec<T>,
76 start: Option<u64>,
77 end: Option<u64>,
78 skip_disjoint_check: bool,
79 label: &str,
80) -> PyResult<String>
81where
82 T: HasTsInit + EncodeToRecordBatch + HasCatalogDataType,
83{
84 catalog
85 .write_to_parquet(
86 &data,
87 start.map(UnixNanos::from),
88 end.map(UnixNanos::from),
89 Some(skip_disjoint_check),
90 )
91 .map(|path| path.to_string_lossy().to_string())
92 .map_err(|e| PyIOError::new_err(format!("Failed to write {label}: {e}")))
93}
94
95fn query_parquet_data<T>(
96 catalog: &mut ParquetDataCatalog,
97 identifiers: Option<Vec<String>>,
98 start: Option<u64>,
99 end: Option<u64>,
100 where_clause: Option<&str>,
101 files: Option<Vec<String>>,
102 optimize_file_loading: bool,
103 error_context: &str,
104) -> PyResult<Vec<T>>
105where
106 T: DecodeTypedFromRecordBatch + HasCatalogDataType + HasTsInit,
107{
108 catalog
109 .query_typed_data::<T>(
110 identifiers,
111 start.map(UnixNanos::from),
112 end.map(UnixNanos::from),
113 where_clause,
114 files,
115 optimize_file_loading,
116 )
117 .map_err(|e| PyIOError::new_err(format!("{error_context}: {e}")))
118}
119
120fn reject_parquet_as_of(as_of: Option<&Bound<'_, PyAny>>) -> PyResult<()> {
121 if as_of.is_some() {
122 return Err(to_pyvalue_err("ParquetDataCatalog does not support as_of"));
123 }
124 Ok(())
125}
126
127#[pyclass(name = "ParquetDataCatalog", module = "nautilus_trader.persistence")]
129#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
130pub struct PyParquetDataCatalog {
131 inner: ParquetDataCatalog,
132}
133
134#[pymethods]
135#[pyo3_stub_gen::derive::gen_stub_pymethods]
136impl PyParquetDataCatalog {
137 #[new]
151 #[pyo3(signature = (base_path, storage_options=None, batch_size=None, compression=None, max_row_group_size=None))]
152 pub fn py_new(
153 base_path: &str,
154 storage_options: Option<HashMap<String, String>>,
155 batch_size: Option<usize>,
156 compression: Option<u8>,
157 max_row_group_size: Option<usize>,
158 ) -> PyResult<Self> {
159 let compression = compression.map(|c| match c {
160 0 => parquet::basic::Compression::UNCOMPRESSED,
161 1 => parquet::basic::Compression::SNAPPY,
162 2 => {
165 let level = parquet::basic::GzipLevel::default();
166 parquet::basic::Compression::GZIP(level)
167 }
168 3 => parquet::basic::Compression::LZO,
169 4 => {
170 let level = parquet::basic::BrotliLevel::default();
171 parquet::basic::Compression::BROTLI(level)
172 }
173 5 => parquet::basic::Compression::LZ4,
174 6 => {
175 let level = parquet::basic::ZstdLevel::default();
176 parquet::basic::Compression::ZSTD(level)
177 }
178 _ => parquet::basic::Compression::SNAPPY,
179 });
180
181 let storage_options = storage_options.map(|m| m.into_iter().collect());
183
184 let inner = ParquetDataCatalog::from_uri(
185 base_path,
186 storage_options,
187 batch_size,
188 compression,
189 max_row_group_size,
190 )
191 .map_err(|e| PyIOError::new_err(format!("Failed to create ParquetDataCatalog: {e}")))?;
192
193 Ok(Self { inner })
194 }
195
196 #[pyo3(signature = (parquet_path, storage_options=None, dry_run=false))]
198 pub fn migrate_from_legacy_parquet_path(
199 mut slf: PyRefMut<'_, Self>,
200 parquet_path: &str,
201 storage_options: Option<HashMap<String, String>>,
202 dry_run: bool,
203 ) -> PyResult<usize> {
204 let storage_options = storage_options.map(|m| m.into_iter().collect());
205 let source = ParquetDataCatalog::from_uri(parquet_path, storage_options, None, None, None)
206 .map_err(to_pyio_err)?;
207 let py = slf.py();
208 if dry_run {
209 return py
210 .detach(|| {
211 let plan = build_catalog_migration_plan(&source)?;
212 plan.ensure_ready()?;
213 Ok::<usize, anyhow::Error>(0)
214 })
215 .map_err(to_pyio_err);
216 }
217 let inner = &mut slf.inner;
218 py.detach(|| {
219 inner
220 .migrate_from_legacy_parquet_catalog(&source)
221 .map(|report| report.migrated_rows)
222 })
223 .map_err(to_pyio_err)
224 }
225
226 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
238 pub fn write_quote_ticks(
239 &self,
240 data: Vec<QuoteTick>,
241 start: Option<u64>,
242 end: Option<u64>,
243 skip_disjoint_check: bool,
244 ) -> PyResult<String> {
245 write_parquet_data(
246 &self.inner,
247 data,
248 start,
249 end,
250 skip_disjoint_check,
251 "quote ticks",
252 )
253 }
254
255 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
267 pub fn write_trade_ticks(
268 &self,
269 data: Vec<TradeTick>,
270 start: Option<u64>,
271 end: Option<u64>,
272 skip_disjoint_check: bool,
273 ) -> PyResult<String> {
274 write_parquet_data(
275 &self.inner,
276 data,
277 start,
278 end,
279 skip_disjoint_check,
280 "trade ticks",
281 )
282 }
283
284 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
296 pub fn write_order_book_deltas(
297 &self,
298 data: Vec<OrderBookDelta>,
299 start: Option<u64>,
300 end: Option<u64>,
301 skip_disjoint_check: bool,
302 ) -> PyResult<String> {
303 write_parquet_data(
304 &self.inner,
305 data,
306 start,
307 end,
308 skip_disjoint_check,
309 "order book deltas",
310 )
311 }
312
313 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
325 pub fn write_bars(
326 &self,
327 data: Vec<Bar>,
328 start: Option<u64>,
329 end: Option<u64>,
330 skip_disjoint_check: bool,
331 ) -> PyResult<String> {
332 write_parquet_data(&self.inner, data, start, end, skip_disjoint_check, "bars")
333 }
334
335 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
347 pub fn write_order_book_depths(
348 &self,
349 data: Vec<OrderBookDepth>,
350 start: Option<u64>,
351 end: Option<u64>,
352 skip_disjoint_check: bool,
353 ) -> PyResult<String> {
354 write_parquet_data(
355 &self.inner,
356 data,
357 start,
358 end,
359 skip_disjoint_check,
360 "order book depths",
361 )
362 }
363
364 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
376 pub fn write_mark_price_updates(
377 &self,
378 data: Vec<MarkPriceUpdate>,
379 start: Option<u64>,
380 end: Option<u64>,
381 skip_disjoint_check: bool,
382 ) -> PyResult<String> {
383 write_parquet_data(
384 &self.inner,
385 data,
386 start,
387 end,
388 skip_disjoint_check,
389 "mark price updates",
390 )
391 }
392
393 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
405 pub fn write_index_price_updates(
406 &self,
407 data: Vec<IndexPriceUpdate>,
408 start: Option<u64>,
409 end: Option<u64>,
410 skip_disjoint_check: bool,
411 ) -> PyResult<String> {
412 write_parquet_data(
413 &self.inner,
414 data,
415 start,
416 end,
417 skip_disjoint_check,
418 "index price updates",
419 )
420 }
421
422 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
434 pub fn write_option_greeks(
435 &self,
436 data: Vec<OptionGreeks>,
437 start: Option<u64>,
438 end: Option<u64>,
439 skip_disjoint_check: bool,
440 ) -> PyResult<String> {
441 write_parquet_data(
442 &self.inner,
443 data,
444 start,
445 end,
446 skip_disjoint_check,
447 "option greeks",
448 )
449 }
450
451 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
453 pub fn write_instrument_statuses(
454 &self,
455 data: Vec<InstrumentStatus>,
456 start: Option<u64>,
457 end: Option<u64>,
458 skip_disjoint_check: bool,
459 ) -> PyResult<String> {
460 write_parquet_data(
461 &self.inner,
462 data,
463 start,
464 end,
465 skip_disjoint_check,
466 "instrument statuses",
467 )
468 }
469
470 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
472 pub fn write_instrument_closes(
473 &self,
474 data: Vec<InstrumentClose>,
475 start: Option<u64>,
476 end: Option<u64>,
477 skip_disjoint_check: bool,
478 ) -> PyResult<String> {
479 write_parquet_data(
480 &self.inner,
481 data,
482 start,
483 end,
484 skip_disjoint_check,
485 "instrument closes",
486 )
487 }
488
489 pub fn write_instruments(&self, instruments: &Bound<'_, PyAny>) -> PyResult<Vec<String>> {
503 let data = instruments;
504 let py = data.py();
505 let list = data.cast::<PyList>()?;
506 let mut instruments = Vec::with_capacity(list.len());
507 for item in list.iter() {
508 let py_item: Py<PyAny> = item.unbind();
509 let instrument = pyobject_to_instrument_any(py, py_item)?;
510 instruments.push(instrument);
511 }
512 self.inner
513 .write_instruments(instruments)
514 .map(|paths| {
515 paths
516 .into_iter()
517 .map(|p| p.to_string_lossy().to_string())
518 .collect()
519 })
520 .map_err(|e| PyIOError::new_err(format!("Failed to write instruments: {e}")))
521 }
522
523 #[pyo3(signature = (
536 instrument_ids=None,
537 start=None,
538 end=None,
539 where_clause=None,
540 instrument_type=None,
541 ))]
542 #[expect(clippy::needless_pass_by_value)]
543 pub fn instruments(
544 &mut self,
545 instrument_ids: Option<Vec<String>>,
546 start: Option<u64>,
547 end: Option<u64>,
548 where_clause: Option<&str>,
549 instrument_type: Option<PyRef<'_, PyNautilusInstrumentType>>,
550 ) -> PyResult<Vec<Py<PyAny>>> {
551 let instrument_type = instrument_type.map(|instrument_type| instrument_type.inner());
552 let rust_instruments = self
553 .inner
554 .query_instruments_filtered_with_where_and_type(
555 instrument_ids.as_deref(),
556 start.map(UnixNanos::from),
557 end.map(UnixNanos::from),
558 where_clause,
559 instrument_type.as_ref(),
560 )
561 .map_err(|e| PyIOError::new_err(format!("Failed to query instruments: {e}")))?;
562 Python::attach(|py| {
563 rust_instruments
564 .into_iter()
565 .map(|inst| instrument_any_to_pyobject(py, inst))
566 .collect()
567 })
568 }
569
570 #[pyo3(signature = (
572 instrument_ids=None,
573 start=None,
574 end=None,
575 where_clause=None,
576 instrument_type=None,
577 ))]
578 #[expect(clippy::needless_pass_by_value)]
579 pub fn query_instrument_arrow_bytes(
580 &mut self,
581 py: Python<'_>,
582 instrument_ids: Option<Vec<String>>,
583 start: Option<u64>,
584 end: Option<u64>,
585 where_clause: Option<&str>,
586 instrument_type: Option<PyRef<'_, PyNautilusInstrumentType>>,
587 ) -> PyResult<Py<PyBytes>> {
588 let instrument_type = instrument_type.map(|instrument_type| instrument_type.inner());
589 let instruments = py
590 .detach(|| {
591 self.inner.query_instruments_filtered_with_where_and_type(
592 instrument_ids.as_deref(),
593 start.map(UnixNanos::from),
594 end.map(UnixNanos::from),
595 where_clause,
596 instrument_type.as_ref(),
597 )
598 })
599 .map_err(|e| PyIOError::new_err(format!("Failed query instruments: {e}")))?;
600 let batch = encode_instruments(&instruments)
601 .map_err(|e| PyIOError::new_err(format!("Failed encode instruments: {e}")))?;
602 let schema = batch.schema().as_ref().clone();
603 arrow_record_batches_to_pybytes(py, &schema, &[batch])
604 }
605
606 #[pyo3(signature = (
608 instrument_ids=None,
609 start=None,
610 end=None,
611 where_clause=None,
612 instrument_type=None,
613 ))]
614 #[expect(clippy::needless_pass_by_value)]
615 pub fn query_instrument_arrow_stream(
616 &mut self,
617 py: Python<'_>,
618 instrument_ids: Option<Vec<String>>,
619 start: Option<u64>,
620 end: Option<u64>,
621 where_clause: Option<&str>,
622 instrument_type: Option<PyRef<'_, PyNautilusInstrumentType>>,
623 ) -> PyResult<Py<PyAny>> {
624 let instrument_type = instrument_type.map(|instrument_type| instrument_type.inner());
625 let instruments = py
626 .detach(|| {
627 self.inner.query_instruments_filtered_with_where_and_type(
628 instrument_ids.as_deref(),
629 start.map(UnixNanos::from),
630 end.map(UnixNanos::from),
631 where_clause,
632 instrument_type.as_ref(),
633 )
634 })
635 .map_err(|e| PyIOError::new_err(format!("Failed query instruments: {e}")))?;
636 let batch = encode_instruments(&instruments)
637 .map_err(|e| PyIOError::new_err(format!("Failed encode instruments: {e}")))?;
638 let schema = batch.schema().as_ref().clone();
639 arrow_record_batches_to_pyarrow_stream(py, &schema, vec![batch])
640 }
641
642 #[pyo3(signature = (data_type, instrument_id=None, *, start, end))]
651 #[expect(clippy::needless_pass_by_value)]
652 pub fn extend_file_name(
653 &self,
654 data_type: PyCatalogDataType,
655 instrument_id: Option<String>,
656 start: u64,
657 end: u64,
658 ) -> PyResult<()> {
659 let data_type = data_type.into_inner();
660 let start_nanos = UnixNanos::from(start);
661 let end_nanos = UnixNanos::from(end);
662
663 self.inner
664 .extend_file_name(&data_type, instrument_id.as_deref(), start_nanos, end_nanos)
665 .map_err(|e| PyIOError::new_err(format!("Failed to extend file name: {e}")))
666 }
667
668 #[pyo3(signature = (start=None, end=None, ensure_contiguous_files=None, deduplicate=None))]
677 pub fn consolidate_catalog(
678 &self,
679 start: Option<u64>,
680 end: Option<u64>,
681 ensure_contiguous_files: Option<bool>,
682 deduplicate: Option<bool>,
683 ) -> PyResult<()> {
684 let start_nanos = start.map(UnixNanos::from);
685 let end_nanos = end.map(UnixNanos::from);
686
687 self.inner
688 .consolidate_catalog(start_nanos, end_nanos, ensure_contiguous_files, deduplicate)
689 .map_err(|e| PyIOError::new_err(format!("Failed to consolidate catalog: {e}")))
690 }
691
692 #[pyo3(signature = (data_type, instrument_id=None, start=None, end=None, ensure_contiguous_files=None, deduplicate=None))]
703 #[expect(clippy::needless_pass_by_value)]
704 pub fn consolidate_data(
705 &mut self,
706 data_type: PyCatalogDataType,
707 instrument_id: Option<String>,
708 start: Option<u64>,
709 end: Option<u64>,
710 ensure_contiguous_files: Option<bool>,
711 deduplicate: Option<bool>,
712 ) -> PyResult<()> {
713 let data_type = data_type.into_inner();
714 let start_nanos = start.map(UnixNanos::from);
715 let end_nanos = end.map(UnixNanos::from);
716
717 self.inner
718 .consolidate_data(
719 &data_type,
720 instrument_id.as_deref(),
721 start_nanos,
722 end_nanos,
723 ensure_contiguous_files,
724 deduplicate,
725 )
726 .map_err(|e| PyIOError::new_err(format!("Failed to consolidate data: {e}")))
727 }
728
729 #[pyo3(signature = (period_nanos=None, start=None, end=None, ensure_contiguous_files=None))]
744 pub fn consolidate_catalog_by_period(
745 &mut self,
746 period_nanos: Option<u64>,
747 start: Option<u64>,
748 end: Option<u64>,
749 ensure_contiguous_files: Option<bool>,
750 ) -> PyResult<()> {
751 let start_nanos = start.map(UnixNanos::from);
752 let end_nanos = end.map(UnixNanos::from);
753
754 self.inner
755 .consolidate_catalog_by_period(
756 period_nanos,
757 start_nanos,
758 end_nanos,
759 ensure_contiguous_files,
760 )
761 .map_err(|e| {
762 PyIOError::new_err(format!("Failed to consolidate catalog by period: {e}"))
763 })
764 }
765
766 #[pyo3(signature = (data_type, identifier=None, period_nanos=None, start=None, end=None, ensure_contiguous_files=None))]
782 #[expect(clippy::needless_pass_by_value)]
783 pub fn consolidate_data_by_period(
784 &mut self,
785 data_type: PyCatalogDataType,
786 identifier: Option<String>,
787 period_nanos: Option<u64>,
788 start: Option<u64>,
789 end: Option<u64>,
790 ensure_contiguous_files: Option<bool>,
791 ) -> PyResult<()> {
792 let data_type = data_type.into_inner();
793 let start_nanos = start.map(UnixNanos::from);
794 let end_nanos = end.map(UnixNanos::from);
795
796 self.inner
797 .consolidate_data_by_period(
798 &data_type,
799 identifier.as_deref(),
800 period_nanos,
801 start_nanos,
802 end_nanos,
803 ensure_contiguous_files,
804 )
805 .map_err(|e| PyIOError::new_err(format!("Failed to consolidate data by period: {e}")))
806 }
807
808 pub fn reset_all_file_names(&self) -> PyResult<()> {
810 self.inner
811 .reset_all_file_names()
812 .map_err(|e| PyIOError::new_err(format!("Failed to reset catalog file names: {e}")))
813 }
814
815 #[pyo3(signature = (data_type, instrument_id=None))]
822 #[expect(clippy::needless_pass_by_value)]
823 pub fn reset_data_file_names(
824 &self,
825 data_type: PyCatalogDataType,
826 instrument_id: Option<String>,
827 ) -> PyResult<()> {
828 let data_type = data_type.into_inner();
829 self.inner
830 .reset_data_file_names(&data_type, instrument_id.as_deref())
831 .map_err(|e| PyIOError::new_err(format!("Failed to reset data file names: {e}")))
832 }
833
834 #[pyo3(signature = (start=None, end=None))]
856 pub fn delete_catalog_range(&mut self, start: Option<u64>, end: Option<u64>) -> PyResult<()> {
857 let start_nanos = start.map(UnixNanos::from);
858 let end_nanos = end.map(UnixNanos::from);
859
860 self.inner
861 .delete_catalog_range(start_nanos, end_nanos)
862 .map_err(|e| PyIOError::new_err(format!("Failed to delete catalog range: {e}")))
863 }
864
865 #[pyo3(signature = (data_type, identifier=None, start=None, end=None))]
887 #[expect(clippy::needless_pass_by_value)]
888 pub fn delete_data_range(
889 &mut self,
890 #[gen_stub(override_type(type_repr = "model.NautilusDataType"))] data_type: &Bound<
891 '_,
892 PyAny,
893 >,
894 identifier: Option<String>,
895 start: Option<u64>,
896 end: Option<u64>,
897 ) -> PyResult<()> {
898 let data_type = nautilus_data_type_from_py(data_type)?;
899 let start_nanos = start.map(UnixNanos::from);
900 let end_nanos = end.map(UnixNanos::from);
901
902 self.inner
903 .delete_data_range(&data_type, identifier.as_deref(), start_nanos, end_nanos)
904 .map_err(|e| PyIOError::new_err(format!("Failed to delete data range: {e}")))
905 }
906
907 #[pyo3(signature = (data, start=None, end=None, skip_disjoint_check=false))]
915 pub fn write_custom_data(
916 &self,
917 _py: Python<'_>,
918 data: Vec<Bound<'_, PyAny>>,
919 start: Option<u64>,
920 end: Option<u64>,
921 skip_disjoint_check: bool,
922 ) -> PyResult<String> {
923 use nautilus_model::data::CustomData;
924
925 let mut custom_items: Vec<CustomData> = Vec::with_capacity(data.len());
926 for obj in data {
927 let custom = obj.extract::<CustomData>().map_err(|_| {
928 to_pytype_err(
929 "write_custom_data requires CustomData wrappers; wrap with CustomData(data_type=DataType(cls, metadata=...), data=...)",
930 )
931 })?;
932 custom_items.push(custom);
933 }
934
935 let start_nanos = start.map(UnixNanos::from);
936 let end_nanos = end.map(UnixNanos::from);
937
938 self.inner
939 .write_custom_data_batch(
940 &custom_items,
941 start_nanos,
942 end_nanos,
943 Some(skip_disjoint_check),
944 )
945 .map(|path| path.to_string_lossy().to_string())
946 .map_err(|e| PyIOError::new_err(format!("Failed to write custom data: {e}")))
947 }
948
949 pub fn list_instruments(&self, data_type: PyCatalogDataType) -> PyResult<Vec<String>> {
951 let data_type = data_type.into_inner();
952 self.inner
953 .list_instruments(&data_type)
954 .map_err(|e| PyIOError::new_err(format!("Failed to list instruments: {e}")))
955 }
956
957 pub fn list_parquet_files(
959 &self,
960 data_type: PyCatalogDataType,
961 instrument_id: &str,
962 ) -> PyResult<Vec<String>> {
963 let data_type = data_type.into_inner();
964 let mut files = Vec::new();
965
966 if let Some(type_name) = custom_type_name(&data_type) {
967 for prefix in custom_data_read_prefixes(type_name) {
968 let directory = format!("data/{prefix}/{instrument_id}");
969 files.extend(self.inner.list_parquet_files(&directory).map_err(|e| {
970 PyIOError::new_err(format!("Failed to list parquet files: {e}"))
971 })?);
972 }
973
974 files.sort();
975 files.dedup();
976 return Ok(files);
977 }
978
979 for prefix in parquet_catalog_data_type_path_prefixes(&data_type) {
980 let prefix = prefix.as_ref();
981 let directory = format!("data/{prefix}/{instrument_id}");
982 files.extend(
983 self.inner.list_parquet_files(&directory).map_err(|e| {
984 PyIOError::new_err(format!("Failed to list parquet files: {e}"))
985 })?,
986 );
987 }
988
989 Ok(files)
990 }
991
992 #[pyo3(signature = (data_type, identifiers=None, start=None, end=None))]
1007 pub fn query_files(
1008 &self,
1009 data_type: PyCatalogDataType,
1010 identifiers: Option<Vec<String>>,
1011 start: Option<u64>,
1012 end: Option<u64>,
1013 ) -> PyResult<Vec<String>> {
1014 let data_type = data_type.into_inner();
1015 let start_nanos = start.map(UnixNanos::from);
1016 let end_nanos = end.map(UnixNanos::from);
1017
1018 self.inner
1019 .query_files(&data_type, identifiers, start_nanos, end_nanos)
1020 .map_err(|e| PyIOError::new_err(format!("Failed to query files list: {e}")))
1021 }
1022
1023 #[pyo3(signature = (data_type, identifiers=None, start=None, end=None, where_clause=None))]
1025 pub fn query_metadata(
1026 &mut self,
1027 py: Python<'_>,
1028 #[gen_stub(override_type(type_repr = "model.NautilusDataType"))] data_type: &Bound<
1029 '_,
1030 PyAny,
1031 >,
1032 identifiers: Option<Vec<String>>,
1033 start: Option<u64>,
1034 end: Option<u64>,
1035 where_clause: Option<&str>,
1036 ) -> PyResult<Py<PyDict>> {
1037 let data_type = nautilus_data_type_from_py(data_type)?;
1038 let metadata = py
1039 .detach(|| {
1040 CatalogReader::query_metadata(
1041 &mut self.inner,
1042 &CatalogQuery::new(data_type)
1043 .with_identifiers(identifiers)
1044 .with_range(start.map(UnixNanos::from), end.map(UnixNanos::from))
1045 .with_where_clause(where_clause.map(str::to_string)),
1046 )
1047 })
1048 .map_err(|e| PyIOError::new_err(format!("Metadata query failed: {e}")))?;
1049
1050 catalog_metadata_to_pydict(py, metadata)
1051 }
1052
1053 #[pyo3(signature = (
1055 data_type,
1056 identifiers=None,
1057 start=None,
1058 end=None,
1059 where_clause=None,
1060 display=true,
1061 as_of=None,
1062 ))]
1063 pub fn query_data_arrow_bytes(
1064 &mut self,
1065 py: Python<'_>,
1066 #[gen_stub(override_type(type_repr = "model.NautilusDataType"))] data_type: &Bound<
1067 '_,
1068 PyAny,
1069 >,
1070 identifiers: Option<Vec<String>>,
1071 start: Option<u64>,
1072 end: Option<u64>,
1073 where_clause: Option<&str>,
1074 display: bool,
1075 as_of: Option<&Bound<'_, PyAny>>,
1076 ) -> PyResult<Py<PyBytes>> {
1077 let data_type = nautilus_data_type_from_py(data_type)?;
1078 reject_parquet_as_of(as_of)?;
1079 let query = CatalogQuery::new(data_type.clone())
1080 .with_identifiers(identifiers)
1081 .with_range(start.map(UnixNanos::from), end.map(UnixNanos::from))
1082 .with_where_clause(where_clause.map(str::to_string));
1083 let batches = py
1084 .detach(|| {
1085 if display {
1086 CatalogReader::query_display_record_batches(&mut self.inner, &query)
1087 } else {
1088 let data = CatalogReader::query_batch(&mut self.inner, &query)?
1089 .to_data_vec_for_compat();
1090 crate::common::arrow::data_to_arrow_batches(&data_type, data)
1091 }
1092 })
1093 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?;
1094 let schema = arrow_ipc_data_schema(&data_type, &batches, display)?;
1095 let batches = arrow_ipc_batches(&schema, batches)?;
1096 arrow_record_batches_to_pybytes(py, &schema, &batches)
1097 }
1098
1099 #[pyo3(signature = (
1101 data_type,
1102 identifiers=None,
1103 start=None,
1104 end=None,
1105 where_clause=None,
1106 display=true,
1107 as_of=None,
1108 ))]
1109 pub fn query_data_arrow_stream(
1110 &mut self,
1111 py: Python<'_>,
1112 #[gen_stub(override_type(type_repr = "model.NautilusDataType"))] data_type: &Bound<
1113 '_,
1114 PyAny,
1115 >,
1116 identifiers: Option<Vec<String>>,
1117 start: Option<u64>,
1118 end: Option<u64>,
1119 where_clause: Option<&str>,
1120 display: bool,
1121 as_of: Option<&Bound<'_, PyAny>>,
1122 ) -> PyResult<Py<PyAny>> {
1123 let data_type = nautilus_data_type_from_py(data_type)?;
1124 reject_parquet_as_of(as_of)?;
1125 let query = CatalogQuery::new(data_type.clone())
1126 .with_identifiers(identifiers)
1127 .with_range(start.map(UnixNanos::from), end.map(UnixNanos::from))
1128 .with_where_clause(where_clause.map(str::to_string));
1129 let batches = py
1130 .detach(|| {
1131 if display {
1132 CatalogReader::query_display_record_batches(&mut self.inner, &query)
1133 } else {
1134 let data = CatalogReader::query_batch(&mut self.inner, &query)?
1135 .to_data_vec_for_compat();
1136 crate::common::arrow::data_to_arrow_batches(&data_type, data)
1137 }
1138 })
1139 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?;
1140 let schema = arrow_ipc_data_schema(&data_type, &batches, display)?;
1141 let batches = arrow_ipc_batches(&schema, batches)?;
1142 arrow_record_batches_to_pyarrow_stream(py, &schema, batches)
1143 }
1144
1145 #[pyo3(signature = (record_type, data, identifier=None, params=None))]
1147 pub fn write_record_arrow_bytes(
1148 &mut self,
1149 py: Python<'_>,
1150 #[gen_stub(override_type(type_repr = "model.NautilusRecordType"))] record_type: &Bound<
1151 '_,
1152 PyAny,
1153 >,
1154 data: Vec<u8>,
1155 identifier: Option<String>,
1156 params: Option<Py<PyDict>>,
1157 ) -> PyResult<()> {
1158 let record_type = catalog_record_type_from_py(record_type)?;
1159 let batches = arrow_record_batches_from_pybytes(data)?;
1160 let params = write_record_params_from_py(py, identifier, params)?;
1161
1162 py.detach(|| CatalogWriter::write_records(&mut self.inner, record_type, &batches, params))
1163 .map_err(|e| PyIOError::new_err(format!("Failed write records: {e}")))
1164 }
1165
1166 #[pyo3(signature = (
1168 record_type,
1169 identifier=None,
1170 start=None,
1171 end=None,
1172 where_clause=None,
1173 display=true,
1174 as_of=None,
1175 ))]
1176 pub fn query_record_arrow_bytes(
1177 &mut self,
1178 py: Python<'_>,
1179 #[gen_stub(override_type(type_repr = "model.NautilusRecordType"))] record_type: &Bound<
1180 '_,
1181 PyAny,
1182 >,
1183 identifier: Option<String>,
1184 start: Option<u64>,
1185 end: Option<u64>,
1186 where_clause: Option<&str>,
1187 display: bool,
1188 as_of: Option<&Bound<'_, PyAny>>,
1189 ) -> PyResult<Py<PyBytes>> {
1190 let record_type = catalog_record_type_from_py(record_type)?;
1191 reject_parquet_as_of(as_of)?;
1192 let start_nanos = start.map(UnixNanos::from);
1193 let end_nanos = end.map(UnixNanos::from);
1194 let batches = py
1195 .detach(|| {
1196 if display {
1197 CatalogReader::query_record_display_batches(
1198 &mut self.inner,
1199 &CatalogRecordQuery::new(record_type)
1200 .with_identifier(identifier)
1201 .with_range(start_nanos, end_nanos)
1202 .with_where_clause(where_clause.map(str::to_string)),
1203 )
1204 } else {
1205 CatalogReader::query_record_batches(
1206 &mut self.inner,
1207 &CatalogRecordQuery::new(record_type)
1208 .with_identifier(identifier)
1209 .with_range(start_nanos, end_nanos)
1210 .with_where_clause(where_clause.map(str::to_string)),
1211 )
1212 }
1213 })
1214 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?;
1215 let schema = arrow_ipc_record_schema(record_type, &batches)?;
1216 let batches = arrow_ipc_batches(&schema, batches)?;
1217
1218 arrow_record_batches_to_pybytes(py, &schema, &batches)
1219 }
1220
1221 #[pyo3(signature = (
1223 record_type,
1224 identifier=None,
1225 start=None,
1226 end=None,
1227 where_clause=None,
1228 display=true,
1229 as_of=None,
1230 ))]
1231 pub fn query_record_arrow_stream(
1232 &mut self,
1233 py: Python<'_>,
1234 #[gen_stub(override_type(type_repr = "model.NautilusRecordType"))] record_type: &Bound<
1235 '_,
1236 PyAny,
1237 >,
1238 identifier: Option<String>,
1239 start: Option<u64>,
1240 end: Option<u64>,
1241 where_clause: Option<&str>,
1242 display: bool,
1243 as_of: Option<&Bound<'_, PyAny>>,
1244 ) -> PyResult<Py<PyAny>> {
1245 let record_type = catalog_record_type_from_py(record_type)?;
1246 reject_parquet_as_of(as_of)?;
1247 let start_nanos = start.map(UnixNanos::from);
1248 let end_nanos = end.map(UnixNanos::from);
1249 let batches = py
1250 .detach(|| {
1251 if display {
1252 CatalogReader::query_record_display_batches(
1253 &mut self.inner,
1254 &CatalogRecordQuery::new(record_type)
1255 .with_identifier(identifier)
1256 .with_range(start_nanos, end_nanos)
1257 .with_where_clause(where_clause.map(str::to_string)),
1258 )
1259 } else {
1260 CatalogReader::query_record_batches(
1261 &mut self.inner,
1262 &CatalogRecordQuery::new(record_type)
1263 .with_identifier(identifier)
1264 .with_range(start_nanos, end_nanos)
1265 .with_where_clause(where_clause.map(str::to_string)),
1266 )
1267 }
1268 })
1269 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?;
1270 let schema = arrow_ipc_record_schema(record_type, &batches)?;
1271 let batches = arrow_ipc_batches(&schema, batches)?;
1272
1273 arrow_record_batches_to_pyarrow_stream(py, &schema, batches)
1274 }
1275
1276 #[pyo3(signature = (start, end, data_type, instrument_id=None))]
1289 #[expect(clippy::needless_pass_by_value)]
1290 pub fn get_missing_intervals_for_request(
1291 &self,
1292 start: u64,
1293 end: u64,
1294 data_type: PyCatalogDataType,
1295 instrument_id: Option<String>,
1296 ) -> PyResult<Vec<(u64, u64)>> {
1297 let data_type = data_type.into_inner();
1298 self.inner
1299 .get_missing_intervals_for_request(start, end, &data_type, instrument_id.as_deref())
1300 .map_err(|e| PyIOError::new_err(format!("Failed to get missing intervals: {e}")))
1301 }
1302
1303 #[pyo3(signature = (data_type, instrument_id=None))]
1314 #[expect(clippy::needless_pass_by_value)]
1315 pub fn query_first_timestamp(
1316 &self,
1317 data_type: PyCatalogDataType,
1318 instrument_id: Option<String>,
1319 ) -> PyResult<Option<u64>> {
1320 let data_type = data_type.into_inner();
1321 self.inner
1322 .query_first_timestamp(&data_type, instrument_id.as_deref())
1323 .map_err(|e| PyIOError::new_err(format!("Failed to query first timestamp: {e}")))
1324 }
1325
1326 #[pyo3(signature = (data_type, instrument_id=None))]
1337 #[expect(clippy::needless_pass_by_value)]
1338 pub fn query_last_timestamp(
1339 &self,
1340 data_type: PyCatalogDataType,
1341 instrument_id: Option<String>,
1342 ) -> PyResult<Option<u64>> {
1343 let data_type = data_type.into_inner();
1344 self.inner
1345 .query_last_timestamp(&data_type, instrument_id.as_deref())
1346 .map_err(|e| PyIOError::new_err(format!("Failed to query last timestamp: {e}")))
1347 }
1348
1349 #[pyo3(signature = (data_type, instrument_id=None))]
1360 #[expect(clippy::needless_pass_by_value)]
1361 pub fn get_intervals(
1362 &self,
1363 data_type: PyCatalogDataType,
1364 instrument_id: Option<String>,
1365 ) -> PyResult<Vec<(u64, u64)>> {
1366 let data_type = data_type.into_inner();
1367 self.inner
1368 .get_intervals(&data_type, instrument_id.as_deref())
1369 .map_err(|e| PyIOError::new_err(format!("Failed to get intervals: {e}")))
1370 }
1371
1372 #[pyo3(signature = (data_type, identifiers=None, start=None, end=None, where_clause=None, files=None, optimize_file_loading=true))]
1374 #[expect(
1375 clippy::too_many_arguments,
1376 reason = "PyO3 signature mirrors the catalog query filters"
1377 )]
1378 pub fn query(
1379 &mut self,
1380 py: Python<'_>,
1381 #[gen_stub(override_type(type_repr = "model.NautilusDataType"))] data_type: &Bound<
1382 '_,
1383 PyAny,
1384 >,
1385 identifiers: Option<Vec<String>>,
1386 start: Option<u64>,
1387 end: Option<u64>,
1388 where_clause: Option<&str>,
1389 files: Option<Vec<String>>,
1390 optimize_file_loading: bool,
1391 ) -> PyResult<Vec<Py<PyAny>>> {
1392 let data_type = nautilus_data_type_from_py(data_type)?;
1393 let start_nanos = start.map(UnixNanos::from);
1394 let end_nanos = end.map(UnixNanos::from);
1395
1396 macro_rules! typed {
1397 ($type:ty) => {
1398 query_parquet_data::<$type>(
1399 &mut self.inner,
1400 identifiers,
1401 start,
1402 end,
1403 where_clause,
1404 files,
1405 optimize_file_loading,
1406 "Query failed",
1407 )?
1408 .into_iter()
1409 .map(Data::from)
1410 .collect::<Vec<Data>>()
1411 };
1412 }
1413
1414 let data = match data_type {
1415 NautilusDataType::QuoteTick => typed!(QuoteTick),
1416 NautilusDataType::TradeTick => typed!(TradeTick),
1417 NautilusDataType::Bar => typed!(Bar),
1418 NautilusDataType::OrderBookDelta => typed!(OrderBookDelta),
1419 NautilusDataType::OrderBookDepth => typed!(OrderBookDepth),
1420 NautilusDataType::IndexPriceUpdate => typed!(IndexPriceUpdate),
1421 NautilusDataType::MarkPriceUpdate => typed!(MarkPriceUpdate),
1422 NautilusDataType::FundingRateUpdate => typed!(FundingRateUpdate),
1423 NautilusDataType::OptionGreeks => typed!(OptionGreeks),
1424 NautilusDataType::InstrumentStatus => typed!(InstrumentStatus),
1425 NautilusDataType::InstrumentClose => typed!(InstrumentClose),
1426 NautilusDataType::Custom { type_name } => py
1427 .detach(|| {
1428 self.inner.query_custom_data_dynamic(
1429 &type_name,
1430 identifiers.as_deref(),
1431 start_nanos,
1432 end_nanos,
1433 where_clause,
1434 files.clone(),
1435 optimize_file_loading,
1436 )
1437 })
1438 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?,
1439 NautilusDataType::Instrument => {
1440 let instruments = py
1441 .detach(|| {
1442 self.inner.query_instruments_filtered_with_where(
1443 identifiers.as_deref(),
1444 start_nanos,
1445 end_nanos,
1446 where_clause,
1447 )
1448 })
1449 .map_err(|e| PyIOError::new_err(format!("Query failed: {e}")))?;
1450 return instruments
1451 .into_iter()
1452 .map(|instrument| instrument_any_to_pyobject(py, instrument))
1453 .collect();
1454 }
1455 #[cfg(feature = "defi")]
1456 NautilusDataType::Defi => {
1457 return Err(to_pytype_err("Defi data is not supported by query"));
1458 }
1459 };
1460
1461 let mut python_objects = Vec::new();
1462 for item in data {
1463 python_objects.push(data_to_pyobject(py, item)?);
1464 }
1465 Ok(python_objects)
1466 }
1467
1468 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1483 pub fn query_quote_ticks(
1484 &mut self,
1485 identifiers: Option<Vec<String>>,
1486 start: Option<u64>,
1487 end: Option<u64>,
1488 where_clause: Option<&str>,
1489 ) -> PyResult<Vec<QuoteTick>> {
1490 query_parquet_data::<QuoteTick>(
1491 &mut self.inner,
1492 identifiers,
1493 start,
1494 end,
1495 where_clause,
1496 None,
1497 true,
1498 "Failed to query data",
1499 )
1500 }
1501
1502 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1517 pub fn query_trade_ticks(
1518 &mut self,
1519 identifiers: Option<Vec<String>>,
1520 start: Option<u64>,
1521 end: Option<u64>,
1522 where_clause: Option<&str>,
1523 ) -> PyResult<Vec<TradeTick>> {
1524 query_parquet_data::<TradeTick>(
1525 &mut self.inner,
1526 identifiers,
1527 start,
1528 end,
1529 where_clause,
1530 None,
1531 true,
1532 "Failed to query data",
1533 )
1534 }
1535
1536 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1551 pub fn query_order_book_deltas(
1552 &mut self,
1553 identifiers: Option<Vec<String>>,
1554 start: Option<u64>,
1555 end: Option<u64>,
1556 where_clause: Option<&str>,
1557 ) -> PyResult<Vec<OrderBookDelta>> {
1558 query_parquet_data::<OrderBookDelta>(
1559 &mut self.inner,
1560 identifiers,
1561 start,
1562 end,
1563 where_clause,
1564 None,
1565 true,
1566 "Failed to query data",
1567 )
1568 }
1569
1570 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1585 pub fn query_bars(
1586 &mut self,
1587 identifiers: Option<Vec<String>>,
1588 start: Option<u64>,
1589 end: Option<u64>,
1590 where_clause: Option<&str>,
1591 ) -> PyResult<Vec<Bar>> {
1592 query_parquet_data::<Bar>(
1593 &mut self.inner,
1594 identifiers,
1595 start,
1596 end,
1597 where_clause,
1598 None,
1599 true,
1600 "Failed to query data",
1601 )
1602 }
1603
1604 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1617 pub fn query_order_book_depths(
1618 &mut self,
1619 identifiers: Option<Vec<String>>,
1620 start: Option<u64>,
1621 end: Option<u64>,
1622 where_clause: Option<&str>,
1623 ) -> PyResult<Vec<OrderBookDepth>> {
1624 query_parquet_data::<OrderBookDepth>(
1625 &mut self.inner,
1626 identifiers,
1627 start,
1628 end,
1629 where_clause,
1630 None,
1631 true,
1632 "Failed to query data",
1633 )
1634 }
1635
1636 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1649 pub fn query_mark_price_updates(
1650 &mut self,
1651 identifiers: Option<Vec<String>>,
1652 start: Option<u64>,
1653 end: Option<u64>,
1654 where_clause: Option<&str>,
1655 ) -> PyResult<Vec<MarkPriceUpdate>> {
1656 query_parquet_data::<MarkPriceUpdate>(
1657 &mut self.inner,
1658 identifiers,
1659 start,
1660 end,
1661 where_clause,
1662 None,
1663 true,
1664 "Failed to query data",
1665 )
1666 }
1667
1668 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1681 pub fn query_index_price_updates(
1682 &mut self,
1683 identifiers: Option<Vec<String>>,
1684 start: Option<u64>,
1685 end: Option<u64>,
1686 where_clause: Option<&str>,
1687 ) -> PyResult<Vec<IndexPriceUpdate>> {
1688 query_parquet_data::<IndexPriceUpdate>(
1689 &mut self.inner,
1690 identifiers,
1691 start,
1692 end,
1693 where_clause,
1694 None,
1695 true,
1696 "Failed to query data",
1697 )
1698 }
1699
1700 #[pyo3(signature = (identifiers=None, start=None, end=None, where_clause=None))]
1702 pub fn query_option_greeks(
1703 &mut self,
1704 identifiers: Option<Vec<String>>,
1705 start: Option<u64>,
1706 end: Option<u64>,
1707 where_clause: Option<&str>,
1708 ) -> PyResult<Vec<OptionGreeks>> {
1709 query_parquet_data::<OptionGreeks>(
1710 &mut self.inner,
1711 identifiers,
1712 start,
1713 end,
1714 where_clause,
1715 None,
1716 true,
1717 "Failed to query data",
1718 )
1719 }
1720
1721 pub fn list_data_types(&self) -> PyResult<Vec<String>> {
1727 self.inner
1728 .list_data_types()
1729 .map_err(|e| PyIOError::new_err(format!("Failed to list data types: {e}")))
1730 }
1731
1732 pub fn list_live_runs(&self) -> PyResult<Vec<String>> {
1738 self.inner
1739 .list_live_runs()
1740 .map_err(|e| PyIOError::new_err(format!("Failed to list live runs: {e}")))
1741 }
1742
1743 pub fn list_backtest_runs(&self) -> PyResult<Vec<String>> {
1749 self.inner
1750 .list_backtest_runs()
1751 .map_err(|e| PyIOError::new_err(format!("Failed to list backtest runs: {e}")))
1752 }
1753
1754 pub fn list_backtests(&self) -> PyResult<Vec<String>> {
1756 self.inner
1757 .list_backtest_runs()
1758 .map_err(|e| PyIOError::new_err(format!("Failed to list backtests: {e}")))
1759 }
1760
1761 #[pyo3(signature = (instance_id))]
1771 pub fn read_live_run(&self, py: Python<'_>, instance_id: &str) -> PyResult<Vec<Py<PyAny>>> {
1772 let data = self
1773 .inner
1774 .read_live_run(instance_id)
1775 .map_err(|e| PyIOError::new_err(format!("Failed to read live run: {e}")))?;
1776
1777 let mut python_objects = Vec::new();
1778 for item in data {
1779 python_objects.push(data_to_pyobject(py, item)?);
1780 }
1781 Ok(python_objects)
1782 }
1783
1784 #[pyo3(signature = (instance_id))]
1794 pub fn read_backtest(&self, py: Python<'_>, instance_id: &str) -> PyResult<Vec<Py<PyAny>>> {
1795 let data = self
1796 .inner
1797 .read_backtest(instance_id)
1798 .map_err(|e| PyIOError::new_err(format!("Failed to read backtest: {e}")))?;
1799
1800 let mut python_objects = Vec::new();
1801 for item in data {
1802 python_objects.push(data_to_pyobject(py, item)?);
1803 }
1804 Ok(python_objects)
1805 }
1806
1807 #[pyo3(signature = (instance_id, data_type, subdirectory=None, identifiers=None, use_ts_event_for_ts_init=false))]
1844 #[expect(clippy::needless_pass_by_value)]
1845 pub fn convert_stream_to_data(
1846 &mut self,
1847 instance_id: &str,
1848 data_type: PyCatalogDataType,
1849 subdirectory: Option<&str>,
1850 identifiers: Option<Vec<String>>,
1851 use_ts_event_for_ts_init: bool,
1852 ) -> PyResult<()> {
1853 let data_type = data_type.into_inner();
1854 let subdir = subdirectory.unwrap_or("backtest");
1855
1856 self.inner
1857 .convert_stream_to_data(
1858 instance_id,
1859 &data_type,
1860 Some(subdir),
1861 identifiers.as_deref(),
1862 use_ts_event_for_ts_init,
1863 )
1864 .map_err(|e| PyIOError::new_err(format!("Failed to convert stream to data: {e}")))
1865 }
1866
1867 #[pyo3(signature = (type_name, identifiers=None, start=None, end=None, where_clause=None))]
1869 #[expect(clippy::needless_pass_by_value)]
1870 pub fn query_custom_data(
1871 &mut self,
1872 py: Python<'_>,
1873 type_name: &str,
1874 identifiers: Option<Vec<String>>,
1875 start: Option<u64>,
1876 end: Option<u64>,
1877 where_clause: Option<&str>,
1878 ) -> PyResult<Vec<Py<PyAny>>> {
1879 let start_nanos = start.map(UnixNanos::from);
1880 let end_nanos = end.map(UnixNanos::from);
1881
1882 let data = py
1883 .detach(|| {
1884 self.inner.query_custom_data_dynamic(
1885 type_name,
1886 identifiers.as_deref(),
1887 start_nanos,
1888 end_nanos,
1889 where_clause,
1890 None,
1891 true,
1892 )
1893 })
1894 .map_err(|e| PyIOError::new_err(format!("Failed to query custom data: {e}")))?;
1895
1896 let mut python_objects = Vec::new();
1897
1898 for item in data {
1899 let py_obj: Py<PyAny> = match item {
1900 Data::Custom(custom) => Py::new(py, custom)?.into_any(),
1901 _ => return Err(PyIOError::new_err("Expected custom data")),
1902 };
1903 python_objects.push(py_obj);
1904 }
1905 Ok(python_objects)
1906 }
1907}