Skip to main content

faucet_core/
columnar.rs

1//! Opt-in Apache Arrow columnar record path (feature `arrow`, RFC 0002 / #375).
2//!
3//! This is the **escape hatch** described in #324: a columnar representation a
4//! connector may produce or consume at the page boundary, *additive* to and
5//! coexisting with the default `serde_json::Value` row model. It touches neither
6//! [`StreamPage`](crate::StreamPage) nor any existing connector — an Arrow-native
7//! source overrides [`Source::stream_batches`](crate::Source::stream_batches) and
8//! an Arrow-native sink overrides
9//! [`Sink::write_batch_columnar`](crate::Sink::write_batch_columnar); the pipeline
10//! uses the columnar path only when **both** sides advertise support (and no
11//! `Value`-shaped stage needs to observe the records), so a
12//! `parquet → parquet` chain never materializes `Value`.
13//!
14//! The `RecordBatch ↔ Value` conversions here are the single source of truth for
15//! the shim (they match `faucet-transform-sql`'s `shovel` byte-for-byte, incl.
16//! `with_explicit_nulls(true)` so an explicit-null field round-trips as
17//! `"key": null` rather than being silently dropped — audit #321 H6).
18
19use crate::FaucetError;
20use arrow::array::RecordBatch;
21use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
22use serde_json::Value;
23use std::sync::Arc;
24
25/// A page of records in **columnar** (Arrow) form, the columnar analogue of
26/// [`StreamPage`](crate::StreamPage).
27///
28/// `bookmark` carries the exact same checkpoint semantics as `StreamPage`:
29/// `Some` triggers flush + bookmark-persist after the batch is durably written;
30/// most sources emit `Some` only on the final batch, CDC-style sources per
31/// committed transaction.
32#[derive(Debug, Clone)]
33pub struct ColumnarPage {
34    /// The record batch to write to the sink for this page.
35    pub batch: RecordBatch,
36    /// Optional bookmark to checkpoint after this batch is durably written.
37    pub bookmark: Option<Value>,
38}
39
40impl ColumnarPage {
41    /// Construct a columnar page from a batch and optional bookmark.
42    pub fn new(batch: RecordBatch, bookmark: Option<Value>) -> Self {
43        Self { batch, bookmark }
44    }
45
46    /// Number of rows in the batch.
47    pub fn num_rows(&self) -> usize {
48        self.batch.num_rows()
49    }
50}
51
52/// Map any display-able error into a [`FaucetError::Transform`] with context.
53fn te<E: std::fmt::Display>(ctx: &str, e: E) -> FaucetError {
54    FaucetError::Transform(format!("columnar shim: {ctx}: {e}"))
55}
56
57/// Infer an Arrow [`Schema`] from a slice of JSON records (each a JSON object).
58///
59/// Returns the inferred schema in an [`Arc`]. An empty slice yields a schema
60/// with no fields.
61pub fn infer_arrow_schema(records: &[Value]) -> Result<SchemaRef, FaucetError> {
62    let iter = records
63        .iter()
64        .map(|v| Ok::<_, arrow::error::ArrowError>(v.clone()));
65    let schema = arrow_json::reader::infer_json_schema_from_iterator(iter)
66        .map_err(|e| te("schema inference", e))?;
67    refine_wide_integers(&schema, records).map(Arc::new)
68}
69
70/// Re-type fields that `arrow-json` widened to `Float64` purely because an
71/// integer did not fit `i64`, and refuse the cases that cannot be re-typed.
72///
73/// `arrow-json` infers `Float64` for any integer outside `i64` range, so a `u64`
74/// id such as `18446744073709551615` silently becomes `1.8446744073709552e19` —
75/// exact value gone, `Ok` returned (#460). Where every observed value for such a
76/// field is a non-negative integer, `UInt64` holds it exactly, so use that.
77/// Where no integer type fits (values spanning negative *and* above `i64::MAX`)
78/// or the value sits somewhere this pass does not re-type (inside a list), fail
79/// with a typed error naming the path rather than approximating it.
80///
81/// Fields arrow-json typed `Float64` because a value genuinely *is* fractional
82/// are left alone, as is a field mixing integers and floats — coercing those to
83/// `Float64` is ordinary JSON-number behaviour, not loss of an exact integer.
84fn refine_wide_integers(schema: &Schema, records: &[Value]) -> Result<Schema, FaucetError> {
85    use arrow::datatypes::{DataType, Field};
86
87    /// Values observed at one field across the page (nulls skipped).
88    fn observed<'a>(records: &'a [Value], name: &str) -> Vec<&'a Value> {
89        records
90            .iter()
91            .filter_map(|r| r.get(name))
92            .filter(|v| !v.is_null())
93            .collect()
94    }
95
96    /// Does this number need more than `i64` *and* is it an exact integer?
97    fn is_wide_integer(v: &Value) -> bool {
98        matches!(v, Value::Number(n) if !n.is_i64() && n.is_u64())
99    }
100
101    fn refine_field(field: &Field, values: Vec<&Value>) -> Result<Field, FaucetError> {
102        match field.data_type() {
103            DataType::Float64 if values.iter().any(|v| is_wide_integer(v)) => {
104                // Every value integral and non-negative → UInt64 is exact.
105                if values
106                    .iter()
107                    .all(|v| matches!(v, Value::Number(n) if n.is_u64()))
108                {
109                    Ok(field.clone().with_data_type(DataType::UInt64))
110                } else {
111                    Err(FaucetError::Transform(format!(
112                        "columnar shim: field {:?} mixes an integer above i64::MAX with values \
113                         no unsigned type can hold, so no exact Arrow type fits. Convert it to a \
114                         string first (a `cast` transform) — it is not silently widened to a \
115                         float because that loses the exact value",
116                        field.name()
117                    )))
118                }
119            }
120            DataType::Struct(children) => {
121                let refined = children
122                    .iter()
123                    .map(|child| {
124                        let child_values = values
125                            .iter()
126                            .filter_map(|v| v.get(child.name()))
127                            .filter(|v| !v.is_null())
128                            .collect();
129                        refine_field(child, child_values).map(Arc::new)
130                    })
131                    .collect::<Result<Vec<_>, _>>()?;
132                Ok(field
133                    .clone()
134                    .with_data_type(DataType::Struct(refined.into())))
135            }
136            // A wide integer inside a list is not re-typed by this pass; refusing
137            // beats writing an approximation nobody asked for.
138            _ if contains_wide_integer_in_container(&values) => {
139                Err(FaucetError::Transform(format!(
140                    "columnar shim: field {:?} contains an integer above i64::MAX inside a list, \
141                     which the columnar path cannot represent exactly. Convert those elements to \
142                     strings first (a `cast` transform)",
143                    field.name()
144                )))
145            }
146            _ => Ok(field.clone()),
147        }
148    }
149
150    /// Any wide integer nested inside an array (at any depth).
151    fn contains_wide_integer_in_container(values: &[&Value]) -> bool {
152        fn walk(v: &Value, in_list: bool) -> bool {
153            match v {
154                Value::Array(items) => items.iter().any(|i| walk(i, true)),
155                Value::Object(map) => map.values().any(|i| walk(i, in_list)),
156                other => in_list && is_wide_integer(other),
157            }
158        }
159        values.iter().any(|v| walk(v, false))
160    }
161
162    let fields = schema
163        .fields()
164        .iter()
165        .map(|f| refine_field(f, observed(records, f.name())).map(Arc::new))
166        .collect::<Result<Vec<_>, _>>()?;
167    Ok(Schema::new(fields).with_metadata(schema.metadata().clone()))
168}
169
170/// Top-level numeric, boolean or temporal columns of `schema` that a masking
171/// pass re-types into strings: a name-based rule covers them (`rewrites`), or
172/// a numeric/boolean column now holds a string in some record (temporal
173/// columns always read back as strings, so only a name rule widens them).
174#[cfg(feature = "masking")]
175pub(crate) fn masked_string_columns(
176    schema: &Schema,
177    records: &[Value],
178    rewrites: impl Fn(&str) -> bool,
179) -> Vec<String> {
180    schema
181        .fields()
182        .iter()
183        .filter(|f| {
184            let t = f.data_type();
185            let scalar = t.is_numeric() || *t == DataType::Boolean;
186            (rewrites(f.name()) && (scalar || t.is_temporal()))
187                || (scalar
188                    && records
189                        .iter()
190                        .any(|r| r.get(f.name()).is_some_and(Value::is_string)))
191        })
192        .map(|f| f.name().clone())
193        .collect()
194}
195
196/// `schema` with the named top-level columns as nullable Utf8.
197pub(crate) fn widen_to_utf8(
198    schema: &SchemaRef,
199    names: &std::collections::BTreeSet<String>,
200) -> SchemaRef {
201    if !schema.fields().iter().any(|f| names.contains(f.name())) {
202        return Arc::clone(schema);
203    }
204    let fields: Vec<Field> = schema
205        .fields()
206        .iter()
207        .map(|f| {
208            if names.contains(f.name()) {
209                Field::new(f.name(), DataType::Utf8, true).with_metadata(f.metadata().clone())
210            } else {
211                f.as_ref().clone()
212            }
213        })
214        .collect();
215    Arc::new(Schema::new(fields).with_metadata(schema.metadata().clone()))
216}
217
218/// Encode a slice of JSON records into a single [`RecordBatch`] against `schema`.
219///
220/// Returns an empty batch if `records` is empty. A fractional number bound
221/// for an integer column is an error (see [`check_integral`]).
222pub fn values_to_record_batch(
223    records: &[Value],
224    schema: SchemaRef,
225) -> Result<RecordBatch, FaucetError> {
226    check_integral(records, &schema)?;
227    let mut decoder = arrow_json::ReaderBuilder::new(schema.clone())
228        .build_decoder()
229        .map_err(|e| te("decoder build", e))?;
230    decoder.serialize(records).map_err(|e| te("encode", e))?;
231    let mut batches = Vec::new();
232    while let Some(b) = decoder.flush().map_err(|e| te("flush", e))? {
233        batches.push(b);
234    }
235    if batches.is_empty() {
236        return Ok(RecordBatch::new_empty(schema));
237    }
238    if batches.len() == 1 {
239        return Ok(batches.pop().unwrap());
240    }
241    arrow::compute::concat_batches(&schema, &batches).map_err(|e| te("concat", e))
242}
243
244/// Refuse a fractional JSON number bound for an integer column of `schema`,
245/// at any depth (struct fields, list items).
246///
247/// arrow-json converts a float into an integer column by truncating it
248/// (`3.7` → `3`), so any JSON → Arrow decode against a declared or inferred
249/// integer type should run this first. Whole floats (`3.0`) pass.
250pub fn check_integral(records: &[Value], schema: &Schema) -> Result<(), FaucetError> {
251    use arrow::datatypes::DataType;
252
253    fn holds_integer(dt: &DataType) -> bool {
254        match dt {
255            DataType::Struct(fields) => fields.iter().any(|f| holds_integer(f.data_type())),
256            DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => {
257                holds_integer(f.data_type())
258            }
259            other => other.is_integer(),
260        }
261    }
262
263    fn walk(value: &Value, dt: &DataType, path: &str) -> Result<(), FaucetError> {
264        match (value, dt) {
265            (Value::Number(n), dt) if dt.is_integer() => {
266                if n.as_i64().is_none()
267                    && n.as_u64().is_none()
268                    && n.as_f64()
269                        .is_none_or(|x| x.fract() != 0.0 || !x.is_finite())
270                {
271                    return Err(FaucetError::Transform(format!(
272                        "field '{path}': {n} is not a whole number, but the column is {dt}"
273                    )));
274                }
275                Ok(())
276            }
277            (Value::Object(map), DataType::Struct(fields)) => {
278                for f in fields {
279                    if let Some(child) = map.get(f.name()) {
280                        walk(child, f.data_type(), &format!("{path}.{}", f.name()))?;
281                    }
282                }
283                Ok(())
284            }
285            (
286                Value::Array(items),
287                DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _),
288            ) => {
289                for item in items {
290                    walk(item, f.data_type(), &format!("{path}[]"))?;
291                }
292                Ok(())
293            }
294            _ => Ok(()),
295        }
296    }
297
298    let fields: Vec<_> = schema
299        .fields()
300        .iter()
301        .filter(|f| holds_integer(f.data_type()))
302        .collect();
303    if fields.is_empty() {
304        return Ok(());
305    }
306    for record in records {
307        let Value::Object(map) = record else { continue };
308        for f in &fields {
309            if let Some(v) = map.get(f.name()) {
310                walk(v, f.data_type(), f.name())?;
311            }
312        }
313    }
314    Ok(())
315}
316
317/// Convenience: infer the schema from `records` and encode them into a batch.
318pub fn values_to_record_batch_inferred(records: &[Value]) -> Result<RecordBatch, FaucetError> {
319    let schema = infer_arrow_schema(records)?;
320    values_to_record_batch(records, schema)
321}
322
323/// Decode a [`RecordBatch`] into JSON objects (one per row).
324///
325/// Uses `arrow-json`'s array writer with **explicit nulls enabled**, so a
326/// null-valued column is emitted as `"key": null` rather than omitted — without
327/// this a `SELECT *`-style identity would silently delete every explicit-null
328/// field (audit #321 H6). An empty batch returns an empty `Vec`.
329///
330/// Decimals are emitted as exact strings (`"12345678901234567890.12"`) rather
331/// than JSON numbers, which `serde_json` would round through `f64`; non-finite
332/// floats become `"NaN"` / `"Infinity"` / `"-Infinity"` instead of `null`.
333pub fn record_batch_to_values(batch: &RecordBatch) -> Result<Vec<Value>, FaucetError> {
334    let mut buf = Vec::new();
335    {
336        let mut writer = arrow_json::writer::WriterBuilder::new()
337            .with_explicit_nulls(true)
338            .with_encoder_factory(Arc::new(ExactEncoderFactory))
339            .build::<_, arrow_json::writer::JsonArray>(&mut buf);
340        writer.write(batch).map_err(|e| te("json write", e))?;
341        writer.finish().map_err(|e| te("json finish", e))?;
342    }
343    serde_json::from_slice(&buf).map_err(|e| te("json parse", e))
344}
345
346#[derive(Debug)]
347struct ExactEncoderFactory;
348
349struct StringEncoder<F: FnMut(usize) -> String>(F);
350
351impl<F: FnMut(usize) -> String> arrow_json::writer::Encoder for StringEncoder<F> {
352    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
353        let text = (self.0)(idx);
354        serde_json::to_writer(&mut *out, &text).expect("writing to a Vec cannot fail");
355    }
356}
357
358struct JsonTextEncoder<F: FnMut(usize) -> String>(F);
359
360impl<F: FnMut(usize) -> String> arrow_json::writer::Encoder for JsonTextEncoder<F> {
361    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
362        out.extend_from_slice((self.0)(idx).as_bytes());
363    }
364}
365
366fn non_finite_json(nan: bool, positive: bool) -> String {
367    let text = if nan {
368        "NaN"
369    } else if positive {
370        "Infinity"
371    } else {
372        "-Infinity"
373    };
374    format!("\"{text}\"")
375}
376
377fn f64_json(v: f64) -> String {
378    if v.is_finite() {
379        Value::from(v).to_string()
380    } else {
381        non_finite_json(v.is_nan(), v.is_sign_positive())
382    }
383}
384
385fn f32_json(v: f32) -> String {
386    if v.is_finite() {
387        serde_json::to_string(&v).expect("a finite f32 always serializes")
388    } else {
389        non_finite_json(v.is_nan(), v.is_sign_positive())
390    }
391}
392
393fn nullable<'a>(
394    encoder: impl arrow_json::writer::Encoder + 'a,
395    array: &'a dyn arrow::array::Array,
396) -> arrow_json::writer::NullableEncoder<'a> {
397    arrow_json::writer::NullableEncoder::new(Box::new(encoder), array.nulls().cloned())
398}
399
400impl arrow_json::writer::EncoderFactory for ExactEncoderFactory {
401    fn make_default_encoder<'a>(
402        &self,
403        _field: &'a arrow::datatypes::FieldRef,
404        array: &'a dyn arrow::array::Array,
405        _options: &'a arrow_json::writer::EncoderOptions,
406    ) -> Result<Option<arrow_json::writer::NullableEncoder<'a>>, arrow::error::ArrowError> {
407        use arrow::array::AsArray;
408        use arrow::datatypes::{
409            DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Float16Type,
410            Float32Type, Float64Type,
411        };
412        let enc = match array.data_type() {
413            DataType::Decimal32(..) => {
414                let a = array.as_primitive::<Decimal32Type>();
415                nullable(StringEncoder(move |i| a.value_as_string(i)), array)
416            }
417            DataType::Decimal64(..) => {
418                let a = array.as_primitive::<Decimal64Type>();
419                nullable(StringEncoder(move |i| a.value_as_string(i)), array)
420            }
421            DataType::Decimal128(..) => {
422                let a = array.as_primitive::<Decimal128Type>();
423                nullable(StringEncoder(move |i| a.value_as_string(i)), array)
424            }
425            DataType::Decimal256(..) => {
426                let a = array.as_primitive::<Decimal256Type>();
427                nullable(StringEncoder(move |i| a.value_as_string(i)), array)
428            }
429            DataType::Float64 => {
430                let a = array.as_primitive::<Float64Type>();
431                if a.values().iter().all(|v| v.is_finite()) {
432                    return Ok(None);
433                }
434                nullable(JsonTextEncoder(move |i| f64_json(a.value(i))), array)
435            }
436            DataType::Float32 => {
437                let a = array.as_primitive::<Float32Type>();
438                if a.values().iter().all(|v| v.is_finite()) {
439                    return Ok(None);
440                }
441                nullable(JsonTextEncoder(move |i| f32_json(a.value(i))), array)
442            }
443            DataType::Float16 => {
444                let a = array.as_primitive::<Float16Type>();
445                if a.values().iter().all(|v| v.is_finite()) {
446                    return Ok(None);
447                }
448                nullable(
449                    JsonTextEncoder(move |i| f32_json(f32::from(a.value(i)))),
450                    array,
451                )
452            }
453            _ => return Ok(None),
454        };
455        Ok(Some(enc))
456    }
457}
458
459/// Compare two schemas for field-level equality (name + data-type + nullability).
460pub fn schema_eq(a: &Schema, b: &Schema) -> bool {
461    a.fields() == b.fields()
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467    use serde_json::json;
468
469    #[test]
470    fn a_fractional_number_in_an_integer_column_is_refused() {
471        use arrow::datatypes::{DataType, Field, Fields};
472        let schema = Arc::new(Schema::new(vec![
473            Field::new("n", DataType::Int64, true),
474            Field::new(
475                "s",
476                DataType::Struct(Fields::from(vec![Field::new("m", DataType::Int32, true)])),
477                true,
478            ),
479            Field::new(
480                "l",
481                DataType::List(Arc::new(Field::new("item", DataType::UInt8, true))),
482                true,
483            ),
484            Field::new("f", DataType::Float64, true),
485        ]));
486        let ok = [
487            json!({"n": 3.0, "s": {"m": 2}, "l": [1, 2], "f": 1.5}),
488            json!("x"),
489        ];
490        let batch = values_to_record_batch(&ok[..1], schema.clone()).unwrap();
491        assert_eq!(batch.num_rows(), 1);
492        check_integral(&ok, &schema).unwrap();
493        for (bad, field) in [
494            (json!({"n": 3.7}), "n"),
495            (json!({"s": {"m": 0.5}}), "s.m"),
496            (json!({"l": [1, 2.25]}), "l[]"),
497        ] {
498            let err = values_to_record_batch(&[bad], schema.clone()).unwrap_err();
499            assert!(err.to_string().contains(&format!("'{field}'")), "{err}");
500        }
501        let floats = Schema::new(vec![Field::new("f", DataType::Float64, true)]);
502        check_integral(&[json!({"f": 0.5})], &floats).unwrap();
503    }
504
505    #[test]
506    fn round_trip_scalars_nulls_nested() {
507        let recs = vec![
508            json!({"id": 1, "name": "a", "score": 1.5, "ok": true, "tags": ["x", "y"]}),
509            json!({"id": 2, "name": null, "score": null, "ok": false, "tags": []}),
510        ];
511        let batch = values_to_record_batch_inferred(&recs).unwrap();
512        assert_eq!(batch.num_rows(), 2);
513        let back = record_batch_to_values(&batch).unwrap();
514        assert_eq!(back[0]["id"], json!(1));
515        assert_eq!(back[0]["tags"], json!(["x", "y"]));
516        // #321 H6: an explicit-null field survives the round-trip.
517        assert!(back[1].as_object().unwrap().contains_key("name"));
518        assert_eq!(back[1]["name"], json!(null));
519    }
520
521    #[test]
522    fn empty_records_yield_empty_batch_and_back() {
523        let schema = infer_arrow_schema(&[json!({"a": 1})]).unwrap();
524        let batch = values_to_record_batch(&[], schema).unwrap();
525        assert_eq!(batch.num_rows(), 0);
526        assert!(record_batch_to_values(&batch).unwrap().is_empty());
527    }
528
529    fn decimal_batch() -> RecordBatch {
530        use arrow::array::Decimal128Array;
531        use arrow::datatypes::{DataType, Field};
532        let amounts = Decimal128Array::from(vec![
533            Some(12_345_678_901_234_567_890_123_456_789_012_345_678_i128),
534            None,
535            Some(-1),
536        ])
537        .with_precision_and_scale(38, 2)
538        .unwrap();
539        let schema = Arc::new(Schema::new(vec![Field::new(
540            "amount",
541            DataType::Decimal128(38, 2),
542            true,
543        )]));
544        RecordBatch::try_new(schema, vec![Arc::new(amounts)]).unwrap()
545    }
546
547    #[test]
548    fn decimals_survive_as_exact_strings_and_round_trip() {
549        let batch = decimal_batch();
550        let values = record_batch_to_values(&batch).unwrap();
551        assert_eq!(
552            values[0]["amount"],
553            json!("123456789012345678901234567890123456.78")
554        );
555        assert_eq!(values[1]["amount"], json!(null));
556        assert_eq!(values[2]["amount"], json!("-0.01"));
557        let back = values_to_record_batch(&values, batch.schema()).unwrap();
558        assert_eq!(back, batch);
559    }
560
561    #[test]
562    fn nested_decimals_are_exact_too() {
563        use arrow::array::{Decimal128Array, ListArray};
564        use arrow::buffer::OffsetBuffer;
565        use arrow::datatypes::{DataType, Field};
566        let item = Arc::new(Field::new("item", DataType::Decimal128(38, 0), true));
567        let values = Decimal128Array::from(vec![99_999_999_999_999_999_999_999_i128])
568            .with_precision_and_scale(38, 0)
569            .unwrap();
570        let list = ListArray::new(
571            item.clone(),
572            OffsetBuffer::from_lengths([1]),
573            Arc::new(values),
574            None,
575        );
576        let schema = Arc::new(Schema::new(vec![Field::new(
577            "ids",
578            DataType::List(item),
579            true,
580        )]));
581        let batch = RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap();
582        let rows = record_batch_to_values(&batch).unwrap();
583        assert_eq!(rows[0]["ids"], json!(["99999999999999999999999"]));
584    }
585
586    #[test]
587    fn non_finite_floats_become_strings_not_nulls() {
588        use arrow::array::{Float32Array, Float64Array};
589        use arrow::datatypes::{DataType, Field};
590        let halves = arrow::compute::cast(
591            &Float32Array::from(vec![Some(f32::INFINITY), Some(0.5), None, Some(f32::NAN)]),
592            &DataType::Float16,
593        )
594        .unwrap();
595        let schema = Arc::new(Schema::new(vec![
596            Field::new("d", DataType::Float64, true),
597            Field::new("f", DataType::Float32, true),
598            Field::new("h", DataType::Float16, true),
599        ]));
600        let batch = RecordBatch::try_new(
601            schema.clone(),
602            vec![
603                Arc::new(Float64Array::from(vec![
604                    Some(f64::NAN),
605                    Some(f64::INFINITY),
606                    Some(1.5),
607                    None,
608                ])),
609                Arc::new(Float32Array::from(vec![
610                    Some(f32::NEG_INFINITY),
611                    Some(0.1),
612                    Some(f32::NAN),
613                    Some(2.0),
614                ])),
615                halves,
616            ],
617        )
618        .unwrap();
619        let rows = record_batch_to_values(&batch).unwrap();
620        assert_eq!(rows[0]["d"], json!("NaN"));
621        assert_eq!(rows[1]["d"], json!("Infinity"));
622        assert_eq!(rows[2]["d"], json!(1.5));
623        assert_eq!(rows[3]["d"], json!(null));
624        assert_eq!(rows[0]["f"], json!("-Infinity"));
625        assert_eq!(rows[1]["f"].to_string(), "0.1");
626        assert_eq!(rows[2]["f"], json!("NaN"));
627        assert_eq!(rows[0]["h"], json!("Infinity"));
628        assert_eq!(rows[1]["h"], json!(0.5));
629        assert_eq!(rows[2]["h"], json!(null));
630        assert_eq!(rows[3]["h"], json!("NaN"));
631
632        let back = values_to_record_batch(&rows, schema).unwrap();
633        let d = back
634            .column(0)
635            .as_any()
636            .downcast_ref::<Float64Array>()
637            .unwrap();
638        assert!(d.value(0).is_nan());
639        assert_eq!(d.value(1), f64::INFINITY);
640        assert_eq!(d.value(2), 1.5);
641    }
642
643    #[test]
644    fn all_finite_floats_keep_the_default_encoding() {
645        let batch = values_to_record_batch_inferred(&[json!({"x": 0.25})]).unwrap();
646        assert_eq!(record_batch_to_values(&batch).unwrap()[0]["x"], json!(0.25));
647    }
648
649    #[test]
650    fn columnar_page_reports_rows() {
651        let batch = values_to_record_batch_inferred(&[json!({"a": 1}), json!({"a": 2})]).unwrap();
652        let page = ColumnarPage::new(batch, Some(json!({"lsn": 42})));
653        assert_eq!(page.num_rows(), 2);
654        assert_eq!(page.bookmark, Some(json!({"lsn": 42})));
655    }
656}
657
658#[cfg(test)]
659mod wide_integer_tests {
660    use super::*;
661    use serde_json::json;
662
663    /// #460: arrow-json types an integer above `i64::MAX` as `Float64`, so
664    /// `18446744073709551615` used to come back as `1.8446744073709552e19` with
665    /// `Ok`. It must now round-trip exactly.
666    #[test]
667    fn u64_above_i64_max_round_trips_exactly() {
668        let recs = vec![json!({"id": u64::MAX})];
669        let back =
670            record_batch_to_values(&values_to_record_batch_inferred(&recs).unwrap()).unwrap();
671        assert_eq!(back[0]["id"], json!(u64::MAX), "exact value must survive");
672        assert!(back[0]["id"].is_u64(), "and stay an integer, not a float");
673    }
674
675    #[test]
676    fn mixed_small_and_wide_integers_all_survive() {
677        let recs = vec![
678            json!({"id": 1}),
679            json!({"id": u64::MAX}),
680            json!({"id": null}),
681        ];
682        let back =
683            record_batch_to_values(&values_to_record_batch_inferred(&recs).unwrap()).unwrap();
684        assert_eq!(back[0]["id"], json!(1));
685        assert_eq!(back[1]["id"], json!(u64::MAX));
686        assert_eq!(back[2]["id"], json!(null));
687    }
688
689    #[test]
690    fn wide_integer_nested_in_a_struct_survives() {
691        let recs = vec![json!({"outer": {"id": u64::MAX, "n": 3}})];
692        let back =
693            record_batch_to_values(&values_to_record_batch_inferred(&recs).unwrap()).unwrap();
694        assert_eq!(back[0]["outer"]["id"], json!(u64::MAX));
695        assert_eq!(back[0]["outer"]["n"], json!(3));
696    }
697
698    /// Genuine floats, and integer/float mixes, keep their existing behaviour —
699    /// coercing those to Float64 is ordinary JSON-number semantics, not loss of
700    /// an exact integer.
701    #[test]
702    fn genuine_floats_are_untouched() {
703        let recs = vec![json!({"a": 1.5}), json!({"a": 2})];
704        let back =
705            record_batch_to_values(&values_to_record_batch_inferred(&recs).unwrap()).unwrap();
706        assert_eq!(back[0]["a"], json!(1.5));
707        assert_eq!(back[1]["a"], json!(2.0));
708
709        // i64 range is unaffected (2^53+1 already round-tripped before).
710        let recs = vec![json!({"n": 9007199254740993i64, "m": i64::MIN})];
711        let back =
712            record_batch_to_values(&values_to_record_batch_inferred(&recs).unwrap()).unwrap();
713        assert_eq!(back[0]["n"], json!(9007199254740993i64));
714        assert_eq!(back[0]["m"], json!(i64::MIN));
715    }
716
717    /// No exact type fits a field spanning negatives and above-i64::MAX, so it
718    /// must fail loudly rather than approximate.
719    #[test]
720    fn unrepresentable_mix_is_refused() {
721        let recs = vec![json!({"v": -1}), json!({"v": u64::MAX})];
722        let err = match values_to_record_batch_inferred(&recs) {
723            Err(e) => e.to_string(),
724            Ok(b) => panic!("must refuse; got {:?}", record_batch_to_values(&b).unwrap()),
725        };
726        assert!(err.contains("\"v\""), "{err}");
727        assert!(err.contains("cast"), "points at the workaround: {err}");
728    }
729
730    /// A wide integer inside a list is not re-typed by this pass, so it is
731    /// refused rather than silently written as a float.
732    #[test]
733    fn wide_integer_inside_a_list_is_refused() {
734        let recs = vec![json!({"ids": [1, u64::MAX]})];
735        let err = match values_to_record_batch_inferred(&recs) {
736            Err(e) => e.to_string(),
737            Ok(_) => panic!("must refuse a wide integer inside a list"),
738        };
739        assert!(err.contains("\"ids\""), "{err}");
740        assert!(err.contains("list"), "{err}");
741    }
742}