1use std::{collections::HashMap, str::FromStr, sync::Arc};
17
18use arrow::{
19 array::{Array, Float64Array, Float64Builder, UInt64Array},
20 datatypes::{DataType, Field, Schema},
21 error::ArrowError,
22 record_batch::RecordBatch,
23};
24use nautilus_model::{
25 data::{Data, greeks::OptionGreekValues, option_chain::OptionGreeks},
26 enums::GreeksConvention,
27 identifiers::InstrumentId,
28};
29
30use super::{
31 ArrowSchemaProvider, DecodeDataFromRecordBatch, DecodeFromRecordBatch, EncodeToRecordBatch,
32 EncodingError, KEY_IDENTIFIER, KEY_INSTRUMENT_ID, decode_required_timestamp,
33 enum_dictionary_array, enum_dictionary_data_type, extract_column, extract_column_string,
34 identifier_array_from_display, json::metadata_for_type, metadata_with_type_name,
35 validate_required_column,
36};
37
38const TYPE_NAME: &str = "OptionGreeks";
39
40impl ArrowSchemaProvider for OptionGreeks {
41 fn get_schema(metadata: Option<HashMap<String, String>>) -> Schema {
42 let fields = vec![
43 Field::new("convention", enum_dictionary_data_type(), false),
44 Field::new("delta", DataType::Float64, false),
45 Field::new("gamma", DataType::Float64, false),
46 Field::new("vega", DataType::Float64, false),
47 Field::new("theta", DataType::Float64, false),
48 Field::new("rho", DataType::Float64, false),
49 Field::new("mark_iv", DataType::Float64, true),
50 Field::new("bid_iv", DataType::Float64, true),
51 Field::new("ask_iv", DataType::Float64, true),
52 Field::new("underlying_price", DataType::Float64, true),
53 Field::new("open_interest", DataType::Float64, true),
54 Field::new("ts_event", crate::arrow::timestamp_data_type(), false),
55 Field::new("ts_init", crate::arrow::timestamp_data_type(), false),
56 Field::new(KEY_IDENTIFIER, DataType::Utf8, true),
57 ];
58
59 Schema::new_with_metadata(fields, metadata_with_type_name(TYPE_NAME, metadata))
60 }
61}
62
63fn append_optional_f64(builder: &mut Float64Builder, value: Option<f64>) {
64 match value {
65 Some(value) => builder.append_value(value),
66 None => builder.append_null(),
67 }
68}
69
70fn optional_f64(values: &Float64Array, row: usize) -> Option<f64> {
71 if values.is_null(row) {
72 None
73 } else {
74 Some(values.value(row))
75 }
76}
77
78impl EncodeToRecordBatch for OptionGreeks {
79 fn encode_batch<T>(
80 metadata: &HashMap<String, String>,
81 data: &[T],
82 ) -> Result<RecordBatch, ArrowError>
83 where
84 T: std::borrow::Borrow<Self>,
85 {
86 let mut delta_builder = Float64Builder::new();
87 let mut gamma_builder = Float64Builder::new();
88 let mut vega_builder = Float64Builder::new();
89 let mut theta_builder = Float64Builder::new();
90 let mut rho_builder = Float64Builder::new();
91 let mut mark_iv_builder = Float64Builder::new();
92 let mut bid_iv_builder = Float64Builder::new();
93 let mut ask_iv_builder = Float64Builder::new();
94 let mut underlying_price_builder = Float64Builder::new();
95 let mut open_interest_builder = Float64Builder::new();
96 let mut ts_event_builder = UInt64Array::builder(data.len());
97 let mut ts_init_builder = UInt64Array::builder(data.len());
98
99 for greeks in data.iter().map(std::borrow::Borrow::borrow) {
100 delta_builder.append_value(greeks.delta);
101 gamma_builder.append_value(greeks.gamma);
102 vega_builder.append_value(greeks.vega);
103 theta_builder.append_value(greeks.theta);
104 rho_builder.append_value(greeks.rho);
105 append_optional_f64(&mut mark_iv_builder, greeks.mark_iv);
106 append_optional_f64(&mut bid_iv_builder, greeks.bid_iv);
107 append_optional_f64(&mut ask_iv_builder, greeks.ask_iv);
108 append_optional_f64(&mut underlying_price_builder, greeks.underlying_price);
109 append_optional_f64(&mut open_interest_builder, greeks.open_interest);
110 ts_event_builder.append_value(greeks.ts_event.as_u64());
111 ts_init_builder.append_value(greeks.ts_init.as_u64());
112 }
113
114 let conventions = enum_dictionary_array(
115 data.iter()
116 .map(std::borrow::Borrow::borrow)
117 .map(|greeks| greeks.convention),
118 )?;
119
120 crate::arrow::record_batch_with_timestamps(
121 Arc::new(Self::get_schema(Some(metadata.clone()))),
122 vec![
123 Arc::new(conventions),
124 Arc::new(delta_builder.finish()),
125 Arc::new(gamma_builder.finish()),
126 Arc::new(vega_builder.finish()),
127 Arc::new(theta_builder.finish()),
128 Arc::new(rho_builder.finish()),
129 Arc::new(mark_iv_builder.finish()),
130 Arc::new(bid_iv_builder.finish()),
131 Arc::new(ask_iv_builder.finish()),
132 Arc::new(underlying_price_builder.finish()),
133 Arc::new(open_interest_builder.finish()),
134 Arc::new(ts_event_builder.finish()),
135 Arc::new(ts_init_builder.finish()),
136 Arc::new(identifier_array_from_display(
137 data.iter()
138 .map(std::borrow::Borrow::borrow)
139 .map(|greeks| greeks.instrument_id),
140 )),
141 ],
142 )
143 }
144
145 fn metadata(&self) -> HashMap<String, String> {
146 let mut metadata = metadata_for_type(TYPE_NAME);
147 metadata.insert(
148 KEY_INSTRUMENT_ID.to_string(),
149 self.instrument_id.to_string(),
150 );
151 metadata
152 }
153}
154
155impl DecodeFromRecordBatch for OptionGreeks {
156 fn decode_batch(
157 metadata: &HashMap<String, String>,
158 record_batch: RecordBatch,
159 ) -> Result<Vec<Self>, EncodingError> {
160 let instrument_id = metadata
161 .get(KEY_INSTRUMENT_ID)
162 .ok_or(EncodingError::MissingMetadata(KEY_INSTRUMENT_ID))?
163 .parse::<InstrumentId>()
164 .map_err(|e| EncodingError::ParseError(KEY_INSTRUMENT_ID, e.to_string()))?;
165 let record_batch = crate::arrow::record_batch_with_u64_timestamps(&record_batch)?;
166 let record_batch = &record_batch;
167 let cols = record_batch.columns();
168
169 let convention_values = extract_column_string(cols, "convention", 0)?;
170 let delta_values = extract_column::<Float64Array>(cols, "delta", 1, DataType::Float64)?;
171 let gamma_values = extract_column::<Float64Array>(cols, "gamma", 2, DataType::Float64)?;
172 let vega_values = extract_column::<Float64Array>(cols, "vega", 3, DataType::Float64)?;
173 let theta_values = extract_column::<Float64Array>(cols, "theta", 4, DataType::Float64)?;
174 let rho_values = extract_column::<Float64Array>(cols, "rho", 5, DataType::Float64)?;
175 for (column, field) in [
176 (0, "convention"),
177 (1, "delta"),
178 (2, "gamma"),
179 (3, "vega"),
180 (4, "theta"),
181 (5, "rho"),
182 ] {
183 validate_required_column(cols[column].as_ref(), field)?;
184 }
185
186 let mark_iv_values = extract_column::<Float64Array>(cols, "mark_iv", 6, DataType::Float64)?;
187 let bid_iv_values = extract_column::<Float64Array>(cols, "bid_iv", 7, DataType::Float64)?;
188 let ask_iv_values = extract_column::<Float64Array>(cols, "ask_iv", 8, DataType::Float64)?;
189 let underlying_price_values =
190 extract_column::<Float64Array>(cols, "underlying_price", 9, DataType::Float64)?;
191 let open_interest_values =
192 extract_column::<Float64Array>(cols, "open_interest", 10, DataType::Float64)?;
193 let ts_event_values =
194 extract_column::<UInt64Array>(cols, "ts_event", 11, DataType::UInt64)?;
195 let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 12, DataType::UInt64)?;
196
197 let result: Result<Vec<Self>, EncodingError> = (0..record_batch.num_rows())
198 .map(|row| {
199 let convention = GreeksConvention::from_str(convention_values.value(row))
200 .map_err(|e| EncodingError::ParseError("convention", e.to_string()))?;
201
202 Ok(Self {
203 instrument_id,
204 convention,
205 greeks: OptionGreekValues {
206 delta: delta_values.value(row),
207 gamma: gamma_values.value(row),
208 vega: vega_values.value(row),
209 theta: theta_values.value(row),
210 rho: rho_values.value(row),
211 },
212 mark_iv: optional_f64(mark_iv_values, row),
213 bid_iv: optional_f64(bid_iv_values, row),
214 ask_iv: optional_f64(ask_iv_values, row),
215 underlying_price: optional_f64(underlying_price_values, row),
216 open_interest: optional_f64(open_interest_values, row),
217 ts_event: decode_required_timestamp(ts_event_values, "ts_event", row)?,
218 ts_init: decode_required_timestamp(ts_init_values, "ts_init", row)?,
219 })
220 })
221 .collect();
222
223 result
224 }
225}
226
227impl DecodeDataFromRecordBatch for OptionGreeks {
228 fn decode_data_batch(
229 metadata: &HashMap<String, String>,
230 record_batch: RecordBatch,
231 ) -> Result<Vec<Data>, EncodingError> {
232 let greeks = Self::decode_batch(metadata, record_batch)?;
233 Ok(greeks.into_iter().map(Data::from).collect())
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use nautilus_model::{enums::GreeksConvention, identifiers::InstrumentId};
240 use rstest::rstest;
241
242 use super::*;
243 use crate::arrow::KEY_TYPE_NAME;
244
245 #[rstest]
246 fn test_encode_decode_round_trip() {
247 let instrument_id = InstrumentId::from("BTC-20260529-100000-C.OKX");
248 let original = vec![
249 OptionGreeks {
250 instrument_id,
251 convention: GreeksConvention::BlackScholes,
252 greeks: OptionGreekValues {
253 delta: 0.55,
254 gamma: 0.012,
255 vega: 3.4,
256 theta: -1.2,
257 rho: 0.01,
258 },
259 mark_iv: Some(0.64),
260 bid_iv: Some(0.62),
261 ask_iv: Some(0.66),
262 underlying_price: Some(100_000.0),
263 open_interest: Some(42.0),
264 ts_event: 1.into(),
265 ts_init: 2.into(),
266 },
267 OptionGreeks {
268 instrument_id,
269 convention: GreeksConvention::PriceAdjusted,
270 greeks: OptionGreekValues {
271 delta: 0.42,
272 gamma: 0.009,
273 vega: 2.9,
274 theta: -0.9,
275 rho: 0.02,
276 },
277 mark_iv: None,
278 bid_iv: None,
279 ask_iv: None,
280 underlying_price: None,
281 open_interest: None,
282 ts_event: 3.into(),
283 ts_init: 4.into(),
284 },
285 ];
286
287 let metadata = original[0].metadata();
288 let record_batch = OptionGreeks::encode_batch(&metadata, &original).unwrap();
289 let schema = record_batch.schema();
290 let decoded = OptionGreeks::decode_batch(schema.metadata(), record_batch).unwrap();
291
292 assert_eq!(
293 schema
294 .fields()
295 .iter()
296 .map(|field| field.name().as_str())
297 .collect::<Vec<_>>(),
298 vec![
299 "convention",
300 "delta",
301 "gamma",
302 "vega",
303 "theta",
304 "rho",
305 "mark_iv",
306 "bid_iv",
307 "ask_iv",
308 "underlying_price",
309 "open_interest",
310 "ts_event",
311 "ts_init",
312 KEY_IDENTIFIER,
313 ]
314 );
315 assert_eq!(schema.field(0).data_type(), &enum_dictionary_data_type());
316 assert_eq!(
317 schema.metadata(),
318 &HashMap::from([
319 (KEY_TYPE_NAME.to_string(), "OptionGreeks".to_string()),
320 (KEY_INSTRUMENT_ID.to_string(), instrument_id.to_string()),
321 ])
322 );
323 assert_eq!(decoded, original);
324 }
325
326 #[rstest]
327 fn test_decode_batch_rejects_null_required_greek_with_field_and_row() {
328 let original = OptionGreeks {
329 instrument_id: InstrumentId::from("BTC-20260529-100000-C.OKX"),
330 convention: GreeksConvention::BlackScholes,
331 greeks: OptionGreekValues {
332 delta: 0.55,
333 gamma: 0.012,
334 vega: 3.4,
335 theta: -1.2,
336 rho: 0.01,
337 },
338 mark_iv: None,
339 bid_iv: None,
340 ask_iv: None,
341 underlying_price: None,
342 open_interest: None,
343 ts_event: 1.into(),
344 ts_init: 2.into(),
345 };
346
347 let metadata = original.metadata();
348 let record_batch = OptionGreeks::encode_batch(&metadata, &[original]).unwrap();
349 let schema = record_batch.schema();
350
351 let fields = schema
352 .fields()
353 .iter()
354 .enumerate()
355 .map(|(index, field)| {
356 if index == 1 {
357 Arc::new(field.as_ref().clone().with_nullable(true))
358 } else {
359 Arc::clone(field)
360 }
361 })
362 .collect::<Vec<_>>();
363
364 let schema = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
365 let mut columns = record_batch.columns().to_vec();
366 columns[1] = Arc::new(Float64Array::from(vec![None]));
367 let record_batch = RecordBatch::try_new(schema, columns).unwrap();
368
369 let result = OptionGreeks::decode_batch(&metadata, record_batch);
370
371 assert!(matches!(
372 result,
373 Err(EncodingError::ParseError("delta", ref message))
374 if message == "row 0: required value is null"
375 ));
376 }
377}