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

nautilus_serialization/arrow/
json.rs
