nautilus_persistence/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::{sync::Arc, vec::IntoIter};
17
18use ahash::{AHashMap, AHashSet};
19use datafusion::{
20    arrow::record_batch::RecordBatch, error::Result, logical_expr::expr::Sort,
21    physical_plan::SendableRecordBatchStream, prelude::*,
22};
23use futures::StreamExt;
24use nautilus_core::{UnixNanos, ffi::cvec::CVec};
25use nautilus_model::data::{Data, HasTsInit};
26use nautilus_serialization::arrow::{
27    DataStreamingError, DecodeDataFromRecordBatch, EncodeToRecordBatch, WriteStream,
28};
29use object_store::ObjectStore;
30use url::Url;
31
32use super::{
33    compare::Compare,
34    kmerge_batch::{EagerStream, ElementBatchIter, KMerge},
35};
36
37#[derive(Debug, Default)]
38pub struct TsInitComparator;
39
40impl<I> Compare<ElementBatchIter<I, Data>> for TsInitComparator
41where
42    I: Iterator<Item = IntoIter<Data>>,
43{
44    fn compare(
45        &self,
46        l: &ElementBatchIter<I, Data>,
47        r: &ElementBatchIter<I, Data>,
48    ) -> std::cmp::Ordering {
49        // Max heap ordering must be reversed
50        l.item.ts_init().cmp(&r.item.ts_init()).reverse()
51    }
52}
53
54pub type QueryResult = KMerge<EagerStream<std::vec::IntoIter<Data>>, Data, TsInitComparator>;
55
56/// Provides a DataFusion session and registers DataFusion queries.
57///
58/// The session is used to register data sources and make queries on them. A
59/// query returns a Chunk of Arrow records. It is decoded and converted into
60/// a Vec of data by types that implement [`DecodeDataFromRecordBatch`].
61#[cfg_attr(
62    feature = "python",
63    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.persistence", unsendable)
64)]
65#[cfg_attr(
66    feature = "python",
67    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")
68)]
69pub struct DataBackendSession {
70    pub chunk_size: usize,
71    pub runtime: Arc<tokio::runtime::Runtime>,
72    session_ctx: SessionContext,
73    batch_streams: Vec<EagerStream<IntoIter<Data>>>,
74    registered_tables: AHashSet<String>,
75}
76
77impl DataBackendSession {
78    /// Creates a new [`DataBackendSession`] instance.
79    ///
80    /// # Panics
81    ///
82    /// Panics if Tokio cannot create the worker runtime.
83    #[must_use]
84    pub fn new(chunk_size: usize) -> Self {
85        let runtime = tokio::runtime::Builder::new_multi_thread()
86            .enable_all()
87            .build()
88            .unwrap();
89        let session_cfg = SessionConfig::new()
90            .set_str("datafusion.optimizer.repartition_file_scans", "false")
91            .set_str("datafusion.optimizer.prefer_existing_sort", "true");
92        let session_ctx = SessionContext::new_with_config(session_cfg);
93        Self {
94            session_ctx,
95            batch_streams: Vec::default(),
96            chunk_size,
97            runtime: Arc::new(runtime),
98            registered_tables: AHashSet::new(),
99        }
100    }
101
102    /// Register an object store with the session context
103    pub fn register_object_store(&mut self, url: &Url, object_store: Arc<dyn ObjectStore>) {
104        self.session_ctx.register_object_store(url, object_store);
105    }
106
107    /// Register an object store with the session context from a URI with optional storage options.
108    ///
109    /// # Errors
110    ///
111    /// Returns an error if the object store URI cannot be normalized or the backend
112    /// cannot be created.
113    pub fn register_object_store_from_uri(
114        &mut self,
115        uri: &str,
116        storage_options: Option<AHashMap<String, String>>,
117    ) -> anyhow::Result<()> {
118        let location =
119            crate::parquet::create_object_store_location_from_path(uri, storage_options)?;
120
121        if let Some(root_url) = location.store_root_url().cloned() {
122            self.register_object_store(&root_url, location.object_store);
123        }
124
125        Ok(())
126    }
127
128    /// Writes encoded data to a streaming sink.
129    ///
130    /// # Errors
131    ///
132    /// Returns an error if Arrow encoding or stream writing fails.
133    pub fn write_data<T: EncodeToRecordBatch>(
134        data: &[T],
135        metadata: &AHashMap<String, String>,
136        stream: &mut dyn WriteStream,
137    ) -> Result<(), DataStreamingError> {
138        // Convert AHashMap to HashMap for Arrow compatibility
139        let metadata: std::collections::HashMap<String, String> = metadata
140            .iter()
141            .map(|(k, v)| (k.clone(), v.clone()))
142            .collect();
143        let record_batch = T::encode_batch(&metadata, data)?;
144        stream.write(&record_batch)?;
145        Ok(())
146    }
147
148    /// Registers a Parquet file and adds a batch stream for decoding.
149    ///
150    /// The caller must specify `T` to indicate the kind of data expected. `table_name` is
151    /// the logical name for queries; `file_path` is the Parquet path; `sql_query` defaults
152    /// to `SELECT * FROM {table_name} ORDER BY ts_init` if `None`.
153    ///
154    /// When `custom_type_name` is `Some`, it is merged into each batch's schema metadata
155    /// before decoding (as `type_name`). Use this for custom data when Parquet/DataFusion
156    /// does not preserve schema metadata so the decoder can look up the type in the registry.
157    ///
158    /// The file data must be ordered by the `ts_init` in ascending order for this
159    /// to work correctly.
160    ///
161    /// # Errors
162    ///
163    /// Returns an error if parquet registration, SQL planning, stream execution, or
164    /// data decoding setup fails.
165    pub fn add_file<T>(
166        &mut self,
167        table_name: &str,
168        file_path: &str,
169        sql_query: Option<&str>,
170        custom_type_name: Option<&str>,
171    ) -> Result<()>
172    where
173        T: DecodeDataFromRecordBatch,
174    {
175        // Check if table is already registered to avoid duplicates
176        let is_new_table = !self.registered_tables.contains(table_name);
177
178        if is_new_table {
179            // Register the table only if it doesn't exist
180            let parquet_options = ParquetReadOptions::<'_> {
181                skip_metadata: Some(false),
182                file_sort_order: vec![vec![Sort {
183                    expr: col("ts_init"),
184                    asc: true,
185                    nulls_first: false,
186                }]],
187                ..Default::default()
188            };
189            self.runtime.block_on(self.session_ctx.register_parquet(
190                table_name,
191                file_path,
192                parquet_options,
193            ))?;
194
195            self.registered_tables.insert(table_name.to_string());
196
197            // Only add batch stream for newly registered tables to avoid duplicates
198            let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
199            let sql_query = sql_query.unwrap_or(&default_query);
200            let query = self.runtime.block_on(self.session_ctx.sql(sql_query))?;
201            let batch_stream = self.runtime.block_on(query.execute_stream())?;
202            self.add_batch_stream::<T>(batch_stream, custom_type_name.map(String::from));
203        }
204
205        Ok(())
206    }
207
208    /// Registers a Parquet file and executes a query, returning the raw record batches.
209    ///
210    /// # Errors
211    ///
212    /// Returns an error if parquet registration, SQL planning, stream execution, or
213    /// batch collection fails.
214    pub fn collect_query_batches(
215        &mut self,
216        table_name: &str,
217        file_path: &str,
218        sql_query: Option<&str>,
219    ) -> Result<Vec<RecordBatch>> {
220        if !self.registered_tables.contains(table_name) {
221            let parquet_options = ParquetReadOptions::<'_> {
222                skip_metadata: Some(false),
223                file_sort_order: vec![vec![Sort {
224                    expr: col("ts_init"),
225                    asc: true,
226                    nulls_first: false,
227                }]],
228                ..Default::default()
229            };
230            self.runtime.block_on(self.session_ctx.register_parquet(
231                table_name,
232                file_path,
233                parquet_options,
234            ))?;
235
236            self.registered_tables.insert(table_name.to_string());
237        }
238
239        let default_query = format!("SELECT * FROM {table_name} ORDER BY ts_init");
240        let sql_query = sql_query.unwrap_or(&default_query);
241        let query = self.runtime.block_on(self.session_ctx.sql(sql_query))?;
242        let mut batch_stream = self.runtime.block_on(query.execute_stream())?;
243
244        self.runtime.block_on(async {
245            let mut batches = Vec::new();
246            while let Some(batch) = batch_stream.next().await {
247                batches.push(batch?);
248            }
249            Ok::<_, datafusion::error::DataFusionError>(batches)
250        })
251    }
252
253    fn add_batch_stream<T>(
254        &mut self,
255        stream: SendableRecordBatchStream,
256        custom_type_name: Option<String>,
257    ) where
258        T: DecodeDataFromRecordBatch,
259    {
260        let transform = stream.map(move |result| match result {
261            Ok(batch) => {
262                let mut metadata: std::collections::HashMap<String, String> =
263                    batch.schema().metadata().clone();
264
265                if let Some(ref tn) = custom_type_name {
266                    metadata.insert("type_name".to_string(), tn.clone());
267                }
268                T::decode_data_batch(&metadata, batch).unwrap().into_iter()
269            }
270            Err(e) => panic!("Error getting next batch from RecordBatchStream: {e}"),
271        });
272
273        self.batch_streams
274            .push(EagerStream::from_stream_with_runtime(
275                transform,
276                self.runtime.clone(),
277            ));
278    }
279
280    // Consumes the registered queries and returns a [`QueryResult].
281    // Passes the output of the query though the a KMerge which sorts the
282    // queries in ascending order of `ts_init`.
283    // QueryResult is an iterator that return Vec<Data>.
284    pub fn get_query_result(&mut self) -> QueryResult {
285        let mut kmerge: KMerge<_, _, _> = KMerge::new(TsInitComparator);
286
287        self.batch_streams
288            .drain(..)
289            .for_each(|eager_stream| kmerge.push_iter(eager_stream));
290
291        kmerge
292    }
293
294    /// Clears all registered tables and batch streams.
295    ///
296    /// This is useful when the underlying files have changed and we need to
297    /// re-register tables with updated data.
298    pub fn clear_registered_tables(&mut self) {
299        self.registered_tables.clear();
300        self.batch_streams.clear();
301
302        // Create a new session context to completely reset the DataFusion state
303        let session_cfg = SessionConfig::new()
304            .set_str("datafusion.optimizer.repartition_file_scans", "false")
305            .set_str("datafusion.optimizer.prefer_existing_sort", "true");
306        self.session_ctx = SessionContext::new_with_config(session_cfg);
307    }
308}
309
310#[must_use]
311pub fn build_query(
312    table: &str,
313    start: Option<UnixNanos>,
314    end: Option<UnixNanos>,
315    where_clause: Option<&str>,
316) -> String {
317    let mut conditions = Vec::new();
318
319    // Add where clause if provided
320    if let Some(clause) = where_clause {
321        conditions.push(clause.to_string());
322    }
323
324    // Add start condition if provided
325    if let Some(start_ts) = start {
326        conditions.push(format!("ts_init >= {start_ts}"));
327    }
328
329    // Add end condition if provided
330    if let Some(end_ts) = end {
331        conditions.push(format!("ts_init <= {end_ts}"));
332    }
333
334    // Build base query
335    let mut query = format!("SELECT * FROM {table}");
336
337    // Add WHERE clause if there are conditions
338    if !conditions.is_empty() {
339        query.push_str(" WHERE ");
340        query.push_str(&conditions.join(" AND "));
341    }
342
343    // Add ORDER BY clause
344    query.push_str(" ORDER BY ts_init");
345
346    query
347}
348
349#[cfg_attr(
350    feature = "python",
351    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.persistence", unsendable)
352)]
353#[cfg_attr(
354    feature = "python",
355    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")
356)]
357pub struct DataQueryResult {
358    pub chunk: Option<CVec>,
359    pub result: QueryResult,
360    pub acc: Vec<Data>,
361    pub size: usize,
362}
363
364impl DataQueryResult {
365    /// Creates a new [`DataQueryResult`] instance.
366    #[must_use]
367    pub const fn new(result: QueryResult, size: usize) -> Self {
368        Self {
369            chunk: None,
370            result,
371            acc: Vec::new(),
372            size,
373        }
374    }
375
376    /// Set new `CVec` backed chunk from data
377    ///
378    /// It also drops previously allocated chunk
379    pub fn set_chunk(&mut self, data: Vec<Data>) -> CVec {
380        self.drop_chunk();
381
382        let chunk: CVec = data.into();
383        self.chunk = Some(chunk);
384        chunk
385    }
386
387    /// Chunks generated by iteration must be dropped after use, otherwise
388    /// it will leak memory. Current chunk is held by the reader,
389    /// drop if exists and reset the field.
390    ///
391    /// # Panics
392    ///
393    /// Panics if the stored chunk has invalid length, capacity, or pointer
394    /// invariants.
395    pub fn drop_chunk(&mut self) {
396        if let Some(CVec { ptr, len, cap }) = self.chunk.take() {
397            assert!(
398                len <= cap,
399                "drop_chunk: len ({len}) > cap ({cap}) - memory corruption or wrong chunk type"
400            );
401            assert!(
402                len == 0 || !ptr.is_null(),
403                "drop_chunk: null ptr with non-zero len ({len}) - memory corruption"
404            );
405
406            // SAFETY: `ptr`, `len`, and `cap` originate from a valid `CVec` and the
407            // assertions above verify the invariants required by `Vec::from_raw_parts`.
408            let data: Vec<Data> = unsafe { Vec::from_raw_parts(ptr.cast::<Data>(), len, cap) };
409            drop(data);
410        }
411    }
412}
413
414impl Iterator for DataQueryResult {
415    type Item = Vec<Data>;
416
417    fn next(&mut self) -> Option<Self::Item> {
418        for _ in 0..self.size {
419            match self.result.next() {
420                Some(item) => self.acc.push(item),
421                None => break,
422            }
423        }
424
425        // TODO: consider using drain here if perf is unchanged
426        // Some(self.acc.drain(0..).collect())
427        let mut acc: Vec<Data> = Vec::new();
428        std::mem::swap(&mut acc, &mut self.acc);
429        Some(acc)
430    }
431}
432
433impl Drop for DataQueryResult {
434    fn drop(&mut self) {
435        self.drop_chunk();
436        self.result.clear();
437    }
438}
nautilus_persistence/backend/session.rs

nautilus_persistence/backend/
session.rs
