nautilus_serialization/arrow/
custom.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 code 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//! Custom data: registration and dynamic decoding.
17//!
18//! - **Registration:** Call [`ensure_custom_data_registered::<T>()`] once (e.g. before using the
19//!   catalog) for each custom data type `T` produced by the `#[custom_data]` macro. When Python
20//!   support is enabled, also call `nautilus_model::data::register_rust_extractor::<T>()`.
21//! - **Decoder:** [`CustomDataDecoder`] provides [`ArrowSchemaProvider`] and
22//!   [`DecodeDataFromRecordBatch`] for Parquet-backed custom data decoded at runtime by type name.
23//!   Types must be registered via [`ensure_custom_data_registered::<T>()`] before use.
24
25use std::{collections::HashMap, sync::Arc};
26
27use arrow::{
28    array::ArrayRef,
29    datatypes::{DataType as ArrowDataType, Field, Schema},
30    record_batch::RecordBatch,
31};
32use nautilus_model::data::{
33    ArrowDecoder, ArrowEncoder, CustomData, CustomDataTrait, Data, DataType,
34    decode_custom_from_arrow, ensure_arrow_registered, ensure_custom_data_json_registered,
35    get_arrow_schema,
36};
37
38use super::{
39    ArrowSchemaProvider, DecodeDataFromRecordBatch, EncodeToRecordBatch, EncodingError,
40    extract_column_string,
41};
42
43const KEY_TYPE_NAME: &str = "type_name";
44const COLUMN_DATA_TYPE: &str = "data_type";
45const FIELD_CUSTOM_DATA: &str = "custom_data";
46
47/// Trait for custom data types that support Arrow schema and record batch encoding.
48/// Used as a type bound by the `#[custom_data]` macro; catalog encoding goes through
49/// the registry, not this trait directly.
50///
51/// Implemented by the `#[custom_data]` macro for Rust custom data types. Python custom
52/// types use the registry encoder registered by `register_custom_data_class` instead.
53pub trait CustomDataSerialize: CustomDataTrait {
54    /// Returns the Arrow schema for this custom data type.
55    ///
56    /// # Errors
57    /// Returns an error if schema construction fails.
58    fn schema(&self) -> anyhow::Result<Schema>;
59
60    /// Encodes a batch of custom data items to an Arrow RecordBatch.
61    ///
62    /// # Errors
63    /// Returns an error if encoding fails (e.g. type mismatch or Arrow error).
64    fn encode_record_batch(
65        &self,
66        items: &[Arc<dyn CustomDataTrait>],
67    ) -> anyhow::Result<RecordBatch>;
68}
69
70/// Registers a custom data type in the JSON and Arrow registries. Call once per type
71/// (e.g. at catalog decode or before querying custom data).
72///
73/// Each distinct type `T` is registered at most once (per process). Safe to call
74/// multiple times for the same `T`.
75///
76/// When Python support is enabled, also call
77/// `nautilus_model::data::register_rust_extractor::<T>()` for types exposed to Python.
78pub fn ensure_custom_data_registered<T>()
79where
80    T: CustomDataTrait
81        + ArrowSchemaProvider
82        + EncodeToRecordBatch
83        + DecodeDataFromRecordBatch
84        + Clone
85        + Send
86        + Sync
87        + 'static,
88{
89    let type_name = T::type_name_static();
90
91    // Skip if already registered
92    if get_arrow_schema(type_name).is_some() {
93        return;
94    }
95
96    let _ = ensure_custom_data_json_registered::<T>();
97
98    let schema = Arc::new(T::get_schema(None));
99
100    let encoder: ArrowEncoder = Box::new(|items: &[Arc<dyn CustomDataTrait>]| {
101        let typed: Result<Vec<T>, _> = items
102            .iter()
103            .map(|b| {
104                b.as_any()
105                    .downcast_ref::<T>()
106                    .cloned()
107                    .ok_or_else(|| anyhow::anyhow!("Expected {}", T::type_name_static()))
108            })
109            .collect();
110        let typed = typed?;
111        let metadata = typed
112            .first()
113            .map(EncodeToRecordBatch::metadata)
114            .unwrap_or_default();
115        EncodeToRecordBatch::encode_batch(&metadata, &typed).map_err(|e| anyhow::anyhow!("{e}"))
116    });
117
118    let decoder: ArrowDecoder = Box::new(|metadata, batch| {
119        T::decode_data_batch(metadata, batch).map_err(|e| anyhow::anyhow!("{e}"))
120    });
121
122    let _ = ensure_arrow_registered(type_name, schema, encoder, decoder);
123}
124
125/// Decoder for custom data types that are identified at runtime by metadata (e.g. `type_name`).
126///
127/// Only Rust-registered custom types (e.g. `RustTestCustomData`, `MacroYieldCurveData`) can be
128/// decoded. Unknown types return an error.
129///
130/// **Important:** The caller must ensure that any Rust custom data types are registered
131/// via [`ensure_custom_data_registered::<T>()`] before use.
132#[derive(Debug)]
133pub struct CustomDataDecoder;
134
135impl ArrowSchemaProvider for CustomDataDecoder {
136    fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
137        if let Some(metadata) = metadata
138            && let Some(type_name) = metadata.get(KEY_TYPE_NAME)
139            && let Some(schema) = get_arrow_schema(type_name)
140        {
141            return (*schema).clone();
142        }
143
144        // Unknown type - return minimal schema (caller should not use this for decode)
145        Schema::new(vec![Field::new("dummy", ArrowDataType::Int64, true)])
146    }
147}
148
149impl DecodeDataFromRecordBatch for CustomDataDecoder {
150    fn decode_data_batch(
151        metadata: &HashMap<String, String>,
152        record_batch: RecordBatch,
153    ) -> Result<Vec<Data>, EncodingError> {
154        let type_name = metadata
155            .get(KEY_TYPE_NAME)
156            .cloned()
157            .unwrap_or_else(|| "Unknown".to_string());
158
159        let (batch, restored_data_type) = strip_data_type_column(&record_batch)?;
160
161        if batch.num_rows() == 0 {
162            return Ok(Vec::new());
163        }
164
165        let data = match decode_custom_from_arrow(&type_name, metadata, batch) {
166            Ok(Some(data)) => data,
167            Ok(None) => {
168                return Err(EncodingError::ParseError(
169                    FIELD_CUSTOM_DATA,
170                    format!(
171                        "unknown custom data type '{type_name}'; only Rust-registered types are supported"
172                    ),
173                ));
174            }
175            Err(e) => {
176                return Err(EncodingError::ParseError(
177                    FIELD_CUSTOM_DATA,
178                    format!("decode_custom_from_arrow: {e}"),
179                ));
180            }
181        };
182
183        let Some(data_type) = restored_data_type else {
184            return Ok(data);
185        };
186
187        Ok(data
188            .into_iter()
189            .map(|item| match item {
190                Data::Custom(custom) => {
191                    Data::Custom(CustomData::new(Arc::clone(&custom.data), data_type.clone()))
192                }
193                other => other,
194            })
195            .collect())
196    }
197}
198
199// Splits the `data_type` column off a batch so the registered decoder sees the schema it
200// registered. Returns the batch unchanged with `None` when the column is absent or the batch
201// is empty.
202fn strip_data_type_column(
203    batch: &RecordBatch,
204) -> Result<(RecordBatch, Option<DataType>), EncodingError> {
205    let Some(column_index) = batch
206        .schema()
207        .fields()
208        .iter()
209        .position(|field| field.name() == COLUMN_DATA_TYPE)
210    else {
211        return Ok((batch.clone(), None));
212    };
213
214    if batch.num_rows() == 0 {
215        return Ok((batch.clone(), None));
216    }
217
218    let columns = batch.columns();
219    let data_type = if columns[column_index].is_null(0) {
220        None
221    } else {
222        let values =
223            extract_column_string(columns, COLUMN_DATA_TYPE, column_index).map_err(|e| {
224                EncodingError::ParseError(FIELD_CUSTOM_DATA, format!("data_type column: {e}"))
225            })?;
226
227        Some(
228            DataType::from_persistence_json(values.value(0))
229                .map_err(|e| EncodingError::ParseError(FIELD_CUSTOM_DATA, e.to_string()))?,
230        )
231    };
232
233    let schema = batch.schema();
234    let fields: Vec<_> = schema
235        .fields()
236        .iter()
237        .enumerate()
238        .filter(|(index, _)| *index != column_index)
239        .map(|(_, field)| field.clone())
240        .collect();
241    let columns: Vec<ArrayRef> = columns
242        .iter()
243        .enumerate()
244        .filter(|(index, _)| *index != column_index)
245        .map(|(_, column)| Arc::clone(column))
246        .collect();
247    let stripped = Schema::new_with_metadata(fields, schema.metadata().clone());
248
249    RecordBatch::try_new(Arc::new(stripped), columns)
250        .map(|batch| (batch, data_type))
251        .map_err(|e| EncodingError::ParseError(FIELD_CUSTOM_DATA, e.to_string()))
252}
253
254#[cfg(test)]
255mod tests {
256    use std::collections::HashMap;
257
258    use arrow::{
259        array::{ArrayRef, Int64Array},
260        datatypes::{DataType as ArrowDataType, Field, Schema},
261    };
262    use rstest::rstest;
263
264    use super::*;
265    use crate::arrow::EncodingError;
266
267    #[rstest]
268    fn test_decode_data_batch_rejects_unregistered_type() {
269        let type_name = "UnregisteredCustomDataForArrowTest";
270        let metadata = HashMap::from([("type_name".to_string(), type_name.to_string())]);
271        let schema = Schema::new(vec![Field::new("value", ArrowDataType::Int64, false)]);
272        let columns: Vec<ArrayRef> = vec![Arc::new(Int64Array::from(vec![7]))];
273        let batch = RecordBatch::try_new(Arc::new(schema), columns).unwrap();
274
275        let error = CustomDataDecoder::decode_data_batch(&metadata, batch)
276            .expect_err("unregistered type must be rejected");
277
278        let EncodingError::ParseError(field, message) = error else {
279            panic!("unexpected error variant: {error:?}");
280        };
281        assert_eq!(field, "custom_data");
282        assert_eq!(
283            message,
284            "unknown custom data type 'UnregisteredCustomDataForArrowTest'; only Rust-registered types are supported"
285        );
286    }
287}
nautilus_serialization/arrow/custom.rs

nautilus_serialization/arrow/
custom.rs
