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            chunk.push(item);
122        }
123
124        let boundary_ts = chunk.last().map(HasTsInit::ts_init);
125
126        loop {
127            match self.next_row()? {
128                Some(item) if Some(item.ts_init()) == boundary_ts => chunk.push(item),
129                Some(item) => {
130                    self.carry.push_front(item);
131                    break;
132                }
133                None => break,
134            }
135        }
136
137        Ok(Some(T::into_batch(chunk)))
138    }
139}
nautilus_persistence/catalog/session.rs

nautilus_persistence/catalog/
session.rs
