nautilus_serialization/arrow/display/
depth.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
16//! Nested display-mode Arrow encoding for [`OrderBookDepth`].
17
18use std::sync::Arc;
19
20use arrow::{
21    array::{StringBuilder, TimestampNanosecondBuilder, UInt8Builder, UInt64Builder},
22    datatypes::Schema,
23    error::ArrowError,
24    record_batch::RecordBatch,
25};
26use nautilus_model::data::OrderBookDepth;
27
28use super::{price_to_f64, quantity_to_f64, unix_nanos_to_i64};
29use crate::arrow::{
30    depth_display::{DepthSideBuilder, schema},
31    timestamp_data_type,
32};
33
34/// Returns the nested depth display schema.
35#[must_use]
36pub fn depths_schema() -> Schema {
37    schema()
38}
39
40/// Encodes every level into nested `bids` and `asks` lists.
41///
42/// Each level contains a display price and size, count, and exact integer order ID.
43/// Empty sides remain empty lists. Mixed instruments and depths share one schema.
44/// Prices and sizes use `Float64`, with nulls for undefined values; use raw catalog
45/// output when exact decimal values are required.
46///
47/// # Errors
48///
49/// Returns an error if Arrow cannot construct the batch or list offsets overflow.
50pub fn encode_depths(data: &[OrderBookDepth]) -> Result<RecordBatch, ArrowError> {
51    let mut instruments = StringBuilder::new();
52    let mut bids = DepthSideBuilder::new();
53    let mut asks = DepthSideBuilder::new();
54    let mut flags = UInt8Builder::with_capacity(data.len());
55    let mut sequence = UInt64Builder::with_capacity(data.len());
56    let mut ts_event =
57        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
58    let mut ts_init =
59        TimestampNanosecondBuilder::with_capacity(data.len()).with_data_type(timestamp_data_type());
60
61    for depth in data {
62        instruments.append_value(depth.instrument_id.to_string());
63        for (side, orders, counts) in [
64            (&mut bids, &depth.bids, &depth.bid_counts),
65            (&mut asks, &depth.asks, &depth.ask_counts),
66        ] {
67            if orders.len() != counts.len() {
68                return Err(ArrowError::InvalidArgumentError(
69                    "Depth orders and counts must have equal lengths".to_string(),
70                ));
71            }
72
73            for (order, count) in orders.iter().zip(counts) {
74                side.prices.append_option(
75                    (!order.price.is_undefined()).then(|| price_to_f64(&order.price)),
76                );
77                side.sizes.append_option(
78                    (!order.size.is_undefined()).then(|| quantity_to_f64(&order.size)),
79                );
80                side.counts.append_value(*count);
81                side.order_ids.append_value(order.order_id);
82            }
83            side.finish_row()?;
84        }
85        flags.append_value(depth.flags);
86        sequence.append_value(depth.sequence);
87        ts_event.append_value(unix_nanos_to_i64(depth.ts_event.as_u64()));
88        ts_init.append_value(unix_nanos_to_i64(depth.ts_init.as_u64()));
89    }
90
91    RecordBatch::try_new(
92        Arc::new(depths_schema()),
93        vec![
94            Arc::new(instruments.finish()),
95            Arc::new(bids.finish()?),
96            Arc::new(asks.finish()?),
97            Arc::new(flags.finish()),
98            Arc::new(sequence.finish()),
99            Arc::new(ts_event.finish()),
100            Arc::new(ts_init.finish()),
101        ],
102    )
103}
104
105#[cfg(test)]
106mod tests {
107    use arrow::{
108        array::{
109            Array, Float64Array, ListArray, StringArray, StructArray, TimestampNanosecondArray,
110            UInt8Array, UInt32Array, UInt64Array,
111        },
112        datatypes::{DataType, Field, Fields, TimeUnit},
113    };
114    use nautilus_model::{
115        data::BookOrder,
116        enums::OrderSide,
117        identifiers::InstrumentId,
118        types::{PRICE_UNDEF, Price, QUANTITY_UNDEF, Quantity},
119    };
120    use rstest::rstest;
121
122    use super::*;
123
124    #[rstest]
125    fn test_depth_display_schema_and_empty_batch() {
126        let batch = encode_depths(&[]).unwrap();
127        let levels: Fields = vec![
128            Field::new("price", DataType::Float64, true),
129            Field::new("size", DataType::Float64, true),
130            Field::new("count", DataType::UInt32, false),
131            Field::new("order_id", DataType::UInt64, false),
132        ]
133        .into();
134        let side = DataType::List(Arc::new(Field::new(
135            "item",
136            DataType::Struct(levels),
137            false,
138        )));
139        let expected = Schema::new(vec![
140            Field::new("instrument_id", DataType::Utf8, false),
141            Field::new("bids", side.clone(), false),
142            Field::new("asks", side, false),
143            Field::new("flags", DataType::UInt8, false),
144            Field::new("sequence", DataType::UInt64, false),
145            Field::new(
146                "ts_event",
147                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
148                false,
149            ),
150            Field::new(
151                "ts_init",
152                DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
153                false,
154            ),
155        ]);
156
157        assert_eq!(batch.num_rows(), 0);
158        assert_eq!(batch.schema().as_ref(), &expected);
159        assert_eq!(depths_schema(), expected);
160    }
161
162    #[rstest]
163    fn test_depth_display_preserves_mixed_depths_and_every_field() {
164        let data = [
165            depth("AAPL.XNAS", 0, 3, 1),
166            depth("MSFT.XNAS", 5, 0, 2),
167            depth("NVDA.XNAS", 25, 27, 3),
168        ];
169        let batch = encode_depths(&data).unwrap();
170        let ids = batch
171            .column_by_name("instrument_id")
172            .unwrap()
173            .as_any()
174            .downcast_ref::<StringArray>()
175            .unwrap();
176        let flags = batch
177            .column_by_name("flags")
178            .unwrap()
179            .as_any()
180            .downcast_ref::<UInt8Array>()
181            .unwrap();
182        let sequence = batch
183            .column_by_name("sequence")
184            .unwrap()
185            .as_any()
186            .downcast_ref::<UInt64Array>()
187            .unwrap();
188        let ts_event = batch
189            .column_by_name("ts_event")
190            .unwrap()
191            .as_any()
192            .downcast_ref::<TimestampNanosecondArray>()
193            .unwrap();
194        let ts_init = batch
195            .column_by_name("ts_init")
196            .unwrap()
197            .as_any()
198            .downcast_ref::<TimestampNanosecondArray>()
199            .unwrap();
200
201        assert_eq!(batch.num_rows(), 3);
202        assert_eq!(batch.schema().as_ref(), &depths_schema());
203
204        for (row, (instrument, bid_len, ask_len, seed)) in [
205            ("AAPL.XNAS", 0, 3, 1_u32),
206            ("MSFT.XNAS", 5, 0, 2),
207            ("NVDA.XNAS", 25, 27, 3),
208        ]
209        .into_iter()
210        .enumerate()
211        {
212            assert_eq!(ids.value(row), instrument);
213            assert_eq!(flags.value(row), u8::try_from(seed).unwrap());
214            assert_eq!(sequence.value(row), u64::from(seed) + 30);
215            assert_eq!(ts_event.value(row), i64::from(seed) + 40);
216            assert_eq!(ts_init.value(row), i64::from(seed) + 50);
217
218            for (name, len, offset) in [("bids", bid_len, 0_u32), ("asks", ask_len, 100)] {
219                let list = batch
220                    .column_by_name(name)
221                    .unwrap()
222                    .as_any()
223                    .downcast_ref::<ListArray>()
224                    .unwrap();
225                let values = list.value(row);
226                let levels = values.as_any().downcast_ref::<StructArray>().unwrap();
227                let prices = levels
228                    .column_by_name("price")
229                    .unwrap()
230                    .as_any()
231                    .downcast_ref::<Float64Array>()
232                    .unwrap();
233                let sizes = levels
234                    .column_by_name("size")
235                    .unwrap()
236                    .as_any()
237                    .downcast_ref::<Float64Array>()
238                    .unwrap();
239                let counts = levels
240                    .column_by_name("count")
241                    .unwrap()
242                    .as_any()
243                    .downcast_ref::<UInt32Array>()
244                    .unwrap();
245                let order_ids = levels
246                    .column_by_name("order_id")
247                    .unwrap()
248                    .as_any()
249                    .downcast_ref::<UInt64Array>()
250                    .unwrap();
251                assert!(!list.is_null(row));
252                assert_eq!(levels.len(), len);
253                assert_eq!(levels.null_count(), 0);
254
255                for i in 0..len {
256                    let n = seed + offset + u32::try_from(i).unwrap();
257                    assert_eq!(prices.value(i), f64::from(n) + 0.25);
258                    assert_eq!(sizes.value(i), f64::from(n) + 0.5);
259                    assert_eq!(counts.value(i), n + 10);
260                    assert_eq!(order_ids.value(i), u64::MAX - u64::from(n));
261                }
262            }
263        }
264    }
265
266    #[rstest]
267    fn test_depth_display_undefined_fields_are_null() {
268        let mut data = depth("AAPL.XNAS", 1, 1, 1);
269        data.bids[0].price = Price::from_raw(PRICE_UNDEF, 0);
270        data.asks[0].size = Quantity::from_raw(QUANTITY_UNDEF, 0);
271        let batch = encode_depths(&[data]).unwrap();
272        for (name, field) in [("bids", "price"), ("asks", "size")] {
273            let list = batch
274                .column_by_name(name)
275                .unwrap()
276                .as_any()
277                .downcast_ref::<ListArray>()
278                .unwrap();
279            let values = list.value(0);
280            let levels = values.as_any().downcast_ref::<StructArray>().unwrap();
281            assert_eq!(levels.len(), 1);
282            assert!(levels.column_by_name(field).unwrap().is_null(0));
283        }
284    }
285
286    fn depth(instrument: &str, bids: usize, asks: usize, seed: u32) -> OrderBookDepth {
287        let side = |len, offset, side| {
288            let orders: Vec<_> = (0..len)
289                .map(|i| {
290                    let n = seed + offset + u32::try_from(i).unwrap();
291                    BookOrder::new(
292                        side,
293                        format!("{n}.25").parse().unwrap(),
294                        format!("{n}.5").parse().unwrap(),
295                        u64::MAX - u64::from(n),
296                    )
297                })
298                .collect();
299            let counts: Vec<_> = (0..len)
300                .map(|i| seed + offset + u32::try_from(i).unwrap() + 10)
301                .collect();
302            (orders, counts)
303        };
304        let (bids, bid_counts) = side(bids, 0, OrderSide::Buy);
305        let (asks, ask_counts) = side(asks, 100, OrderSide::Sell);
306        OrderBookDepth::new_checked(
307            InstrumentId::from(instrument),
308            bids,
309            asks,
310            bid_counts,
311            ask_counts,
312            u8::try_from(seed).unwrap(),
313            u64::from(seed) + 30,
314            (u64::from(seed) + 40).into(),
315            (u64::from(seed) + 50).into(),
316        )
317        .unwrap()
318    }
319}
nautilus_serialization/arrow/display/depth.rs

nautilus_serialization/arrow/display/
depth.rs
