1use crate::FaucetError;
20use arrow::array::RecordBatch;
21use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
22use serde_json::Value;
23use std::sync::Arc;
24
25#[derive(Debug, Clone)]
33pub struct ColumnarPage {
34 pub batch: RecordBatch,
36 pub bookmark: Option<Value>,
38}
39
40impl ColumnarPage {
41 pub fn new(batch: RecordBatch, bookmark: Option<Value>) -> Self {
43 Self { batch, bookmark }
44 }
45
46 pub fn num_rows(&self) -> usize {
48 self.batch.num_rows()
49 }
50}
51
52fn te<E: std::fmt::Display>(ctx: &str, e: E) -> FaucetError {
54 FaucetError::Transform(format!("columnar shim: {ctx}: {e}"))
55}
56
57pub 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
70fn refine_wide_integers(schema: &Schema, records: &[Value]) -> Result<Schema, FaucetError> {
85 use arrow::datatypes::{DataType, Field};
86
87 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 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 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 _ 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 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#[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
196pub(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
218pub 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
244pub 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
317pub 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
323pub 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
459pub 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 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 #[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 #[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 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 #[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 #[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}