nautilus_serialization/arrow/
json.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
16use std::{
17    collections::{HashMap, HashSet},
18    fmt::Display,
19    sync::Arc,
20};
21
22use arrow::{
23    array::{
24        Array, ArrayRef, BooleanArray, BooleanBuilder, Float64Array, Float64Builder, StringBuilder,
25        StringDictionaryBuilder, TimestampNanosecondArray, UInt64Array, UInt64Builder,
26    },
27    datatypes::{DataType, Field, Int8Type, Schema},
28    error::ArrowError,
29    record_batch::RecordBatch,
30};
31use serde::{Serialize, de::DeserializeOwned};
32use serde_json::{Map, Number, Value};
33
34use super::{
35    EncodingError, KEY_IDENTIFIER, KEY_TYPE_NAME, StringColumnRef, enum_dictionary_data_type,
36    extract_column, extract_column_string, identifier_array_from_display, json_string_field,
37    metadata_with_type_name,
38};
39
40#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub enum JsonFieldEncoding {
42    Utf8,
43    Utf8Json,
44    EnumDictionary,
45    UInt64,
46    Timestamp,
47    Float64,
48    Boolean,
49}
50
51#[derive(Clone, Copy, Debug, PartialEq, Eq)]
52pub struct JsonFieldSpec {
53    pub name: &'static str,
54    pub encoding: JsonFieldEncoding,
55    pub nullable: bool,
56}
57
58impl JsonFieldSpec {
59    #[must_use]
60    pub const fn utf8(name: &'static str, nullable: bool) -> Self {
61        Self {
62            name,
63            encoding: JsonFieldEncoding::Utf8,
64            nullable,
65        }
66    }
67
68    #[must_use]
69    pub const fn utf8_json(name: &'static str, nullable: bool) -> Self {
70        Self {
71            name,
72            encoding: JsonFieldEncoding::Utf8Json,
73            nullable,
74        }
75    }
76
77    #[must_use]
78    pub const fn enum_dictionary(name: &'static str, nullable: bool) -> Self {
79        Self {
80            name,
81            encoding: JsonFieldEncoding::EnumDictionary,
82            nullable,
83        }
84    }
85
86    #[must_use]
87    pub const fn u64(name: &'static str, nullable: bool) -> Self {
88        Self {
89            name,
90            encoding: JsonFieldEncoding::UInt64,
91            nullable,
92        }
93    }
94
95    #[must_use]
96    pub const fn timestamp(name: &'static str, nullable: bool) -> Self {
97        Self {
98            name,
99            encoding: JsonFieldEncoding::Timestamp,
100            nullable,
101        }
102    }
103
104    #[must_use]
105    pub const fn f64(name: &'static str, nullable: bool) -> Self {
106        Self {
107            name,
108            encoding: JsonFieldEncoding::Float64,
109            nullable,
110        }
111    }
112
113    #[must_use]
114    pub const fn boolean(name: &'static str, nullable: bool) -> Self {
115        Self {
116            name,
117            encoding: JsonFieldEncoding::Boolean,
118            nullable,
119        }
120    }
121
122    fn field(self) -> Field {
123        let data_type = match self.encoding {
124            JsonFieldEncoding::Utf8 | JsonFieldEncoding::Utf8Json => DataType::Utf8,
125            JsonFieldEncoding::EnumDictionary => enum_dictionary_data_type(),
126            JsonFieldEncoding::UInt64 => DataType::UInt64,
127            JsonFieldEncoding::Timestamp => super::timestamp_data_type(),
128            JsonFieldEncoding::Float64 => DataType::Float64,
129            JsonFieldEncoding::Boolean => DataType::Boolean,
130        };
131
132        if self.encoding == JsonFieldEncoding::Utf8Json {
133            json_string_field(self.name, self.nullable)
134        } else {
135            Field::new(self.name, data_type, self.nullable)
136        }
137    }
138}
139
140#[must_use]
141pub fn metadata_for_type(type_name: &'static str) -> HashMap<String, String> {
142    metadata_with_type_name(type_name, None)
143}
144
145#[must_use]
146pub fn schema_for_type(
147    type_name: &'static str,
148    metadata: Option<HashMap<String, String>>,
149    fields: &[JsonFieldSpec],
150) -> Schema {
151    Schema::new_with_metadata(
152        fields
153            .iter()
154            .copied()
155            .map(JsonFieldSpec::field)
156            .collect::<Vec<_>>(),
157        metadata_with_type_name(type_name, metadata),
158    )
159}
160
161#[must_use]
162pub fn schema_for_type_with_identifier(
163    type_name: &'static str,
164    metadata: Option<HashMap<String, String>>,
165    fields: &[JsonFieldSpec],
166) -> Schema {
167    let mut fields = fields.to_vec();
168    fields.push(JsonFieldSpec::utf8(KEY_IDENTIFIER, true));
169    schema_for_type(type_name, metadata, &fields)
170}
171
172/// Encodes typed records into an Arrow record batch with the supplied schema metadata.
173///
174/// # Errors
175///
176/// Returns an error if JSON serialization fails or if a field cannot be encoded into
177/// the requested Arrow column type.
178pub fn encode_batch<'a, T: Serialize + 'a>(
179    type_name: &'static str,
180    metadata: &HashMap<String, String>,
181    data: impl IntoIterator<Item = &'a T>,
182    fields: &[JsonFieldSpec],
183) -> Result<RecordBatch, ArrowError> {
184    if let Some(name) = duplicate_field_name(fields) {
185        return Err(invalid_argument(format!(
186            "Duplicate field specification `{name}`"
187        )));
188    }
189
190    let rows = serialize_rows(data)?;
191    let arrays: Result<Vec<ArrayRef>, ArrowError> = fields
192        .iter()
193        .copied()
194        .map(|field| encode_column(field, &rows))
195        .collect();
196
197    RecordBatch::try_new(
198        Arc::new(schema_for_type(type_name, Some(metadata.clone()), fields)),
199        arrays?,
200    )
201}
202
203/// Encodes typed records with the catalog row identifier column.
204///
205/// # Errors
206///
207/// Returns an error if the number of identifiers differs from the number of data
208/// rows or if any field cannot be encoded into the requested Arrow column type.
209pub fn encode_batch_with_identifier<'a, T, D, I>(
210    type_name: &'static str,
211    metadata: &HashMap<String, String>,
212    data: D,
213    fields: &[JsonFieldSpec],
214    identifiers: impl IntoIterator<Item = I>,
215) -> Result<RecordBatch, ArrowError>
216where
217    T: Serialize + 'a,
218    D: IntoIterator<Item = &'a T>,
219    D::IntoIter: ExactSizeIterator,
220    I: Display,
221{
222    if let Some(name) = duplicate_field_name(fields) {
223        return Err(invalid_argument(format!(
224            "Duplicate field specification `{name}`"
225        )));
226    }
227
228    let data = data.into_iter();
229    let data_len = data.len();
230    let identifier_array = identifier_array_from_display(identifiers);
231    if identifier_array.len() != data_len {
232        return Err(invalid_argument(format!(
233            "identifier values length {} does not match data length {}",
234            identifier_array.len(),
235            data_len
236        )));
237    }
238
239    let rows = serialize_rows(data)?;
240    let mut arrays = fields
241        .iter()
242        .copied()
243        .map(|field| encode_column(field, &rows))
244        .collect::<Result<Vec<ArrayRef>, ArrowError>>()?;
245    arrays.push(Arc::new(identifier_array));
246
247    RecordBatch::try_new(
248        Arc::new(schema_for_type_with_identifier(
249            type_name,
250            Some(metadata.clone()),
251            fields,
252        )),
253        arrays,
254    )
255}
256
257/// Decodes typed records from an Arrow record batch produced by encode_batch.
258///
259/// # Errors
260///
261/// Returns an error if a required column is missing, has the wrong type, contains
262/// invalid JSON, or cannot be deserialized into the target type.
263pub fn decode_batch<T: DeserializeOwned>(
264    metadata: &HashMap<String, String>,
265    record_batch: &RecordBatch,
266    fields: &[JsonFieldSpec],
267    fallback_type_name: Option<&'static str>,
268) -> Result<Vec<T>, EncodingError> {
269    decode_batch_with_metadata_fields(metadata, record_batch, fields, &[], fallback_type_name)
270}
271
272/// Decodes typed records whose `metadata_fields` are stored once in schema metadata rather
273/// than as columns.
274///
275/// # Errors
276///
277/// Returns an error if a metadata field is absent, or for any reason [`decode_batch`] fails.
278pub fn decode_batch_with_metadata_fields<T: DeserializeOwned>(
279    metadata: &HashMap<String, String>,
280    record_batch: &RecordBatch,
281    fields: &[JsonFieldSpec],
282    metadata_fields: &[&'static str],
283    fallback_type_name: Option<&'static str>,
284) -> Result<Vec<T>, EncodingError> {
285    if let Some(name) = duplicate_field_name(fields) {
286        return Err(EncodingError::ParseError(
287            name,
288            "duplicate field specification".to_string(),
289        ));
290    }
291
292    let schema = record_batch.schema();
293    let columns: Result<Vec<_>, EncodingError> = fields
294        .iter()
295        .enumerate()
296        .map(|(expected_index, field)| {
297            let index = column_index(&schema, field.name, expected_index)?;
298            decode_column_ref(record_batch.columns(), *field, index)
299        })
300        .collect();
301    let columns = columns?;
302
303    let metadata_values = metadata_fields
304        .iter()
305        .map(|&name| {
306            metadata
307                .get(name)
308                .map(|value| (name, Value::String(value.clone())))
309                .ok_or(EncodingError::MissingMetadata(name))
310        })
311        .collect::<Result<Vec<_>, _>>()?;
312
313    let mut decoded = Vec::with_capacity(record_batch.num_rows());
314    let type_name = metadata
315        .get(KEY_TYPE_NAME)
316        .cloned()
317        .or_else(|| fallback_type_name.map(str::to_string));
318
319    for row in 0..record_batch.num_rows() {
320        let mut value = Map::new();
321        if let Some(type_name) = &type_name {
322            value.insert("type".to_string(), Value::String(type_name.clone()));
323        }
324
325        for (name, metadata_value) in &metadata_values {
326            value.insert((*name).to_string(), metadata_value.clone());
327        }
328
329        for column in &columns {
330            value.insert(column.name().to_string(), column.to_json(row)?);
331        }
332
333        let json = serde_json::to_vec(&Value::Object(value))
334            .map_err(|e| EncodingError::ParseError("record_batch", format!("row {row}: {e}")))?;
335        decoded.push(
336            serde_json::from_slice(&json).map_err(|e| {
337                EncodingError::ParseError("record_batch", format!("row {row}: {e}"))
338            })?,
339        );
340    }
341
342    Ok(decoded)
343}
344
345fn duplicate_field_name(fields: &[JsonFieldSpec]) -> Option<&'static str> {
346    let mut names = HashSet::with_capacity(fields.len());
347    fields
348        .iter()
349        .find_map(|field| (!names.insert(field.name)).then_some(field.name))
350}
351
352fn column_index(
353    schema: &Schema,
354    name: &'static str,
355    expected_index: usize,
356) -> Result<usize, EncodingError> {
357    let mut matches = schema
358        .fields()
359        .iter()
360        .enumerate()
361        .filter(|(_, field)| field.name() == name);
362    let Some((index, _)) = matches.next() else {
363        return Err(EncodingError::MissingColumn(name, expected_index));
364    };
365
366    if matches.next().is_some() {
367        return Err(EncodingError::ParseError(
368            name,
369            "duplicate column name".to_string(),
370        ));
371    }
372    Ok(index)
373}
374
375fn serialize_rows<'a, T: Serialize + 'a>(
376    data: impl IntoIterator<Item = &'a T>,
377) -> Result<Vec<Map<String, Value>>, ArrowError> {
378    data.into_iter()
379        .map(|item| match serde_json::to_value(item) {
380            Ok(Value::Object(map)) => Ok(map),
381            Ok(_) => Err(invalid_argument(
382                "Expected serialized value to be a JSON object".to_string(),
383            )),
384            Err(e) => Err(invalid_argument(e.to_string())),
385        })
386        .collect()
387}
388
389fn encode_column(
390    field: JsonFieldSpec,
391    rows: &[Map<String, Value>],
392) -> Result<ArrayRef, ArrowError> {
393    match field.encoding {
394        JsonFieldEncoding::Utf8 => encode_utf8_column(field, rows),
395        JsonFieldEncoding::Utf8Json => encode_utf8_json_column(field, rows),
396        JsonFieldEncoding::EnumDictionary => encode_enum_dictionary_column(field, rows),
397        JsonFieldEncoding::UInt64 => encode_u64_column(field, rows),
398        JsonFieldEncoding::Timestamp => encode_timestamp_column(field, rows),
399        JsonFieldEncoding::Float64 => encode_f64_column(field, rows),
400        JsonFieldEncoding::Boolean => encode_bool_column(field, rows),
401    }
402}
403
404fn encode_utf8_column(
405    field: JsonFieldSpec,
406    rows: &[Map<String, Value>],
407) -> Result<ArrayRef, ArrowError> {
408    let mut builder = StringBuilder::new();
409
410    for row in rows {
411        match require_value(field, row.get(field.name))? {
412            Some(value) => builder.append_value(value_to_string(value)?),
413            None => builder.append_null(),
414        }
415    }
416
417    Ok(Arc::new(builder.finish()))
418}
419
420fn encode_utf8_json_column(
421    field: JsonFieldSpec,
422    rows: &[Map<String, Value>],
423) -> Result<ArrayRef, ArrowError> {
424    let mut builder = StringBuilder::new();
425
426    for row in rows {
427        match require_value(field, row.get(field.name))? {
428            Some(value) => builder.append_value(
429                serde_json::to_string(value).map_err(|e| invalid_argument(e.to_string()))?,
430            ),
431            None => builder.append_null(),
432        }
433    }
434
435    Ok(Arc::new(builder.finish()))
436}
437
438fn encode_enum_dictionary_column(
439    field: JsonFieldSpec,
440    rows: &[Map<String, Value>],
441) -> Result<ArrayRef, ArrowError> {
442    let mut builder = StringDictionaryBuilder::<Int8Type>::new();
443
444    for row in rows {
445        match require_value(field, row.get(field.name))? {
446            Some(value) => {
447                builder.append(value_to_string(value)?)?;
448            }
449            None => builder.append_null(),
450        }
451    }
452
453    Ok(Arc::new(builder.finish()))
454}
455
456fn encode_u64_column(
457    field: JsonFieldSpec,
458    rows: &[Map<String, Value>],
459) -> Result<ArrayRef, ArrowError> {
460    let mut builder = UInt64Builder::new();
461
462    for row in rows {
463        match require_value(field, row.get(field.name))? {
464            Some(value) => builder.append_value(parse_u64(value)?),
465            None => builder.append_null(),
466        }
467    }
468
469    Ok(Arc::new(builder.finish()))
470}
471
472fn encode_timestamp_column(
473    field: JsonFieldSpec,
474    rows: &[Map<String, Value>],
475) -> Result<ArrayRef, ArrowError> {
476    let values = rows
477        .iter()
478        .map(|row| {
479            require_value(field, row.get(field.name))?
480                .map(parse_u64)
481                .transpose()
482        })
483        .collect::<Result<Vec<_>, ArrowError>>()?;
484    Ok(Arc::new(super::optional_timestamp_array(values)?))
485}
486
487fn encode_f64_column(
488    field: JsonFieldSpec,
489    rows: &[Map<String, Value>],
490) -> Result<ArrayRef, ArrowError> {
491    let mut builder = Float64Builder::new();
492
493    for row in rows {
494        match require_value(field, row.get(field.name))? {
495            Some(value) => builder.append_value(parse_f64(value)?),
496            None => builder.append_null(),
497        }
498    }
499
500    Ok(Arc::new(builder.finish()))
501}
502
503fn encode_bool_column(
504    field: JsonFieldSpec,
505    rows: &[Map<String, Value>],
506) -> Result<ArrayRef, ArrowError> {
507    let mut builder = BooleanBuilder::new();
508
509    for row in rows {
510        match require_value(field, row.get(field.name))? {
511            Some(value) => builder.append_value(parse_bool(value)?),
512            None => builder.append_null(),
513        }
514    }
515
516    Ok(Arc::new(builder.finish()))
517}
518
519fn require_value(
520    field: JsonFieldSpec,
521    value: Option<&Value>,
522) -> Result<Option<&Value>, ArrowError> {
523    match value {
524        Some(Value::Null) | None if !field.nullable => Err(invalid_argument(format!(
525            "Missing required field `{}`",
526            field.name
527        ))),
528        Some(Value::Null) | None => Ok(None),
529        Some(value) => Ok(Some(value)),
530    }
531}
532
533fn value_to_string(value: &Value) -> Result<String, ArrowError> {
534    match value {
535        Value::String(value) => Ok(value.clone()),
536        Value::Null => Err(invalid_argument("Unexpected null value".to_string())),
537        Value::Bool(_) | Value::Number(_) => Ok(value.to_string()),
538        Value::Array(_) | Value::Object(_) => {
539            serde_json::to_string(value).map_err(|e| invalid_argument(e.to_string()))
540        }
541    }
542}
543
544fn parse_u64(value: &Value) -> Result<u64, ArrowError> {
545    match value {
546        Value::Number(number) => number
547            .as_u64()
548            .ok_or_else(|| invalid_argument(format!("Expected u64, found `{number}`"))),
549        Value::String(value) => value
550            .parse::<u64>()
551            .map_err(|e| invalid_argument(format!("Failed to parse u64 from `{value}`: {e}"))),
552        _ => Err(invalid_argument(format!(
553            "Expected u64-compatible value, found `{value}`"
554        ))),
555    }
556}
557
558fn parse_f64(value: &Value) -> Result<f64, ArrowError> {
559    match value {
560        Value::Number(number) => number
561            .as_f64()
562            .ok_or_else(|| invalid_argument(format!("Expected f64, found `{number}`"))),
563        Value::String(value) => value
564            .parse::<f64>()
565            .map_err(|e| invalid_argument(format!("Failed to parse f64 from `{value}`: {e}"))),
566        _ => Err(invalid_argument(format!(
567            "Expected f64-compatible value, found `{value}`"
568        ))),
569    }
570}
571
572fn parse_bool(value: &Value) -> Result<bool, ArrowError> {
573    match value {
574        Value::Bool(value) => Ok(*value),
575        Value::String(value) => value
576            .parse::<bool>()
577            .map_err(|e| invalid_argument(format!("Failed to parse bool from `{value}`: {e}"))),
578        _ => Err(invalid_argument(format!(
579            "Expected bool-compatible value, found `{value}`"
580        ))),
581    }
582}
583
584enum ColumnRef<'a> {
585    Utf8 {
586        name: &'static str,
587        values: StringColumnRef<'a>,
588    },
589    Utf8Json {
590        name: &'static str,
591        values: StringColumnRef<'a>,
592    },
593    UInt64 {
594        name: &'static str,
595        values: &'a UInt64Array,
596    },
597    Timestamp {
598        name: &'static str,
599        values: &'a TimestampNanosecondArray,
600    },
601    Float64 {
602        name: &'static str,
603        values: &'a Float64Array,
604    },
605    Boolean {
606        name: &'static str,
607        values: &'a BooleanArray,
608    },
609}
610
611impl ColumnRef<'_> {
612    fn name(&self) -> &'static str {
613        match self {
614            Self::Utf8 { name, .. }
615            | Self::Utf8Json { name, .. }
616            | Self::UInt64 { name, .. }
617            | Self::Timestamp { name, .. }
618            | Self::Float64 { name, .. }
619            | Self::Boolean { name, .. } => name,
620        }
621    }
622
623    fn to_json(&self, row: usize) -> Result<Value, EncodingError> {
624        match self {
625            Self::Utf8 { values, .. } => Ok(string_to_json(values, row)),
626            Self::Utf8Json { values, .. } => {
627                if values_is_null(values, row) {
628                    Ok(Value::Null)
629                } else {
630                    serde_json::from_str(values.value(row)).map_err(|e| {
631                        EncodingError::ParseError(self.name(), format!("row {row}: {e}"))
632                    })
633                }
634            }
635            Self::UInt64 { values, .. } => {
636                if values.is_null(row) {
637                    Ok(Value::Null)
638                } else {
639                    Ok(Value::Number(Number::from(values.value(row))))
640                }
641            }
642            Self::Timestamp { values, .. } => {
643                if values.is_null(row) {
644                    Ok(Value::Null)
645                } else {
646                    Ok(Value::Number(Number::from(super::decode_timestamp(
647                        values,
648                        self.name(),
649                        row,
650                    )?)))
651                }
652            }
653            Self::Float64 { values, .. } => f64_to_json(self.name(), values, row),
654            Self::Boolean { values, .. } => {
655                if values.is_null(row) {
656                    Ok(Value::Null)
657                } else {
658                    Ok(Value::Bool(values.value(row)))
659                }
660            }
661        }
662    }
663}
664
665fn decode_column_ref(
666    columns: &[ArrayRef],
667    field: JsonFieldSpec,
668    index: usize,
669) -> Result<ColumnRef<'_>, EncodingError> {
670    match field.encoding {
671        JsonFieldEncoding::Utf8 | JsonFieldEncoding::EnumDictionary => Ok(ColumnRef::Utf8 {
672            name: field.name,
673            values: extract_column_string(columns, field.name, index)?,
674        }),
675        JsonFieldEncoding::Utf8Json => Ok(ColumnRef::Utf8Json {
676            name: field.name,
677            values: extract_column_string(columns, field.name, index)?,
678        }),
679        JsonFieldEncoding::UInt64 => Ok(ColumnRef::UInt64 {
680            name: field.name,
681            values: extract_column::<UInt64Array>(columns, field.name, index, DataType::UInt64)?,
682        }),
683        JsonFieldEncoding::Timestamp => Ok(ColumnRef::Timestamp {
684            name: field.name,
685            values: extract_column::<TimestampNanosecondArray>(
686                columns,
687                field.name,
688                index,
689                super::timestamp_data_type(),
690            )?,
691        }),
692        JsonFieldEncoding::Float64 => Ok(ColumnRef::Float64 {
693            name: field.name,
694            values: extract_column::<Float64Array>(columns, field.name, index, DataType::Float64)?,
695        }),
696        JsonFieldEncoding::Boolean => Ok(ColumnRef::Boolean {
697            name: field.name,
698            values: extract_column::<BooleanArray>(columns, field.name, index, DataType::Boolean)?,
699        }),
700    }
701}
702
703fn string_to_json(values: &StringColumnRef<'_>, row: usize) -> Value {
704    if values_is_null(values, row) {
705        Value::Null
706    } else {
707        Value::String(values.value(row).to_string())
708    }
709}
710
711fn f64_to_json(
712    name: &'static str,
713    values: &Float64Array,
714    row: usize,
715) -> Result<Value, EncodingError> {
716    if values.is_null(row) {
717        return Ok(Value::Null);
718    }
719
720    Number::from_f64(values.value(row))
721        .map(Value::Number)
722        .ok_or_else(|| EncodingError::ParseError(name, format!("row {row}: invalid f64 value")))
723}
724
725fn values_is_null(values: &StringColumnRef<'_>, row: usize) -> bool {
726    values.is_null(row)
727}
728
729fn invalid_argument(message: String) -> ArrowError {
730    ArrowError::InvalidArgumentError(message)
731}
732
733#[cfg(test)]
734mod tests {
735    use rstest::rstest;
736    use serde::Deserialize;
737
738    use super::*;
739
740    const FIELDS: [JsonFieldSpec; 2] = [
741        JsonFieldSpec::u64("left", false),
742        JsonFieldSpec::u64("right", false),
743    ];
744
745    #[derive(Debug, PartialEq, Serialize, Deserialize)]
746    struct Record {
747        left: u64,
748        right: u64,
749    }
750
751    #[rstest]
752    fn decode_batch_matches_columns_by_name() {
753        let rows = [Record {
754            left: 11,
755            right: 29,
756        }];
757        let metadata = HashMap::new();
758        let batch = encode_batch("Record", &metadata, &rows, &FIELDS).unwrap();
759        let reordered = batch.project(&[1, 0]).unwrap();
760
761        let decoded = decode_batch::<Record>(&metadata, &reordered, &FIELDS, None).unwrap();
762
763        assert_eq!(decoded, rows);
764    }
765
766    #[rstest]
767    fn encode_batch_rejects_duplicate_field_specs() {
768        let rows = [Record {
769            left: 11,
770            right: 29,
771        }];
772        let duplicate = [FIELDS[0], FIELDS[0]];
773
774        let error = encode_batch("Record", &HashMap::new(), &rows, &duplicate).unwrap_err();
775
776        assert!(matches!(error, ArrowError::InvalidArgumentError(message)
777            if message == "Duplicate field specification `left`"));
778    }
779
780    #[rstest]
781    fn encode_batch_with_identifier_rejects_duplicate_field_specs() {
782        let rows = [Record {
783            left: 11,
784            right: 29,
785        }];
786        let duplicate = [FIELDS[0], FIELDS[0]];
787
788        let error = encode_batch_with_identifier(
789            "Record",
790            &HashMap::new(),
791            &rows,
792            &duplicate,
793            ["record-1"],
794        )
795        .unwrap_err();
796
797        assert!(matches!(error, ArrowError::InvalidArgumentError(message)
798            if message == "Duplicate field specification `left`"));
799    }
800
801    #[rstest]
802    fn decode_batch_rejects_duplicate_field_specs() {
803        let rows = [Record {
804            left: 11,
805            right: 29,
806        }];
807        let metadata = HashMap::new();
808        let batch = encode_batch("Record", &metadata, &rows, &FIELDS).unwrap();
809        let duplicate = [FIELDS[0], FIELDS[0]];
810
811        let error = decode_batch::<Record>(&metadata, &batch, &duplicate, None).unwrap_err();
812
813        assert!(matches!(error, EncodingError::ParseError("left", message)
814            if message == "duplicate field specification"));
815    }
816
817    #[rstest]
818    fn decode_batch_rejects_duplicate_column_names() {
819        let rows = [Record {
820            left: 11,
821            right: 29,
822        }];
823        let metadata = HashMap::new();
824        let batch = encode_batch("Record", &metadata, &rows, &FIELDS).unwrap();
825        let duplicate = batch.project(&[0, 0, 1]).unwrap();
826
827        let error = decode_batch::<Record>(&metadata, &duplicate, &FIELDS, None).unwrap_err();
828
829        assert!(matches!(error, EncodingError::ParseError("left", message)
830            if message == "duplicate column name"));
831    }
832
833    #[rstest]
834    #[case::missing(None)]
835    #[case::null(Some(Value::Null))]
836    fn encode_batch_rejects_missing_required_value(#[case] value: Option<Value>) {
837        let rows = [value
838            .into_iter()
839            .map(|value| ("left".to_string(), value))
840            .collect::<Map<_, _>>()];
841        let error = encode_batch("Record", &HashMap::new(), &rows, &FIELDS[..1]).unwrap_err();
842
843        assert!(matches!(error, ArrowError::InvalidArgumentError(message)
844            if message == "Missing required field `left`"));
845    }
846
847    #[rstest]
848    fn encode_batch_rejects_non_object() {
849        let error = encode_batch("Record", &HashMap::new(), &[17_u64], &FIELDS).unwrap_err();
850
851        assert!(matches!(error, ArrowError::InvalidArgumentError(message)
852            if message == "Expected serialized value to be a JSON object"));
853    }
854
855    #[rstest]
856    #[case(Value::from(-1), "Expected u64, found `-1`")]
857    #[case(Value::from(true), "Expected u64-compatible value, found `true`")]
858    #[case(
859        Value::from("18446744073709551616"),
860        "Failed to parse u64 from `18446744073709551616`: number too large to fit in target type"
861    )]
862    fn encode_batch_rejects_invalid_u64(#[case] value: Value, #[case] expected: &str) {
863        let rows = [Map::from_iter([("left".to_string(), value)])];
864        let error = encode_batch("Record", &HashMap::new(), &rows, &FIELDS[..1]).unwrap_err();
865
866        assert!(matches!(error, ArrowError::InvalidArgumentError(message) if message == expected));
867    }
868
869    #[rstest]
870    fn encode_decode_string_values_and_nulls() {
871        let fields = [
872            JsonFieldSpec::u64("count", true),
873            JsonFieldSpec::f64("ratio", true),
874            JsonFieldSpec::boolean("active", true),
875        ];
876        let rows = [
877            Map::from_iter([
878                ("count".to_string(), Value::from(u64::MAX.to_string())),
879                ("ratio".to_string(), Value::from("-1.25")),
880                ("active".to_string(), Value::from("false")),
881            ]),
882            Map::new(),
883        ];
884        let metadata = HashMap::new();
885        let batch = encode_batch("Record", &metadata, &rows, &fields).unwrap();
886        let decoded = decode_batch::<Map<String, Value>>(&metadata, &batch, &fields, None).unwrap();
887
888        assert_eq!(
889            decoded,
890            vec![
891                Map::from_iter([
892                    ("count".to_string(), Value::from(u64::MAX)),
893                    ("ratio".to_string(), Value::from(-1.25)),
894                    ("active".to_string(), Value::from(false)),
895                ]),
896                Map::from_iter([
897                    ("count".to_string(), Value::Null),
898                    ("ratio".to_string(), Value::Null),
899                    ("active".to_string(), Value::Null),
900                ])
901            ]
902        );
903    }
904
905    #[rstest]
906    fn encode_decode_enum_dictionary_with_metadata_fields() {
907        let fields = [JsonFieldSpec::enum_dictionary("side", true)];
908        let rows = [
909            Map::from_iter([("side".to_string(), Value::from("BUY"))]),
910            Map::new(),
911            Map::from_iter([("side".to_string(), Value::from("BUY"))]),
912        ];
913        let metadata = HashMap::from([("venue".to_string(), "SIM".to_string())]);
914        let batch = encode_batch("Record", &metadata, &rows, &fields).unwrap();
915        let schema = batch.schema();
916        let decoded = decode_batch_with_metadata_fields::<Map<String, Value>>(
917            schema.metadata(),
918            &batch,
919            &fields,
920            &["venue"],
921            None,
922        )
923        .unwrap();
924
925        assert_eq!(schema.field(0).data_type(), &enum_dictionary_data_type());
926        assert_eq!(schema.metadata()[KEY_TYPE_NAME], "Record");
927        assert_eq!(
928            decoded,
929            vec![
930                Map::from_iter([
931                    ("type".to_string(), Value::from("Record")),
932                    ("venue".to_string(), Value::from("SIM")),
933                    ("side".to_string(), Value::from("BUY")),
934                ]),
935                Map::from_iter([
936                    ("type".to_string(), Value::from("Record")),
937                    ("venue".to_string(), Value::from("SIM")),
938                    ("side".to_string(), Value::Null),
939                ]),
940                Map::from_iter([
941                    ("type".to_string(), Value::from("Record")),
942                    ("venue".to_string(), Value::from("SIM")),
943                    ("side".to_string(), Value::from("BUY")),
944                ]),
945            ]
946        );
947    }
948
949    #[rstest]
950    fn decode_batch_with_metadata_fields_rejects_missing_metadata() {
951        let rows = [Record {
952            left: 11,
953            right: 29,
954        }];
955
956        let batch = encode_batch("Record", &HashMap::new(), &rows, &FIELDS).unwrap();
957
958        let error = decode_batch_with_metadata_fields::<Record>(
959            &HashMap::new(),
960            &batch,
961            &FIELDS,
962            &["venue"],
963            None,
964        )
965        .unwrap_err();
966
967        assert!(matches!(error, EncodingError::MissingMetadata("venue")));
968    }
969
970    #[rstest]
971    #[case(f64::NAN)]
972    #[case(f64::INFINITY)]
973    #[case(f64::NEG_INFINITY)]
974    fn decode_batch_rejects_nonfinite_float(#[case] value: f64) {
975        let fields = [JsonFieldSpec::f64("ratio", false)];
976        let batch = RecordBatch::try_new(
977            Arc::new(schema_for_type("Record", None, &fields)),
978            vec![Arc::new(Float64Array::from(vec![1.25, value]))],
979        )
980        .unwrap();
981        let error =
982            decode_batch::<Map<String, Value>>(&HashMap::new(), &batch, &fields, None).unwrap_err();
983
984        assert!(matches!(error, EncodingError::ParseError("ratio", message)
985            if message == "row 1: invalid f64 value"));
986    }
987}
nautilus_serialization/arrow/json.rs

nautilus_serialization/arrow/
json.rs
