nautilus_persistence/python/backend/
session.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
16use std::collections::HashMap;
17
18use nautilus_core::python::{IntoPyObjectNautilusExt, to_pyruntime_err};
19use nautilus_model::{
20    data::{
21        Bar, Data, DataFFI, InstrumentStatus, MarkPriceUpdate, OptionGreeks, OrderBookDelta,
22        OrderBookDepth10, QuoteTick, TradeTick,
23    },
24    python::data::{DATA_FFI_CVEC_CAPSULE_NAME, DataFfiCVec},
25};
26use nautilus_serialization::arrow::{ArrowSchemaProvider, custom::CustomDataDecoder};
27use pyo3::{prelude::*, types::PyCapsule};
28
29use crate::backend::session::{DataBackendSession, DataQueryResult};
30
31/// Wrapper to pass a raw pointer across the GIL release boundary.
32struct SendPtr<T>(*mut T);
33
34// SAFETY: Access is serialized by the calling `PyRefMut`
35unsafe impl<T> Send for SendPtr<T> {}
36
37/// Converts a `Data` variant into a Python object via PyO3.
38#[allow(
39    clippy::match_wildcard_for_single_variants,
40    reason = "Data::Defi appears through nautilus-model feature unification"
41)]
42fn data_to_pyobject(py: Python<'_>, item: Data) -> PyResult<Py<PyAny>> {
43    match item {
44        Data::Quote(quote) => Py::new(py, quote).map(pyo3::Py::into_any),
45        Data::Trade(trade) => Py::new(py, trade).map(pyo3::Py::into_any),
46        Data::Bar(bar) => Py::new(py, bar).map(pyo3::Py::into_any),
47        Data::Delta(delta) => Py::new(py, delta).map(pyo3::Py::into_any),
48        Data::Deltas(deltas) => Py::new(py, (*deltas).clone()).map(pyo3::Py::into_any),
49        Data::Depth10(depth) => Py::new(py, *depth).map(pyo3::Py::into_any),
50        Data::IndexPriceUpdate(price) => Py::new(py, price).map(pyo3::Py::into_any),
51        Data::MarkPriceUpdate(price) => Py::new(py, price).map(pyo3::Py::into_any),
52        Data::FundingRateUpdate(funding_rate) => Py::new(py, funding_rate).map(pyo3::Py::into_any),
53        Data::OptionGreeks(greeks) => Py::new(py, greeks).map(pyo3::Py::into_any),
54        Data::InstrumentStatus(status) => Py::new(py, status).map(pyo3::Py::into_any),
55        Data::InstrumentClose(close) => Py::new(py, close).map(pyo3::Py::into_any),
56        Data::Custom(custom) => Py::new(py, custom).map(pyo3::Py::into_any),
57        #[cfg(feature = "defi")]
58        Data::Defi(_) => Err(to_pyruntime_err("Unsupported Data::Defi variant")),
59        #[allow(unreachable_patterns)]
60        _ => Err(to_pyruntime_err("Unsupported Data variant")),
61    }
62}
63
64#[repr(C)]
65#[pyclass(frozen, eq, eq_int, from_py_object)]
66#[pyo3_stub_gen::derive::gen_stub_pyclass_enum(module = "nautilus_trader.persistence")]
67#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
68pub enum NautilusDataType {
69    // Custom = 0,  # First slot reserved for custom data
70    OrderBookDelta = 1,
71    OrderBookDepth10 = 2,
72    QuoteTick = 3,
73    TradeTick = 4,
74    Bar = 5,
75    MarkPriceUpdate = 6,
76    OptionGreeks = 7,
77    InstrumentStatus = 8,
78}
79
80#[pymethods]
81#[pyo3_stub_gen::derive::gen_stub_pymethods]
82impl NautilusDataType {
83    #[expect(
84        clippy::trivially_copy_pass_by_ref,
85        reason = "PyO3 special methods use a borrowed receiver"
86    )]
87    const fn __hash__(&self) -> isize {
88        *self as isize
89    }
90}
91
92#[pymethods]
93#[pyo3_stub_gen::derive::gen_stub_pymethods]
94impl DataBackendSession {
95    #[new]
96    #[pyo3(signature=(chunk_size=10_000))]
97    fn new_session(chunk_size: usize) -> Self {
98        Self::new(chunk_size)
99    }
100
101    /// Registers a Parquet file and adds a batch stream for decoding.
102    ///
103    /// The caller must specify `T` to indicate the kind of data expected. `table_name` is
104    /// the logical name for queries; `file_path` is the Parquet path; `sql_query` defaults
105    /// to `SELECT * FROM {table_name} ORDER BY ts_init` if `None`.
106    ///
107    /// When `custom_type_name` is `Some`, it is merged into each batch's schema metadata
108    /// before decoding (as `type_name`). Use this for custom data when Parquet/DataFusion
109    /// does not preserve schema metadata so the decoder can look up the type in the registry.
110    ///
111    /// The file data must be ordered by the `ts_init` in ascending order for this
112    /// to work correctly.
113    #[pyo3(name = "add_file")]
114    #[pyo3(signature = (data_type, table_name, file_path, sql_query=None))]
115    fn py_add_file(
116        mut slf: PyRefMut<'_, Self>,
117        data_type: NautilusDataType,
118        table_name: &str,
119        file_path: &str,
120        sql_query: Option<&str>,
121    ) -> PyResult<()> {
122        let _guard = slf.runtime.enter();
123
124        match data_type {
125            NautilusDataType::OrderBookDelta => slf
126                .add_file::<OrderBookDelta>(table_name, file_path, sql_query, None)
127                .map_err(to_pyruntime_err),
128            NautilusDataType::OrderBookDepth10 => slf
129                .add_file::<OrderBookDepth10>(table_name, file_path, sql_query, None)
130                .map_err(to_pyruntime_err),
131            NautilusDataType::QuoteTick => slf
132                .add_file::<QuoteTick>(table_name, file_path, sql_query, None)
133                .map_err(to_pyruntime_err),
134            NautilusDataType::TradeTick => slf
135                .add_file::<TradeTick>(table_name, file_path, sql_query, None)
136                .map_err(to_pyruntime_err),
137            NautilusDataType::Bar => slf
138                .add_file::<Bar>(table_name, file_path, sql_query, None)
139                .map_err(to_pyruntime_err),
140            NautilusDataType::MarkPriceUpdate => slf
141                .add_file::<MarkPriceUpdate>(table_name, file_path, sql_query, None)
142                .map_err(to_pyruntime_err),
143            NautilusDataType::OptionGreeks => slf
144                .add_file::<OptionGreeks>(table_name, file_path, sql_query, None)
145                .map_err(to_pyruntime_err),
146            NautilusDataType::InstrumentStatus => slf
147                .add_file::<InstrumentStatus>(table_name, file_path, sql_query, None)
148                .map_err(to_pyruntime_err),
149        }
150    }
151
152    /// Registers a Parquet file for a custom data type identified by `type_name`.
153    ///
154    /// The custom data type must have been registered via
155    /// `ensure_custom_data_registered::<T>()` before calling this method.
156    #[pyo3(name = "add_custom_file")]
157    #[pyo3(signature = (type_name, table_name, file_path, sql_query=None))]
158    fn py_add_custom_file(
159        mut slf: PyRefMut<'_, Self>,
160        type_name: &str,
161        table_name: &str,
162        file_path: &str,
163        sql_query: Option<&str>,
164    ) -> PyResult<()> {
165        let _guard = slf.runtime.enter();
166        let mut metadata = HashMap::new();
167        metadata.insert("type_name".to_string(), type_name.to_string());
168        let base_schema = CustomDataDecoder::get_schema(Some(metadata));
169        base_schema.field_with_name("ts_init").map_err(|_| {
170            to_pyruntime_err(format!(
171                "custom data type '{type_name}' is not registered with an Arrow schema containing ts_init"
172            ))
173        })?;
174        // Use schemaless registration so DataFusion preserves the parquet file's
175        // schema metadata (e.g. `bar_type`) on output batches, since the
176        // explicit-schema variant strips per-batch metadata that decoders rely on.
177        slf.add_file::<CustomDataDecoder>(table_name, file_path, sql_query, Some(type_name))
178            .map_err(to_pyruntime_err)
179    }
180
181    fn to_query_result(mut slf: PyRefMut<'_, Self>) -> DataQueryResult {
182        let py = slf.py();
183        let chunk_size = slf.chunk_size;
184        let ptr = SendPtr(&raw mut *slf);
185
186        // SAFETY: see comment on `__next__` for the safety argument.
187        // The GIL release is needed here because `get_query_result` eagerly
188        // pulls the first element from each stream (via `KMerge::push_iter`),
189        // which blocks on the tokio channel while workers may need the GIL.
190        let query_result = unsafe {
191            py.detach(move || {
192                let p = ptr;
193                (*p.0).get_query_result()
194            })
195        };
196
197        DataQueryResult::new(query_result, chunk_size)
198    }
199
200    /// Register an object store with the session context from a URI with optional storage options.
201    #[pyo3(name = "register_object_store_from_uri")]
202    #[pyo3(signature = (uri, storage_options=None))]
203    fn py_register_object_store_from_uri(
204        mut slf: PyRefMut<'_, Self>,
205        uri: &str,
206        storage_options: Option<HashMap<String, String>>,
207    ) -> PyResult<()> {
208        // Convert HashMap to AHashMap for internal use
209        let storage_options = storage_options.map(|m| m.into_iter().collect());
210        slf.register_object_store_from_uri(uri, storage_options)
211            .map_err(to_pyruntime_err)
212    }
213}
214
215#[pymethods]
216#[pyo3_stub_gen::derive::gen_stub_pymethods]
217impl DataQueryResult {
218    /// The reader implements an iterator.
219    const fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
220        slf
221    }
222
223    /// Each iteration returns a chunk of values read from the parquet file.
224    ///
225    /// For built-in types, returns a `PyCapsule` containing a `CVec` of `DataFFI` (C layout)
226    /// consumed by Cython `capsule_to_list`. For custom data types (which are not
227    /// FFI-safe), returns a Python list of PyO3 objects directly.
228    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<Py<PyAny>>> {
229        let py = slf.py();
230        let ptr = SendPtr(&raw mut *slf);
231
232        // SAFETY: `PyRefMut` guarantees exclusive access to the underlying
233        // object for the duration of this method call. The runtime borrow
234        // flag prevents any other Python thread from accessing it.
235        //
236        // The GIL must be released here so that tokio worker threads can
237        // acquire it when decoding custom data types via `Python::attach`.
238        // Without this, custom-type streaming deadlocks: the main thread
239        // holds the GIL while blocking on `recv`, and workers block on
240        // `Python::attach` waiting for the GIL.
241        let acc = unsafe {
242            py.detach(move || {
243                let p = ptr;
244                (*p.0).next()
245            })
246        };
247
248        match acc {
249            Some(acc) if !acc.is_empty() => {
250                let has_non_ffi = acc.iter().any(|d| {
251                    matches!(
252                        d,
253                        Data::Custom(_)
254                            | Data::FundingRateUpdate(_)
255                            | Data::OptionGreeks(_)
256                            | Data::InstrumentStatus(_)
257                    )
258                });
259
260                if has_non_ffi {
261                    // Non-FFI data: convert directly to Python objects.
262                    let objects: Vec<Py<PyAny>> = acc
263                        .into_iter()
264                        .map(|item| data_to_pyobject(py, item))
265                        .collect::<PyResult<_>>()?;
266                    Ok(Some(objects.into_py_any_unwrap(py)))
267                } else {
268                    // Built-in types: FFI capsule path
269                    let ffi_data: Vec<DataFFI> = acc
270                        .into_iter()
271                        .map(DataFFI::try_from)
272                        .collect::<Result<Vec<_>, _>>()
273                        .map_err(to_pyruntime_err)?;
274                    let cvec: DataFfiCVec = ffi_data.into();
275                    match PyCapsule::new_with_value_and_destructor::<DataFfiCVec, _>(
276                        py,
277                        cvec,
278                        DATA_FFI_CVEC_CAPSULE_NAME,
279                        |_, _| {},
280                    ) {
281                        Ok(capsule) => Ok(Some(capsule.into_py_any_unwrap(py))),
282                        Err(e) => Err(to_pyruntime_err(e)),
283                    }
284                }
285            }
286            _ => Ok(None),
287        }
288    }
289}
nautilus_persistence/python/backend/session.rs

nautilus_persistence/python/backend/
session.rs
