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::{to_pyruntime_err, to_pyvalue_err};
19use nautilus_model::{
20    data::{
21        Bar, FundingRateUpdate, IndexPriceUpdate, InstrumentClose, InstrumentStatus,
22        MarkPriceUpdate, NautilusDataType, OptionGreeks, OrderBookDelta, OrderBookDepth, QuoteTick,
23        TradeTick,
24    },
25    python::data::data_to_pyobject,
26};
27use nautilus_serialization::arrow::{ArrowSchemaProvider, custom::CustomDataDecoder};
28use pyo3::{IntoPyObjectExt, prelude::*};
29
30use super::conversion::nautilus_data_type_from_py;
31use crate::backend::session::{DataBackendSession, DataQueryResult, QueryError};
32
33/// Wrapper to pass a raw pointer across the GIL release boundary.
34struct SendPtr<T>(*mut T);
35
36// SAFETY: Access is serialized by the calling `PyRefMut`
37unsafe impl<T> Send for SendPtr<T> {}
38
39#[pymethods]
40#[pyo3_stub_gen::derive::gen_stub_pymethods]
41impl DataBackendSession {
42    /// Provides a DataFusion session for registering and querying catalog table sources.
43    #[new]
44    #[pyo3(signature=(chunk_size=10_000))]
45    fn py_new(chunk_size: usize) -> PyResult<Self> {
46        if chunk_size == 0 {
47            return Err(to_pyvalue_err("chunk_size must be positive"));
48        }
49
50        Ok(Self::new(chunk_size))
51    }
52
53    /// Registers a Parquet file and adds a batch stream for decoding.
54    ///
55    /// The caller must specify `T` to indicate the kind of data expected. `table_name` is
56    /// the logical name for queries; `file_path` is the Parquet path; `sql_query` defaults
57    /// to `SELECT * FROM {table_name} ORDER BY ts_init` if `None`.
58    ///
59    /// When `custom_type_name` is `Some`, it is merged into each batch's schema metadata
60    /// before decoding (as `type_name`). Use this for custom data when Parquet/DataFusion
61    /// does not preserve schema metadata so the decoder can look up the type in the registry.
62    ///
63    /// The file data must be ordered by the `ts_init` in ascending order for this
64    /// to work correctly.
65    ///
66    /// # Errors
67    ///
68    /// Returns an error if parquet registration, SQL planning, stream execution, or
69    /// data decoding setup fails.
70    #[pyo3(name = "add_file")]
71    #[pyo3(signature = (data_type, table_name, file_path, sql_query=None))]
72    fn py_add_file(
73        mut slf: PyRefMut<'_, Self>,
74        data_type: &Bound<'_, PyAny>,
75        table_name: &str,
76        file_path: &str,
77        sql_query: Option<&str>,
78    ) -> PyResult<()> {
79        let _guard = slf.runtime.enter();
80        let data_type = nautilus_data_type_from_py(data_type)?;
81
82        slf.add_file_for_data_type(&data_type, table_name, file_path, sql_query)
83    }
84
85    /// Registers a Parquet file for a custom data type identified by `type_name`.
86    ///
87    /// The custom data type must have been registered via
88    /// `ensure_custom_data_registered::<T>()` before calling this method.
89    #[pyo3(name = "add_custom_file")]
90    #[pyo3(signature = (type_name, table_name, file_path, sql_query=None))]
91    fn py_add_custom_file(
92        mut slf: PyRefMut<'_, Self>,
93        type_name: &str,
94        table_name: &str,
95        file_path: &str,
96        sql_query: Option<&str>,
97    ) -> PyResult<()> {
98        let _guard = slf.runtime.enter();
99
100        let data_type = NautilusDataType::Custom {
101            type_name: type_name.to_string(),
102        };
103
104        slf.add_file_for_data_type(&data_type, table_name, file_path, sql_query)
105    }
106
107    fn to_query_result(mut slf: PyRefMut<'_, Self>) -> DataQueryResult {
108        let py = slf.py();
109        let chunk_size = slf.chunk_size;
110        let ptr = SendPtr(&raw mut *slf);
111
112        // SAFETY: see comment on `__next__` for the safety argument.
113        // The GIL release is needed here because `get_query_result` eagerly
114        // pulls the first element from each stream (via `KMerge::push_iter`),
115        // which blocks on the tokio channel while workers may need the GIL.
116        let query_result = unsafe {
117            py.detach(move || {
118                let p = ptr;
119                (*p.0).get_query_result()
120            })
121        };
122
123        DataQueryResult::new(query_result, chunk_size)
124    }
125
126    /// Register an object store with the session context from a URI with optional storage options.
127    ///
128    /// # Errors
129    ///
130    /// Returns an error if the object store URI cannot be normalized or the backend
131    /// cannot be created.
132    #[pyo3(name = "register_object_store_from_uri")]
133    #[pyo3(signature = (uri, storage_options=None))]
134    fn py_register_object_store_from_uri(
135        mut slf: PyRefMut<'_, Self>,
136        uri: &str,
137        storage_options: Option<HashMap<String, String>>,
138    ) -> PyResult<()> {
139        // Convert HashMap to AHashMap for internal use
140        let storage_options = storage_options.map(|m| m.into_iter().collect());
141        slf.register_object_store_from_uri(uri, storage_options)
142            .map_err(to_pyruntime_err)
143    }
144}
145
146impl DataBackendSession {
147    fn add_file_for_data_type(
148        &mut self,
149        data_type: &NautilusDataType,
150        table_name: &str,
151        file_path: &str,
152        sql_query: Option<&str>,
153    ) -> PyResult<()> {
154        match data_type {
155            NautilusDataType::OrderBookDelta => self
156                .add_file::<OrderBookDelta>(table_name, file_path, sql_query, None)
157                .map_err(to_pyruntime_err),
158            NautilusDataType::OrderBookDepth => self
159                .add_file::<OrderBookDepth>(table_name, file_path, sql_query, None)
160                .map_err(to_pyruntime_err),
161            NautilusDataType::QuoteTick => self
162                .add_file::<QuoteTick>(table_name, file_path, sql_query, None)
163                .map_err(to_pyruntime_err),
164            NautilusDataType::TradeTick => self
165                .add_file::<TradeTick>(table_name, file_path, sql_query, None)
166                .map_err(to_pyruntime_err),
167            NautilusDataType::Bar => self
168                .add_file::<Bar>(table_name, file_path, sql_query, None)
169                .map_err(to_pyruntime_err),
170            NautilusDataType::MarkPriceUpdate => self
171                .add_file::<MarkPriceUpdate>(table_name, file_path, sql_query, None)
172                .map_err(to_pyruntime_err),
173            NautilusDataType::IndexPriceUpdate => self
174                .add_file::<IndexPriceUpdate>(table_name, file_path, sql_query, None)
175                .map_err(to_pyruntime_err),
176            NautilusDataType::FundingRateUpdate => self
177                .add_file::<FundingRateUpdate>(table_name, file_path, sql_query, None)
178                .map_err(to_pyruntime_err),
179            NautilusDataType::OptionGreeks => self
180                .add_file::<OptionGreeks>(table_name, file_path, sql_query, None)
181                .map_err(to_pyruntime_err),
182            NautilusDataType::InstrumentStatus => self
183                .add_file::<InstrumentStatus>(table_name, file_path, sql_query, None)
184                .map_err(to_pyruntime_err),
185            NautilusDataType::InstrumentClose => self
186                .add_file::<InstrumentClose>(table_name, file_path, sql_query, None)
187                .map_err(to_pyruntime_err),
188            NautilusDataType::Custom { type_name } => {
189                let mut metadata = HashMap::new();
190                metadata.insert("type_name".to_string(), type_name.clone());
191                let base_schema = CustomDataDecoder::get_schema(Some(metadata));
192                base_schema.field_with_name("ts_init").map_err(|_| {
193                    to_pyruntime_err(format!(
194                        "custom data type '{type_name}' is not registered with an Arrow schema containing ts_init"
195                    ))
196                })?;
197
198                // Use schemaless registration so DataFusion preserves the parquet file's
199                // schema metadata (e.g. `bar_type`) on output batches, since the
200                // explicit-schema variant strips per-batch metadata that decoders rely on.
201                self.add_file::<CustomDataDecoder>(
202                    table_name,
203                    file_path,
204                    sql_query,
205                    Some(type_name.as_str()),
206                )
207                .map_err(to_pyruntime_err)
208            }
209            NautilusDataType::Instrument => Err(to_pyvalue_err(format!(
210                "DataBackendSession does not support data type {data_type}"
211            ))),
212            #[cfg(feature = "defi")]
213            NautilusDataType::Defi => Err(to_pyvalue_err(format!(
214                "DataBackendSession does not support data type {data_type}"
215            ))),
216        }
217    }
218}
219
220#[pymethods]
221#[pyo3_stub_gen::derive::gen_stub_pymethods]
222impl DataQueryResult {
223    /// Collects the remaining query records as native Python objects.
224    ///
225    /// # Errors
226    ///
227    /// Returns an error if a query stream fails or a record batch cannot be decoded.
228    #[pyo3(name = "to_list")]
229    fn py_to_list(mut slf: PyRefMut<'_, Self>) -> PyResult<Vec<Py<PyAny>>> {
230        let py = slf.py();
231        let ptr = SendPtr(&raw mut *slf);
232
233        // SAFETY: `PyRefMut` guarantees exclusive access to the underlying query result for the
234        // duration of this method call. As with `__next__`, release the GIL while waiting for
235        // decoder workers that may need to acquire it for custom data.
236        let data = unsafe {
237            py.detach(move || -> Result<Vec<_>, QueryError> {
238                let p = ptr;
239                let result = &mut *p.0;
240                let mut data = Vec::new();
241
242                for chunk in result.by_ref() {
243                    let chunk = chunk?;
244
245                    if chunk.is_empty() {
246                        break;
247                    }
248                    data.extend(chunk);
249                }
250
251                Ok(data)
252            })
253        }
254        .map_err(to_pyruntime_err)?;
255
256        data.into_iter()
257            .map(|item| data_to_pyobject(py, item))
258            .collect()
259    }
260
261    /// The reader implements an iterator.
262    const fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
263        slf
264    }
265
266    /// Each iteration returns a chunk of values read from the parquet file.
267    ///
268    /// # Errors
269    ///
270    /// Returns an error if a query stream fails or a record batch cannot be decoded.
271    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<Py<PyAny>>> {
272        let py = slf.py();
273        let ptr = SendPtr(&raw mut *slf);
274
275        // SAFETY: `PyRefMut` guarantees exclusive access to the underlying
276        // object for the duration of this method call. The runtime borrow
277        // flag prevents any other Python thread from accessing it.
278        //
279        // The GIL must be released here so that tokio worker threads can
280        // acquire it when decoding custom data types via `Python::attach`.
281        // Without this, custom-type streaming deadlocks: the main thread
282        // holds the GIL while blocking on `recv`, and workers block on
283        // `Python::attach` waiting for the GIL.
284        let acc = unsafe {
285            py.detach(move || {
286                let p = ptr;
287                (*p.0).next()
288            })
289        };
290
291        match acc {
292            Some(Ok(acc)) if !acc.is_empty() => {
293                let objects: Vec<Py<PyAny>> = acc
294                    .into_iter()
295                    .map(|item| data_to_pyobject(py, item))
296                    .collect::<PyResult<_>>()?;
297                Ok(Some(objects.into_py_any(py)?))
298            }
299            Some(Err(e)) => Err(to_pyruntime_err(e)),
300            _ => Ok(None),
301        }
302    }
303}
nautilus_persistence/python/backend/session.rs

nautilus_persistence/python/backend/
session.rs
