nautilus_persistence/catalog/
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. See the License for
12//  the specific language governing permissions and limitations under the License.
13// -------------------------------------------------------------------------------------------------
14
15//! Typed catalog query sessions.
16//!
17//! Every typed batch session yields rows in non-decreasing `ts_init` order. Pages must preserve
18//! that order across the whole source; debug builds assert the invariant for every yielded row.
19
20use std::collections::VecDeque;
21
22use nautilus_core::UnixNanos;
23use nautilus_model::data::{DataBatch, HasTsInit, IntoDataBatch};
24
25pub const DEFAULT_DATA_BATCH_CHUNK_SIZE: usize = 10_000;
26
27/// Streaming query session that yields ordered [`DataBatch`] chunks.
28pub trait DataBatchQuery: Send {
29    /// Returns the next typed batch, or `None` after the session is exhausted.
30    ///
31    /// # Errors
32    ///
33    /// Returns an error if the underlying page source or typed conversion fails.
34    fn next_batch(&mut self) -> anyhow::Result<Option<DataBatch>>;
35
36    /// Resets the session when the backend supports replaying its source.
37    ///
38    /// # Errors
39    ///
40    /// Returns an error if the backend cannot reset its source state.
41    fn reset(&mut self) -> anyhow::Result<bool> {
42        Ok(false)
43    }
44}
45
46pub type DataBatchQueryResult = Box<dyn DataBatchQuery>;
47
48/// Adapts pages of ordered typed values into timestamp-aligned [`DataBatch`] chunks.
49pub struct TypedDataBatchSession<T> {
50    pages: Box<dyn Iterator<Item = anyhow::Result<Vec<T>>> + Send>,
51    carry: VecDeque<T>,
52    chunk_size: usize,
53    previous_ts_init: Option<UnixNanos>,
54}
55
56impl<T> TypedDataBatchSession<T>
57where
58    T: IntoDataBatch + HasTsInit + Send,
59{
60    /// Creates a session over a page source.
61    #[must_use]
62    pub fn new(
63        pages: Box<dyn Iterator<Item = anyhow::Result<Vec<T>>> + Send>,
64        chunk_size: Option<usize>,
65    ) -> Self {
66        Self {
67            pages,
68            carry: VecDeque::new(),
69            chunk_size: chunk_size.unwrap_or(DEFAULT_DATA_BATCH_CHUNK_SIZE).max(1),
70            previous_ts_init: None,
71        }
72    }
73
74    /// Creates a session over one collected, ordered vector.
75    #[must_use]
76    pub fn from_vec(data: Vec<T>, chunk_size: Option<usize>) -> Self
77    where
78        T: 'static,
79    {
80        Self::new(Box::new(std::iter::once(Ok(data))), chunk_size)
81    }
82
83    fn next_row(&mut self) -> anyhow::Result<Option<T>> {
84        loop {
85            if let Some(row) = self.carry.pop_front() {
86                debug_assert!(
87                    self.previous_ts_init
88                        .is_none_or(|previous| previous <= row.ts_init()),
89                    "TypedDataBatchSession ordering invariant violated: ts_init {} after {:?}",
90                    row.ts_init(),
91                    self.previous_ts_init,
92                );
93                self.previous_ts_init = Some(row.ts_init());
94                return Ok(Some(row));
95            }
96
97            match self.pages.next() {
98                Some(page) => self.carry.extend(page?),
99                None => return Ok(None),
100            }
101        }
102    }
103}
104
105impl<T> DataBatchQuery for TypedDataBatchSession<T>
106where
107    T: IntoDataBatch + HasTsInit + Send,
108{
109    fn next_batch(&mut self) -> anyhow::Result<Option<DataBatch>> {
110        let Some(first) = self.next_row()? else {
111            return Ok(None);
112        };
113
114        let mut chunk = Vec::with_capacity(self.chunk_size.min(1024));
115        chunk.push(first);
116
117        while chunk.len() < self.chunk_size {
118            let Some(item) = self.next_row()? else {
119                return Ok(Some(T::into_batch(chunk)));
120            };
121
122            chunk.push(item);
123        }
124
125        let boundary_ts = chunk.last().map(HasTsInit::ts_init);
126
127        loop {
128            match self.next_row()? {
129                Some(item) if Some(item.ts_init()) == boundary_ts => chunk.push(item),
130                Some(item) => {
131                    self.carry.push_front(item);
132                    break;
133                }
134                None => break,
135            }
136        }
137
138        Ok(Some(T::into_batch(chunk)))
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use nautilus_model::{
145        data::QuoteTick,
146        identifiers::InstrumentId,
147        types::{Price, Quantity},
148    };
149    use rstest::rstest;
150
151    use super::*;
152
153    fn typed_quote(ts_init: u64) -> QuoteTick {
154        QuoteTick::new(
155            InstrumentId::from("AUD/USD.SIM"),
156            Price::from("1.0"),
157            Price::from("1.1"),
158            Quantity::from("1000"),
159            Quantity::from("1000"),
160            UnixNanos::from(ts_init),
161            UnixNanos::from(ts_init),
162        )
163    }
164
165    fn batch_ts(batch: &DataBatch) -> Vec<u64> {
166        match batch {
167            DataBatch::Quote(quotes) => quotes
168                .as_ref()
169                .iter()
170                .map(|quote| quote.ts_init.as_u64())
171                .collect(),
172            other => panic!("expected quote batch, found {other:?}"),
173        }
174    }
175
176    #[rstest]
177    fn typed_session_chunks_pages_with_carry_across_pulls() {
178        let pages: Vec<anyhow::Result<Vec<QuoteTick>>> = vec![
179            Ok(vec![typed_quote(1), typed_quote(2), typed_quote(3)]),
180            Ok(Vec::new()),
181            Ok(vec![typed_quote(4), typed_quote(5)]),
182        ];
183        let mut session = TypedDataBatchSession::new(Box::new(pages.into_iter()), Some(2));
184
185        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [1, 2]);
186        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [3, 4]);
187        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [5]);
188        assert!(session.next_batch().unwrap().is_none());
189    }
190
191    #[rstest]
192    fn typed_session_extends_chunk_across_equal_boundary_ts() {
193        let data = vec![
194            typed_quote(1),
195            typed_quote(2),
196            typed_quote(2),
197            typed_quote(2),
198            typed_quote(3),
199        ];
200        let mut session = TypedDataBatchSession::from_vec(data, Some(2));
201
202        assert_eq!(
203            batch_ts(&session.next_batch().unwrap().unwrap()),
204            [1, 2, 2, 2]
205        );
206        assert_eq!(batch_ts(&session.next_batch().unwrap().unwrap()), [3]);
207        assert!(session.next_batch().unwrap().is_none());
208    }
209
210    #[rstest]
211    fn typed_session_empty_source_yields_none() {
212        let mut session = TypedDataBatchSession::<QuoteTick>::from_vec(Vec::new(), None);
213
214        assert!(session.next_batch().unwrap().is_none());
215    }
216
217    #[rstest]
218    fn typed_session_propagates_page_error() {
219        let pages: Vec<anyhow::Result<Vec<QuoteTick>>> = vec![
220            Ok(vec![typed_quote(1)]),
221            Err(anyhow::anyhow!("page failed")),
222        ];
223        let mut session = TypedDataBatchSession::new(Box::new(pages.into_iter()), Some(4));
224
225        assert_eq!(session.next_batch().unwrap_err().to_string(), "page failed");
226    }
227
228    #[rstest]
229    fn typed_session_reports_reset_as_unsupported() {
230        let mut session = TypedDataBatchSession::from_vec(vec![typed_quote(1)], None);
231
232        assert!(!session.reset().unwrap());
233    }
234}
nautilus_persistence/catalog/session.rs

nautilus_persistence/catalog/
session.rs
