nautilus_persistence/backend/
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 persistence: shared helpers and orchestration.
17//!
18//! Centralizes the logic for appending the `data_type` column and metadata to Arrow batches
19//! (Parquet/Feather), and custom-data write preparation, path construction, and decode logic
20//! so the catalog delegates here instead of inlining custom-specific branching.
21
22use std::{collections::HashMap, sync::Arc};
23
24use datafusion::arrow::{
25    array::{Array, StringArray},
26    datatypes::{DataType as ArrowDataType, Field, Schema},
27    record_batch::RecordBatch,
28};
29use nautilus_core::UnixNanos;
30use nautilus_model::data::{
31    Bar, CustomData, CustomDataTrait, Data, IndexPriceUpdate, MarkPriceUpdate, OptionGreeks,
32    OrderBookDelta, OrderBookDepth10, QuoteTick, TradeTick, close::InstrumentClose,
33    encode_custom_to_arrow,
34};
35use nautilus_serialization::arrow::DecodeDataFromRecordBatch;
36#[cfg(feature = "python")]
37use nautilus_serialization::arrow::custom::CustomDataDecoder;
38
39/// Builds a schema that adds the `data_type` column and `type_name` metadata to a base schema.
40/// Used when creating a Feather buffer for custom data (single type per writer).
41#[must_use]
42pub fn schema_with_data_type_column(base_schema: &Schema, type_name: &str) -> Schema {
43    let mut fields: Vec<_> = base_schema.fields().iter().cloned().collect();
44    fields.push(Arc::new(Field::new("data_type", ArrowDataType::Utf8, true)));
45    let mut meta = base_schema.metadata().clone();
46    meta.insert("type_name".to_string(), type_name.to_string());
47    Schema::new_with_metadata(fields, meta)
48}
49
50/// Appends a `data_type` column (JSON string per row) and `type_name` + optional metadata to the
51/// batch schema. Used by both the Parquet catalog and Feather writer for catalog-compatible output.
52///
53/// # Errors
54///
55/// Returns an error if the new `RecordBatch` cannot be created.
56#[expect(
57    clippy::implicit_hasher,
58    reason = "Arrow schema metadata uses the standard HashMap type"
59)]
60pub fn augment_batch_with_data_type_column(
61    batch: &RecordBatch,
62    data_type_json: &str,
63    type_name: &str,
64    dt_meta: Option<&HashMap<String, String>>,
65) -> anyhow::Result<RecordBatch> {
66    let num_rows = batch.num_rows();
67    let data_type_array: Arc<dyn Array> = Arc::new(StringArray::from(
68        (0..num_rows)
69            .map(|_| Some(data_type_json))
70            .collect::<Vec<_>>(),
71    ));
72    let schema = batch.schema();
73    let mut fields: Vec<_> = schema.fields().iter().cloned().collect();
74    fields.push(Arc::new(Field::new(
75        "data_type",
76        ArrowDataType::Utf8,
77        false,
78    )));
79    let mut meta = schema.metadata().clone();
80    meta.insert("type_name".to_string(), type_name.to_string());
81
82    if let Some(m) = dt_meta {
83        meta.extend(m.clone());
84    }
85    let new_schema = Arc::new(Schema::new_with_metadata(fields, meta));
86    let mut columns = batch.columns().to_vec();
87    columns.push(data_type_array);
88    let new_batch = RecordBatch::try_new(new_schema, columns)
89        .map_err(|e| anyhow::anyhow!("Failed to merge custom data type metadata: {e}"))?;
90    Ok(new_batch)
91}
92
93/// Normalizes a custom data identifier for use in directory paths.
94/// Replaces `//` with `/`, and filters out empty segments and `..` to prevent path traversal.
95#[must_use]
96fn safe_directory_identifier(identifier: &str) -> String {
97    let normalized = identifier.replace("//", "/");
98    let segments: Vec<&str> = normalized
99        .split('/')
100        .filter(|s| !s.is_empty() && *s != "..")
101        .collect();
102    segments.join("/")
103}
104
105/// Returns path components for custom data: `["data", "custom", type_name, ...identifier segments]`.
106/// Used by the catalog to build full object-store paths via `make_object_store_path_owned`.
107#[must_use]
108pub fn custom_data_path_components(type_name: &str, identifier: Option<&str>) -> Vec<String> {
109    let mut components = vec![
110        "data".to_string(),
111        "custom".to_string(),
112        type_name.to_string(),
113    ];
114
115    if let Some(id) = identifier {
116        let safe = safe_directory_identifier(id);
117        if !safe.is_empty() {
118            for segment in safe.split('/') {
119                components.push(segment.to_string());
120            }
121        }
122    }
123    components
124}
125
126/// Prepares a batch of custom data for writing: encodes to Arrow, augments with `data_type` column,
127/// and returns type identity and timestamp range so the catalog can build path and perform I/O.
128///
129/// # Errors
130///
131/// Returns an error if encoding or augmentation fails, or if the type is not registered.
132pub fn prepare_custom_data_batch(
133    data: Vec<CustomData>,
134) -> anyhow::Result<(RecordBatch, String, Option<String>, UnixNanos, UnixNanos)> {
135    let Some(first_custom) = data.first() else {
136        anyhow::bail!("prepare_custom_data_batch called with empty data");
137    };
138
139    let type_name = first_custom.data.type_name();
140    let identifier = first_custom.data_type.identifier().map(String::from);
141    let dt_meta = first_custom.data_type.metadata_string_map();
142    let data_type_json = first_custom
143        .data_type
144        .to_persistence_json()
145        .map_err(|e| anyhow::anyhow!("Failed to serialize data_type for persistence: {e}"))?;
146
147    let start_ts = first_custom.data.ts_init();
148    let end_ts = data.last().map_or(start_ts, |custom| custom.data.ts_init());
149    let items: Vec<Arc<dyn CustomDataTrait>> =
150        data.into_iter().map(|c| Arc::clone(&c.data)).collect();
151
152    let batch = encode_custom_to_arrow(type_name, &items)
153        .map_err(|e| anyhow::anyhow!("Failed to encode custom data to Arrow: {e}"))?
154        .ok_or_else(|| {
155            anyhow::anyhow!(
156                "Custom data type \"{type_name}\" is not registered for Arrow encoding; \
157                 call register_custom_data_class or ensure_custom_data_registered before writing"
158            )
159        })?;
160    let batch =
161        augment_batch_with_data_type_column(&batch, &data_type_json, type_name, dt_meta.as_ref())?;
162
163    Ok((batch, type_name.to_string(), identifier, start_ts, end_ts))
164}
165
166/// Decodes a `RecordBatch` to Data objects based on metadata.
167///
168/// Supports both standard data types and custom data types when `allow_custom_fallback`
169/// is true (e.g. when decoding files under `custom/`). When false, unknown type names
170/// produce an error instead of attempting custom decode.
171///
172/// # Errors
173///
174/// Returns an error if decoding fails or the type is unknown (and custom fallback not allowed).
175#[expect(
176    clippy::implicit_hasher,
177    reason = "Arrow schema metadata uses the standard HashMap type"
178)]
179pub fn decode_batch_to_data(
180    metadata: &HashMap<String, String>,
181    batch: RecordBatch,
182    allow_custom_fallback: bool,
183) -> anyhow::Result<Vec<Data>> {
184    let type_name = metadata
185        .get("type_name")
186        .cloned()
187        .or_else(|| metadata.get("bar_type").map(|_| "bars".to_string()))
188        .ok_or_else(|| anyhow::anyhow!("Missing type_name in metadata"))?;
189
190    match type_name.as_str() {
191        "QuoteTick" | "quotes" => Ok(QuoteTick::decode_data_batch(metadata, batch)?),
192        "TradeTick" | "trades" => Ok(TradeTick::decode_data_batch(metadata, batch)?),
193        "Bar" | "bars" => Ok(Bar::decode_data_batch(metadata, batch)?),
194        "OrderBookDelta" | "order_book_deltas" => {
195            Ok(OrderBookDelta::decode_data_batch(metadata, batch)?)
196        }
197        "OrderBookDepth10" | "order_book_depths" => {
198            Ok(OrderBookDepth10::decode_data_batch(metadata, batch)?)
199        }
200        "MarkPriceUpdate" | "mark_price_updates" => {
201            Ok(MarkPriceUpdate::decode_data_batch(metadata, batch)?)
202        }
203        "IndexPriceUpdate" | "index_price_updates" => {
204            Ok(IndexPriceUpdate::decode_data_batch(metadata, batch)?)
205        }
206        "OptionGreeks" | "option_greeks" => Ok(OptionGreeks::decode_data_batch(metadata, batch)?),
207        "InstrumentClose" | "instrument_closes" => {
208            Ok(InstrumentClose::decode_data_batch(metadata, batch)?)
209        }
210        _ => {
211            if allow_custom_fallback {
212                #[cfg(feature = "python")]
213                {
214                    return Ok(CustomDataDecoder::decode_data_batch(metadata, batch)?);
215                }
216                #[cfg(not(feature = "python"))]
217                {
218                    anyhow::bail!("Unknown data type: {type_name}");
219                }
220            }
221            anyhow::bail!(
222                "Unknown data type: {type_name}; custom decode only allowed in custom data context"
223            )
224        }
225    }
226}
227
228/// Decodes multiple `RecordBatches` (e.g. from custom data files) into a single `Vec<Data>`.
229/// Optionally replaces `ts_init` column with `ts_event` before decoding each batch.
230///
231/// # Errors
232///
233/// Returns an error if any batch fails to decode.
234pub fn decode_custom_batches_to_data(
235    batches: Vec<RecordBatch>,
236    use_ts_event_for_ts_init: bool,
237) -> anyhow::Result<Vec<Data>> {
238    let mut file_data = Vec::new();
239    let schema = batches
240        .first()
241        .map(arrow::array::RecordBatch::schema)
242        .ok_or_else(|| {
243            anyhow::anyhow!("decode_custom_batches_to_data called with empty batches")
244        })?;
245
246    for mut batch in batches {
247        if use_ts_event_for_ts_init {
248            let column_names: Vec<String> =
249                schema.fields().iter().map(|f| f.name().clone()).collect();
250
251            if let (Some(ts_event_idx), Some(ts_init_idx)) = (
252                column_names.iter().position(|n| n == "ts_event"),
253                column_names.iter().position(|n| n == "ts_init"),
254            ) {
255                let mut new_columns = batch.columns().to_vec();
256                new_columns[ts_init_idx] = new_columns[ts_event_idx].clone();
257                batch = RecordBatch::try_new(schema.clone(), new_columns)
258                    .map_err(|e| anyhow::anyhow!("Failed to create new batch: {e}"))?;
259            }
260        }
261        let metadata = batch.schema().metadata().clone();
262        let decoded = decode_batch_to_data(&metadata, batch, true)?;
263        file_data.extend(decoded);
264    }
265    Ok(file_data)
266}
nautilus_persistence/backend/custom.rs

nautilus_persistence/backend/
custom.rs
