1use std::borrow::Cow;
137use std::io::BufRead;
138use std::sync::Arc;
139
140use arrow_array::cast::AsArray;
141use arrow_array::timezone::Tz;
142use arrow_array::types::*;
143use arrow_array::{ArrayRef, RecordBatch, RecordBatchReader, downcast_integer};
144use arrow_schema::{ArrowError, DataType, FieldRef, Schema, SchemaRef, TimeUnit};
145use chrono::Utc;
146use serde_core::Serialize;
147
148use crate::StructMode;
149use crate::reader::binary_array::{
150 BinaryArrayDecoder, BinaryViewDecoder, FixedSizeBinaryArrayDecoder,
151};
152use crate::reader::boolean_array::BooleanArrayDecoder;
153use crate::reader::decimal_array::DecimalArrayDecoder;
154use crate::reader::list_array::{
155 FixedSizeListArrayDecoder, ListArrayDecoder, ListViewArrayDecoder,
156};
157use crate::reader::map_array::MapArrayDecoder;
158use crate::reader::null_array::NullArrayDecoder;
159use crate::reader::primitive_array::PrimitiveArrayDecoder;
160use crate::reader::run_end_array::RunEndEncodedArrayDecoder;
161use crate::reader::string_array::StringArrayDecoder;
162use crate::reader::string_view_array::StringViewArrayDecoder;
163use crate::reader::struct_array::StructArrayDecoder;
164use crate::reader::tape::{Tape, TapeDecoder};
165use crate::reader::timestamp_array::TimestampArrayDecoder;
166
167pub use schema::*;
168pub use value_iter::ValueIter;
169
170mod binary_array;
171mod boolean_array;
172mod decimal_array;
173mod list_array;
174mod map_array;
175mod null_array;
176mod primitive_array;
177mod run_end_array;
178mod schema;
179mod serializer;
180mod string_array;
181mod string_view_array;
182mod struct_array;
183mod tape;
184mod timestamp_array;
185mod value_iter;
186
187pub struct ReaderBuilder {
189 batch_size: usize,
190 coerce_primitive: bool,
191 strict_mode: bool,
192 ignore_type_conflicts: bool,
193 is_field: bool,
194 struct_mode: StructMode,
195
196 schema: SchemaRef,
197}
198
199impl ReaderBuilder {
200 pub fn new(schema: SchemaRef) -> Self {
209 Self {
210 batch_size: 1024,
211 coerce_primitive: false,
212 strict_mode: false,
213 ignore_type_conflicts: false,
214 is_field: false,
215 struct_mode: Default::default(),
216 schema,
217 }
218 }
219
220 pub fn new_with_field(field: impl Into<FieldRef>) -> Self {
251 Self {
252 batch_size: 1024,
253 coerce_primitive: false,
254 strict_mode: false,
255 ignore_type_conflicts: false,
256 is_field: true,
257 struct_mode: Default::default(),
258 schema: Arc::new(Schema::new([field.into()])),
259 }
260 }
261
262 pub fn with_batch_size(self, batch_size: usize) -> Self {
264 Self { batch_size, ..self }
265 }
266
267 pub fn with_coerce_primitive(self, coerce_primitive: bool) -> Self {
270 Self {
271 coerce_primitive,
272 ..self
273 }
274 }
275
276 pub fn with_strict_mode(self, strict_mode: bool) -> Self {
282 Self {
283 strict_mode,
284 ..self
285 }
286 }
287
288 pub fn with_struct_mode(self, struct_mode: StructMode) -> Self {
292 Self {
293 struct_mode,
294 ..self
295 }
296 }
297
298 pub fn with_ignore_type_conflicts(self, ignore_type_conflicts: bool) -> Self {
311 Self {
312 ignore_type_conflicts,
313 ..self
314 }
315 }
316
317 pub fn build<R: BufRead>(self, reader: R) -> Result<Reader<R>, ArrowError> {
319 Ok(Reader {
320 reader,
321 decoder: self.build_decoder()?,
322 })
323 }
324
325 pub fn build_decoder(self) -> Result<Decoder, ArrowError> {
327 let (data_type, nullable) = if self.is_field {
328 let field = &self.schema.fields[0];
329 let data_type = Cow::Borrowed(field.data_type());
330 (data_type, field.is_nullable())
331 } else {
332 let data_type = Cow::Owned(DataType::Struct(self.schema.fields.clone()));
333 (data_type, false)
334 };
335
336 let ctx = DecoderContext {
337 coerce_primitive: self.coerce_primitive,
338 strict_mode: self.strict_mode,
339 struct_mode: self.struct_mode,
340 ignore_type_conflicts: self.ignore_type_conflicts,
341 };
342 let decoder = ctx.make_decoder(data_type.as_ref(), nullable)?;
343
344 let num_fields = self.schema.flattened_fields().len();
345
346 Ok(Decoder {
347 decoder,
348 is_field: self.is_field,
349 tape_decoder: TapeDecoder::new(self.batch_size, num_fields),
350 batch_size: self.batch_size,
351 schema: self.schema,
352 })
353 }
354}
355
356pub struct Reader<R> {
360 reader: R,
361 decoder: Decoder,
362}
363
364impl<R> std::fmt::Debug for Reader<R> {
365 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
366 f.debug_struct("Reader")
367 .field("decoder", &self.decoder)
368 .finish()
369 }
370}
371
372impl<R: BufRead> Reader<R> {
373 fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
375 loop {
376 let buf = self.reader.fill_buf()?;
377 if buf.is_empty() {
378 break;
379 }
380 let read = buf.len();
381
382 let decoded = self.decoder.decode(buf)?;
383 self.reader.consume(decoded);
384 if decoded != read {
385 break;
386 }
387 }
388 self.decoder.flush()
389 }
390}
391
392impl<R: BufRead> Iterator for Reader<R> {
393 type Item = Result<RecordBatch, ArrowError>;
394
395 fn next(&mut self) -> Option<Self::Item> {
396 self.read().transpose()
397 }
398}
399
400impl<R: BufRead> RecordBatchReader for Reader<R> {
401 fn schema(&self) -> SchemaRef {
402 self.decoder.schema.clone()
403 }
404}
405
406pub struct Decoder {
447 tape_decoder: TapeDecoder,
448 decoder: Box<dyn ArrayDecoder>,
449 batch_size: usize,
450 is_field: bool,
451 schema: SchemaRef,
452}
453
454impl std::fmt::Debug for Decoder {
455 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
456 f.debug_struct("Decoder")
457 .field("schema", &self.schema)
458 .field("batch_size", &self.batch_size)
459 .finish()
460 }
461}
462
463impl Decoder {
464 pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
473 self.tape_decoder.decode(buf)
474 }
475
476 pub fn serialize<S: Serialize>(&mut self, rows: &[S]) -> Result<(), ArrowError> {
653 self.tape_decoder.serialize(rows)
654 }
655
656 pub fn has_partial_record(&self) -> bool {
658 self.tape_decoder.has_partial_row()
659 }
660
661 pub fn len(&self) -> usize {
663 self.tape_decoder.num_buffered_rows()
664 }
665
666 pub fn is_empty(&self) -> bool {
668 self.len() == 0
669 }
670
671 pub fn flush(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
678 let tape = self.tape_decoder.finish()?;
679
680 if tape.num_rows() == 0 {
681 return Ok(None);
682 }
683
684 let mut next_object = 1;
686 let pos: Vec<_> = (0..tape.num_rows())
687 .map(|_| {
688 let next = tape.next(next_object, "row").unwrap();
689 std::mem::replace(&mut next_object, next)
690 })
691 .collect();
692
693 let decoded = self.decoder.decode(&tape, &pos)?;
694 self.tape_decoder.clear();
695
696 let batch = match self.is_field {
697 true => RecordBatch::try_new(self.schema.clone(), vec![decoded])?,
698 false => {
699 RecordBatch::from(decoded.as_struct().clone()).with_schema(self.schema.clone())?
700 }
701 };
702
703 Ok(Some(batch))
704 }
705}
706
707trait ArrayDecoder: Send {
708 fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError>;
710}
711
712pub struct DecoderContext {
717 coerce_primitive: bool,
719 strict_mode: bool,
721 struct_mode: StructMode,
723 ignore_type_conflicts: bool,
725}
726
727impl DecoderContext {
728 pub fn coerce_primitive(&self) -> bool {
730 self.coerce_primitive
731 }
732
733 pub fn strict_mode(&self) -> bool {
735 self.strict_mode
736 }
737
738 pub fn struct_mode(&self) -> StructMode {
740 self.struct_mode
741 }
742
743 pub fn ignore_type_conflicts(&self) -> bool {
745 self.ignore_type_conflicts
746 }
747
748 fn make_decoder(
753 &self,
754 data_type: &DataType,
755 is_nullable: bool,
756 ) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
757 make_decoder(self, data_type, is_nullable)
758 }
759}
760
761fn make_decoder(
762 ctx: &DecoderContext,
763 data_type: &DataType,
764 is_nullable: bool,
765) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
766 macro_rules! primitive_decoder {
767 ($t:ty, $data_type:expr) => {
768 Ok(Box::new(PrimitiveArrayDecoder::<$t>::new(ctx, $data_type)))
769 };
770 }
771 macro_rules! timestamp_decoder {
772 ($t:ty, $data_type:expr, $tz:expr) => {{
773 Ok(Box::new(TimestampArrayDecoder::<$t, _>::new(
774 ctx, $data_type, $tz,
775 )))
776 }};
777 }
778 macro_rules! decimal_decoder {
779 ($t:ty, $p:expr, $s:expr) => {
780 Ok(Box::new(DecimalArrayDecoder::<$t>::new(ctx, $p, $s)))
781 };
782 }
783
784 downcast_integer! {
785 *data_type => (primitive_decoder, data_type),
786 DataType::Null => Ok(Box::new(NullArrayDecoder::new(ctx))),
787 DataType::Float16 => primitive_decoder!(Float16Type, data_type),
788 DataType::Float32 => primitive_decoder!(Float32Type, data_type),
789 DataType::Float64 => primitive_decoder!(Float64Type, data_type),
790 DataType::Timestamp(TimeUnit::Second, None) => {
791 timestamp_decoder!(TimestampSecondType, data_type, Utc)
792 },
793 DataType::Timestamp(TimeUnit::Millisecond, None) => {
794 timestamp_decoder!(TimestampMillisecondType, data_type, Utc)
795 },
796 DataType::Timestamp(TimeUnit::Microsecond, None) => {
797 timestamp_decoder!(TimestampMicrosecondType, data_type, Utc)
798 },
799 DataType::Timestamp(TimeUnit::Nanosecond, None) => {
800 timestamp_decoder!(TimestampNanosecondType, data_type, Utc)
801 },
802 DataType::Timestamp(TimeUnit::Second, Some(ref tz)) => {
803 let tz: Tz = tz.parse()?;
804 timestamp_decoder!(TimestampSecondType, data_type, tz)
805 },
806 DataType::Timestamp(TimeUnit::Millisecond, Some(ref tz)) => {
807 let tz: Tz = tz.parse()?;
808 timestamp_decoder!(TimestampMillisecondType, data_type, tz)
809 },
810 DataType::Timestamp(TimeUnit::Microsecond, Some(ref tz)) => {
811 let tz: Tz = tz.parse()?;
812 timestamp_decoder!(TimestampMicrosecondType, data_type, tz)
813 },
814 DataType::Timestamp(TimeUnit::Nanosecond, Some(ref tz)) => {
815 let tz: Tz = tz.parse()?;
816 timestamp_decoder!(TimestampNanosecondType, data_type, tz)
817 },
818 DataType::Date32 => primitive_decoder!(Date32Type, data_type),
819 DataType::Date64 => primitive_decoder!(Date64Type, data_type),
820 DataType::Time32(TimeUnit::Second) => primitive_decoder!(Time32SecondType, data_type),
821 DataType::Time32(TimeUnit::Millisecond) => primitive_decoder!(Time32MillisecondType, data_type),
822 DataType::Time64(TimeUnit::Microsecond) => primitive_decoder!(Time64MicrosecondType, data_type),
823 DataType::Time64(TimeUnit::Nanosecond) => primitive_decoder!(Time64NanosecondType, data_type),
824 DataType::Duration(TimeUnit::Nanosecond) => primitive_decoder!(DurationNanosecondType, data_type),
825 DataType::Duration(TimeUnit::Microsecond) => primitive_decoder!(DurationMicrosecondType, data_type),
826 DataType::Duration(TimeUnit::Millisecond) => primitive_decoder!(DurationMillisecondType, data_type),
827 DataType::Duration(TimeUnit::Second) => primitive_decoder!(DurationSecondType, data_type),
828 DataType::Decimal32(p, s) => decimal_decoder!(Decimal32Type, p, s),
829 DataType::Decimal64(p, s) => decimal_decoder!(Decimal64Type, p, s),
830 DataType::Decimal128(p, s) => decimal_decoder!(Decimal128Type, p, s),
831 DataType::Decimal256(p, s) => decimal_decoder!(Decimal256Type, p, s),
832 DataType::Boolean => Ok(Box::new(BooleanArrayDecoder::new(ctx))),
833 DataType::Utf8 => Ok(Box::new(StringArrayDecoder::<i32>::new(ctx))),
834 DataType::Utf8View => Ok(Box::new(StringViewArrayDecoder::new(ctx))),
835 DataType::LargeUtf8 => Ok(Box::new(StringArrayDecoder::<i64>::new(ctx))),
836 DataType::List(_) => Ok(Box::new(ListArrayDecoder::<i32>::new(ctx, data_type, is_nullable)?)),
837 DataType::LargeList(_) => Ok(Box::new(ListArrayDecoder::<i64>::new(ctx, data_type, is_nullable)?)),
838 DataType::ListView(_) => Ok(Box::new(ListViewArrayDecoder::<i32>::new(ctx, data_type, is_nullable)?)),
839 DataType::LargeListView(_) => Ok(Box::new(ListViewArrayDecoder::<i64>::new(ctx, data_type, is_nullable)?)),
840 DataType::FixedSizeList(_, _) => Ok(Box::new(FixedSizeListArrayDecoder::new(ctx, data_type, is_nullable)?)),
841 DataType::Struct(_) => Ok(Box::new(StructArrayDecoder::new(ctx, data_type, is_nullable)?)),
842 DataType::Binary => Ok(Box::new(BinaryArrayDecoder::<i32>::default())),
843 DataType::LargeBinary => Ok(Box::new(BinaryArrayDecoder::<i64>::default())),
844 DataType::FixedSizeBinary(len) => Ok(Box::new(FixedSizeBinaryArrayDecoder::new(len))),
845 DataType::BinaryView => Ok(Box::new(BinaryViewDecoder::default())),
846 DataType::Map(_, _) => Ok(Box::new(MapArrayDecoder::new(ctx, data_type, is_nullable)?)),
847 DataType::RunEndEncoded(ref r, _) => match r.data_type() {
848 DataType::Int16 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int16Type>::new(ctx, data_type, is_nullable)?)),
849 DataType::Int32 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int32Type>::new(ctx, data_type, is_nullable)?)),
850 DataType::Int64 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int64Type>::new(ctx, data_type, is_nullable)?)),
851 d => unreachable!("unsupported run end index type: {d}"),
852 },
853 _ => Err(ArrowError::NotYetImplemented(format!("Support for {data_type} in JSON reader")))
854 }
855}
856
857#[cfg(test)]
858mod tests {
859 use arrow_array::cast::AsArray;
860 use arrow_array::{
861 Array, BooleanArray, Float64Array, GenericListViewArray, Int32Array, ListArray, MapArray,
862 NullArray, OffsetSizeTrait, StringArray, StringViewArray, StructArray,
863 };
864 use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer, ScalarBuffer};
865 use arrow_cast::display::{ArrayFormatter, FormatOptions};
866 use arrow_schema::{Field, Fields};
867 use serde_json::json;
868 use std::fs::File;
869 use std::io::{BufReader, Cursor, Seek};
870
871 use super::*;
872
873 fn do_read(
874 buf: &str,
875 batch_size: usize,
876 coerce_primitive: bool,
877 strict_mode: bool,
878 schema: SchemaRef,
879 ) -> Vec<RecordBatch> {
880 let mut unbuffered = vec![];
881
882 for batch_size in [1, 3, 100, batch_size] {
884 unbuffered = ReaderBuilder::new(schema.clone())
885 .with_batch_size(batch_size)
886 .with_coerce_primitive(coerce_primitive)
887 .build(Cursor::new(buf.as_bytes()))
888 .unwrap()
889 .collect::<Result<Vec<_>, _>>()
890 .unwrap();
891
892 for b in unbuffered.iter().take(unbuffered.len() - 1) {
893 assert_eq!(b.num_rows(), batch_size)
894 }
895
896 for b in [1, 3, 5] {
898 let buffered = ReaderBuilder::new(schema.clone())
899 .with_batch_size(batch_size)
900 .with_coerce_primitive(coerce_primitive)
901 .with_strict_mode(strict_mode)
902 .build(BufReader::with_capacity(b, Cursor::new(buf.as_bytes())))
903 .unwrap()
904 .collect::<Result<Vec<_>, _>>()
905 .unwrap();
906 assert_eq!(unbuffered, buffered);
907 }
908 }
909
910 unbuffered
911 }
912
913 #[test]
914 fn test_basic() {
915 let buf = r#"
916 {"a": 1, "b": 2, "c": true, "d": 1}
917 {"a": 2E0, "b": 4, "c": false, "d": 2, "e": 254}
918
919 {"b": 6, "a": 2.0, "d": 45}
920 {"b": "5", "a": 2}
921 {"b": 4e0}
922 {"b": 7, "a": null}
923 "#;
924
925 let schema = Arc::new(Schema::new(vec![
926 Field::new("a", DataType::Int64, true),
927 Field::new("b", DataType::Int32, true),
928 Field::new("c", DataType::Boolean, true),
929 Field::new("d", DataType::Date32, true),
930 Field::new("e", DataType::Date64, true),
931 ]));
932
933 let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder().unwrap();
934 assert!(decoder.is_empty());
935 assert_eq!(decoder.len(), 0);
936 assert!(!decoder.has_partial_record());
937 assert_eq!(decoder.decode(buf.as_bytes()).unwrap(), 221);
938 assert!(!decoder.is_empty());
939 assert_eq!(decoder.len(), 6);
940 assert!(!decoder.has_partial_record());
941 let batch = decoder.flush().unwrap().unwrap();
942 assert_eq!(batch.num_rows(), 6);
943 assert!(decoder.is_empty());
944 assert_eq!(decoder.len(), 0);
945 assert!(!decoder.has_partial_record());
946
947 let batches = do_read(buf, 1024, false, false, schema);
948 assert_eq!(batches.len(), 1);
949
950 let col1 = batches[0].column(0).as_primitive::<Int64Type>();
951 assert_eq!(col1.null_count(), 2);
952 assert_eq!(col1.values(), &[1, 2, 2, 2, 0, 0]);
953 assert!(col1.is_null(4));
954 assert!(col1.is_null(5));
955
956 let col2 = batches[0].column(1).as_primitive::<Int32Type>();
957 assert_eq!(col2.null_count(), 0);
958 assert_eq!(col2.values(), &[2, 4, 6, 5, 4, 7]);
959
960 let col3 = batches[0].column(2).as_boolean();
961 assert_eq!(col3.null_count(), 4);
962 assert!(col3.value(0));
963 assert!(!col3.is_null(0));
964 assert!(!col3.value(1));
965 assert!(!col3.is_null(1));
966
967 let col4 = batches[0].column(3).as_primitive::<Date32Type>();
968 assert_eq!(col4.null_count(), 3);
969 assert!(col4.is_null(3));
970 assert_eq!(col4.values(), &[1, 2, 45, 0, 0, 0]);
971
972 let col5 = batches[0].column(4).as_primitive::<Date64Type>();
973 assert_eq!(col5.null_count(), 5);
974 assert!(col5.is_null(0));
975 assert!(col5.is_null(2));
976 assert!(col5.is_null(3));
977 assert_eq!(col5.values(), &[0, 254, 0, 0, 0, 0]);
978 }
979
980 #[test]
981 fn test_string() {
982 let buf = r#"
983 {"a": "1", "b": "2"}
984 {"a": "hello", "b": "shoo"}
985 {"b": "\t😁foo", "a": "\nfoobar\ud83d\ude00\u0061\u0073\u0066\u0067\u00FF"}
986
987 {"b": null}
988 {"b": "", "a": null}
989
990 "#;
991 let schema = Arc::new(Schema::new(vec![
992 Field::new("a", DataType::Utf8, true),
993 Field::new("b", DataType::LargeUtf8, true),
994 ]));
995
996 let batches = do_read(buf, 1024, false, false, schema);
997 assert_eq!(batches.len(), 1);
998
999 let col1 = batches[0].column(0).as_string::<i32>();
1000 assert_eq!(col1.null_count(), 2);
1001 assert_eq!(col1.value(0), "1");
1002 assert_eq!(col1.value(1), "hello");
1003 assert_eq!(col1.value(2), "\nfoobar😀asfgÿ");
1004 assert!(col1.is_null(3));
1005 assert!(col1.is_null(4));
1006
1007 let col2 = batches[0].column(1).as_string::<i64>();
1008 assert_eq!(col2.null_count(), 1);
1009 assert_eq!(col2.value(0), "2");
1010 assert_eq!(col2.value(1), "shoo");
1011 assert_eq!(col2.value(2), "\t😁foo");
1012 assert!(col2.is_null(3));
1013 assert_eq!(col2.value(4), "");
1014 }
1015
1016 #[test]
1017 fn test_long_string_view_allocation() {
1018 let expected_capacity: usize = 41;
1028
1029 let buf = r#"
1030 {"a": "short", "b": "dummy"}
1031 {"a": "this is definitely long", "b": "dummy"}
1032 {"a": "hello", "b": "dummy"}
1033 {"a": "\nfoobar😀asfgÿ", "b": "dummy"}
1034 "#;
1035
1036 let schema = Arc::new(Schema::new(vec![
1037 Field::new("a", DataType::Utf8View, true),
1038 Field::new("b", DataType::LargeUtf8, true),
1039 ]));
1040
1041 let batches = do_read(buf, 1024, false, false, schema);
1042 assert_eq!(batches.len(), 1, "Expected one record batch");
1043
1044 let col_a = batches[0].column(0);
1046 let string_view_array = col_a
1047 .as_any()
1048 .downcast_ref::<StringViewArray>()
1049 .expect("Column should be a StringViewArray");
1050
1051 let data_buffer = string_view_array.to_data().buffers()[0].len();
1054
1055 assert!(
1058 data_buffer >= expected_capacity,
1059 "Data buffer length ({data_buffer}) should be at least {expected_capacity}",
1060 );
1061
1062 assert_eq!(string_view_array.value(0), "short");
1064 assert_eq!(string_view_array.value(1), "this is definitely long");
1065 assert_eq!(string_view_array.value(2), "hello");
1066 assert_eq!(string_view_array.value(3), "\nfoobar😀asfgÿ");
1067 }
1068
1069 #[test]
1071 fn test_numeric_view_allocation() {
1072 let expected_capacity: usize = 33;
1080
1081 let buf = r#"
1082 {"n": 123456789}
1083 {"n": 1000000000000}
1084 {"n": 3.1415}
1085 {"n": 2.718281828459045}
1086 "#;
1087
1088 let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Utf8View, true)]));
1089
1090 let batches = do_read(buf, 1024, true, false, schema);
1091 assert_eq!(batches.len(), 1, "Expected one record batch");
1092
1093 let col_n = batches[0].column(0);
1094 let string_view_array = col_n
1095 .as_any()
1096 .downcast_ref::<StringViewArray>()
1097 .expect("Column should be a StringViewArray");
1098
1099 let data_buffer = string_view_array.to_data().buffers()[0].len();
1101 assert!(
1102 data_buffer >= expected_capacity,
1103 "Data buffer length ({data_buffer}) should be at least {expected_capacity}",
1104 );
1105
1106 assert_eq!(string_view_array.value(0), "123456789");
1109 assert_eq!(string_view_array.value(1), "1000000000000");
1110 assert_eq!(string_view_array.value(2), "3.1415");
1111 assert_eq!(string_view_array.value(3), "2.718281828459045");
1112 }
1113
1114 #[test]
1115 fn test_string_with_uft8view() {
1116 let buf = r#"
1117 {"a": "1", "b": "2"}
1118 {"a": "hello", "b": "shoo"}
1119 {"b": "\t😁foo", "a": "\nfoobar\ud83d\ude00\u0061\u0073\u0066\u0067\u00FF"}
1120
1121 {"b": null}
1122 {"b": "", "a": null}
1123
1124 "#;
1125 let schema = Arc::new(Schema::new(vec![
1126 Field::new("a", DataType::Utf8View, true),
1127 Field::new("b", DataType::LargeUtf8, true),
1128 ]));
1129
1130 let batches = do_read(buf, 1024, false, false, schema);
1131 assert_eq!(batches.len(), 1);
1132
1133 let col1 = batches[0].column(0).as_string_view();
1134 assert_eq!(col1.null_count(), 2);
1135 assert_eq!(col1.value(0), "1");
1136 assert_eq!(col1.value(1), "hello");
1137 assert_eq!(col1.value(2), "\nfoobar😀asfgÿ");
1138 assert!(col1.is_null(3));
1139 assert!(col1.is_null(4));
1140 assert_eq!(col1.data_type(), &DataType::Utf8View);
1141
1142 let col2 = batches[0].column(1).as_string::<i64>();
1143 assert_eq!(col2.null_count(), 1);
1144 assert_eq!(col2.value(0), "2");
1145 assert_eq!(col2.value(1), "shoo");
1146 assert_eq!(col2.value(2), "\t😁foo");
1147 assert!(col2.is_null(3));
1148 assert_eq!(col2.value(4), "");
1149 }
1150
1151 #[test]
1152 fn test_complex() {
1153 let buf = r#"
1154 {"list": [], "nested": {"a": 1, "b": 2}, "nested_list": {"list2": [{"c": 3}, {"c": 4}]}}
1155 {"list": [5, 6], "nested": {"a": 7}, "nested_list": {"list2": []}}
1156 {"list": null, "nested": {"a": null}}
1157 "#;
1158
1159 let schema = Arc::new(Schema::new(vec![
1160 Field::new_list("list", Field::new("element", DataType::Int32, false), true),
1161 Field::new_struct(
1162 "nested",
1163 vec![
1164 Field::new("a", DataType::Int32, true),
1165 Field::new("b", DataType::Int32, true),
1166 ],
1167 true,
1168 ),
1169 Field::new_struct(
1170 "nested_list",
1171 vec![Field::new_list(
1172 "list2",
1173 Field::new_struct(
1174 "element",
1175 vec![Field::new("c", DataType::Int32, false)],
1176 false,
1177 ),
1178 true,
1179 )],
1180 true,
1181 ),
1182 ]));
1183
1184 let batches = do_read(buf, 1024, false, false, schema);
1185 assert_eq!(batches.len(), 1);
1186
1187 let list = batches[0].column(0).as_list::<i32>();
1188 assert_eq!(list.len(), 3);
1189 assert_eq!(list.value_offsets(), &[0, 0, 2, 2]);
1190 assert_eq!(list.null_count(), 1);
1191 assert!(list.is_null(2));
1192 let list_values = list.values().as_primitive::<Int32Type>();
1193 assert_eq!(list_values.values(), &[5, 6]);
1194
1195 let nested = batches[0].column(1).as_struct();
1196 let a = nested.column(0).as_primitive::<Int32Type>();
1197 assert_eq!(list.null_count(), 1);
1198 assert_eq!(a.values(), &[1, 7, 0]);
1199 assert!(list.is_null(2));
1200
1201 let b = nested.column(1).as_primitive::<Int32Type>();
1202 assert_eq!(b.null_count(), 2);
1203 assert_eq!(b.len(), 3);
1204 assert_eq!(b.value(0), 2);
1205 assert!(b.is_null(1));
1206 assert!(b.is_null(2));
1207
1208 let nested_list = batches[0].column(2).as_struct();
1209 assert_eq!(nested_list.len(), 3);
1210 assert_eq!(nested_list.null_count(), 1);
1211 assert!(nested_list.is_null(2));
1212
1213 let list2 = nested_list.column(0).as_list::<i32>();
1214 assert_eq!(list2.len(), 3);
1215 assert_eq!(list2.null_count(), 1);
1216 assert_eq!(list2.value_offsets(), &[0, 2, 2, 2]);
1217 assert!(list2.is_null(2));
1218
1219 let list2_values = list2.values().as_struct();
1220
1221 let c = list2_values.column(0).as_primitive::<Int32Type>();
1222 assert_eq!(c.values(), &[3, 4]);
1223 }
1224
1225 #[test]
1226 fn test_projection() {
1227 let buf = r#"
1228 {"list": [], "nested": {"a": 1, "b": 2}, "nested_list": {"list2": [{"c": 3, "d": 5}, {"c": 4}]}}
1229 {"list": [5, 6], "nested": {"a": 7}, "nested_list": {"list2": []}}
1230 "#;
1231
1232 let schema = Arc::new(Schema::new(vec![
1233 Field::new_struct(
1234 "nested",
1235 vec![Field::new("a", DataType::Int32, false)],
1236 true,
1237 ),
1238 Field::new_struct(
1239 "nested_list",
1240 vec![Field::new_list(
1241 "list2",
1242 Field::new_struct(
1243 "element",
1244 vec![Field::new("d", DataType::Int32, true)],
1245 false,
1246 ),
1247 true,
1248 )],
1249 true,
1250 ),
1251 ]));
1252
1253 let batches = do_read(buf, 1024, false, false, schema);
1254 assert_eq!(batches.len(), 1);
1255
1256 let nested = batches[0].column(0).as_struct();
1257 assert_eq!(nested.num_columns(), 1);
1258 let a = nested.column(0).as_primitive::<Int32Type>();
1259 assert_eq!(a.null_count(), 0);
1260 assert_eq!(a.values(), &[1, 7]);
1261
1262 let nested_list = batches[0].column(1).as_struct();
1263 assert_eq!(nested_list.num_columns(), 1);
1264 assert_eq!(nested_list.null_count(), 0);
1265
1266 let list2 = nested_list.column(0).as_list::<i32>();
1267 assert_eq!(list2.value_offsets(), &[0, 2, 2]);
1268 assert_eq!(list2.null_count(), 0);
1269
1270 let child = list2.values().as_struct();
1271 assert_eq!(child.num_columns(), 1);
1272 assert_eq!(child.len(), 2);
1273 assert_eq!(child.null_count(), 0);
1274
1275 let c = child.column(0).as_primitive::<Int32Type>();
1276 assert_eq!(c.values(), &[5, 0]);
1277 assert_eq!(c.null_count(), 1);
1278 assert!(c.is_null(1));
1279 }
1280
1281 #[test]
1282 fn test_map() {
1283 let buf = r#"
1284 {"map": {"a": ["foo", null]}}
1285 {"map": {"a": [null], "b": []}}
1286 {"map": {"c": null, "a": ["baz"]}}
1287 "#;
1288 let map = Field::new_map(
1289 "map",
1290 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1291 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1292 Field::new_list(
1293 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
1294 Field::new("element", DataType::Utf8, true),
1295 true,
1296 ),
1297 false,
1298 true,
1299 );
1300
1301 let schema = Arc::new(Schema::new(vec![map]));
1302
1303 let batches = do_read(buf, 1024, false, false, schema);
1304 assert_eq!(batches.len(), 1);
1305
1306 let map = batches[0].column(0).as_map();
1307 let map_keys = map.keys().as_string::<i32>();
1308 let map_values = map.values().as_list::<i32>();
1309 assert_eq!(map.value_offsets(), &[0, 1, 3, 5]);
1310
1311 let k: Vec<_> = map_keys.iter().flatten().collect();
1312 assert_eq!(&k, &["a", "a", "b", "c", "a"]);
1313
1314 let list_values = map_values.values().as_string::<i32>();
1315 let lv: Vec<_> = list_values.iter().collect();
1316 assert_eq!(&lv, &[Some("foo"), None, None, Some("baz")]);
1317 assert_eq!(map_values.value_offsets(), &[0, 2, 3, 3, 3, 4]);
1318 assert_eq!(map_values.null_count(), 1);
1319 assert!(map_values.is_null(3));
1320
1321 let options = FormatOptions::default().with_null("null");
1322 let formatter = ArrayFormatter::try_new(map, &options).unwrap();
1323 assert_eq!(formatter.value(0).to_string(), "{a: [foo, null]}");
1324 assert_eq!(formatter.value(1).to_string(), "{a: [null], b: []}");
1325 assert_eq!(formatter.value(2).to_string(), "{c: null, a: [baz]}");
1326 }
1327
1328 #[test]
1329 fn test_not_coercing_primitive_into_string_without_flag() {
1330 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
1331
1332 let buf = r#"{"a": 1}"#;
1333 let err = ReaderBuilder::new(schema.clone())
1334 .with_batch_size(1024)
1335 .build(Cursor::new(buf.as_bytes()))
1336 .unwrap()
1337 .read()
1338 .unwrap_err();
1339
1340 assert_eq!(
1341 err.to_string(),
1342 "Json error: whilst decoding field 'a': expected string got 1"
1343 );
1344
1345 let buf = r#"{"a": true}"#;
1346 let err = ReaderBuilder::new(schema)
1347 .with_batch_size(1024)
1348 .build(Cursor::new(buf.as_bytes()))
1349 .unwrap()
1350 .read()
1351 .unwrap_err();
1352
1353 assert_eq!(
1354 err.to_string(),
1355 "Json error: whilst decoding field 'a': expected string got true"
1356 );
1357 }
1358
1359 #[test]
1360 fn test_coercing_primitive_into_string() {
1361 let buf = r#"
1362 {"a": 1, "b": 2, "c": true}
1363 {"a": 2E0, "b": 4, "c": false}
1364
1365 {"b": 6, "a": 2.0}
1366 {"b": "5", "a": 2}
1367 {"b": 4e0}
1368 {"b": 7, "a": null}
1369 "#;
1370
1371 let schema = Arc::new(Schema::new(vec![
1372 Field::new("a", DataType::Utf8, true),
1373 Field::new("b", DataType::Utf8, true),
1374 Field::new("c", DataType::Utf8, true),
1375 ]));
1376
1377 let batches = do_read(buf, 1024, true, false, schema);
1378 assert_eq!(batches.len(), 1);
1379
1380 let col1 = batches[0].column(0).as_string::<i32>();
1381 assert_eq!(col1.null_count(), 2);
1382 assert_eq!(col1.value(0), "1");
1383 assert_eq!(col1.value(1), "2E0");
1384 assert_eq!(col1.value(2), "2.0");
1385 assert_eq!(col1.value(3), "2");
1386 assert!(col1.is_null(4));
1387 assert!(col1.is_null(5));
1388
1389 let col2 = batches[0].column(1).as_string::<i32>();
1390 assert_eq!(col2.null_count(), 0);
1391 assert_eq!(col2.value(0), "2");
1392 assert_eq!(col2.value(1), "4");
1393 assert_eq!(col2.value(2), "6");
1394 assert_eq!(col2.value(3), "5");
1395 assert_eq!(col2.value(4), "4e0");
1396 assert_eq!(col2.value(5), "7");
1397
1398 let col3 = batches[0].column(2).as_string::<i32>();
1399 assert_eq!(col3.null_count(), 4);
1400 assert_eq!(col3.value(0), "true");
1401 assert_eq!(col3.value(1), "false");
1402 assert!(col3.is_null(2));
1403 assert!(col3.is_null(3));
1404 assert!(col3.is_null(4));
1405 assert!(col3.is_null(5));
1406 }
1407
1408 fn test_decimal<T: DecimalType>(data_type: DataType) {
1409 let buf = r#"
1410 {"a": 1, "b": 2, "c": 38.30}
1411 {"a": 2, "b": 4, "c": 123.456}
1412
1413 {"b": 1337, "a": "2.0452"}
1414 {"b": "5", "a": "11034.2"}
1415 {"b": 40}
1416 {"b": 1234, "a": null}
1417 "#;
1418
1419 let schema = Arc::new(Schema::new(vec![
1420 Field::new("a", data_type.clone(), true),
1421 Field::new("b", data_type.clone(), true),
1422 Field::new("c", data_type, true),
1423 ]));
1424
1425 let batches = do_read(buf, 1024, true, false, schema);
1426 assert_eq!(batches.len(), 1);
1427
1428 let col1 = batches[0].column(0).as_primitive::<T>();
1429 assert_eq!(col1.null_count(), 2);
1430 assert!(col1.is_null(4));
1431 assert!(col1.is_null(5));
1432 assert_eq!(
1433 col1.values(),
1434 &[100, 200, 204, 1103420, 0, 0].map(T::Native::usize_as)
1435 );
1436
1437 let col2 = batches[0].column(1).as_primitive::<T>();
1438 assert_eq!(col2.null_count(), 0);
1439 assert_eq!(
1440 col2.values(),
1441 &[200, 400, 133700, 500, 4000, 123400].map(T::Native::usize_as)
1442 );
1443
1444 let col3 = batches[0].column(2).as_primitive::<T>();
1445 assert_eq!(col3.null_count(), 4);
1446 assert!(!col3.is_null(0));
1447 assert!(!col3.is_null(1));
1448 assert!(col3.is_null(2));
1449 assert!(col3.is_null(3));
1450 assert!(col3.is_null(4));
1451 assert!(col3.is_null(5));
1452 assert_eq!(
1453 col3.values(),
1454 &[3830, 12345, 0, 0, 0, 0].map(T::Native::usize_as)
1455 );
1456 }
1457
1458 #[test]
1459 fn test_decimals() {
1460 test_decimal::<Decimal32Type>(DataType::Decimal32(8, 2));
1461 test_decimal::<Decimal64Type>(DataType::Decimal64(10, 2));
1462 test_decimal::<Decimal128Type>(DataType::Decimal128(10, 2));
1463 test_decimal::<Decimal256Type>(DataType::Decimal256(10, 2));
1464 }
1465
1466 fn test_timestamp<T: ArrowTimestampType>() {
1467 let buf = r#"
1468 {"a": 1, "b": "2020-09-08T13:42:29.190855+00:00", "c": 38.30, "d": "1997-01-31T09:26:56.123"}
1469 {"a": 2, "b": "2020-09-08T13:42:29.190855Z", "c": 123.456, "d": 123.456}
1470
1471 {"b": 1337, "b": "2020-09-08T13:42:29Z", "c": "1997-01-31T09:26:56.123", "d": "1997-01-31T09:26:56.123Z"}
1472 {"b": 40, "c": "2020-09-08T13:42:29.190855+00:00", "d": "1997-01-31 09:26:56.123-05:00"}
1473 {"b": 1234, "a": null, "c": "1997-01-31 09:26:56.123Z", "d": "1997-01-31 092656"}
1474 {"c": "1997-01-31T14:26:56.123-05:00", "d": "1997-01-31"}
1475 "#;
1476
1477 let with_timezone = DataType::Timestamp(T::UNIT, Some("+08:00".into()));
1478 let schema = Arc::new(Schema::new(vec![
1479 Field::new("a", T::DATA_TYPE, true),
1480 Field::new("b", T::DATA_TYPE, true),
1481 Field::new("c", T::DATA_TYPE, true),
1482 Field::new("d", with_timezone, true),
1483 ]));
1484
1485 let batches = do_read(buf, 1024, true, false, schema);
1486 assert_eq!(batches.len(), 1);
1487
1488 let unit_in_nanos: i64 = match T::UNIT {
1489 TimeUnit::Second => 1_000_000_000,
1490 TimeUnit::Millisecond => 1_000_000,
1491 TimeUnit::Microsecond => 1_000,
1492 TimeUnit::Nanosecond => 1,
1493 };
1494
1495 let col1 = batches[0].column(0).as_primitive::<T>();
1496 assert_eq!(col1.null_count(), 4);
1497 assert!(col1.is_null(2));
1498 assert!(col1.is_null(3));
1499 assert!(col1.is_null(4));
1500 assert!(col1.is_null(5));
1501 assert_eq!(col1.values(), &[1, 2, 0, 0, 0, 0].map(T::Native::usize_as));
1502
1503 let col2 = batches[0].column(1).as_primitive::<T>();
1504 assert_eq!(col2.null_count(), 1);
1505 assert!(col2.is_null(5));
1506 assert_eq!(
1507 col2.values(),
1508 &[
1509 1599572549190855000 / unit_in_nanos,
1510 1599572549190855000 / unit_in_nanos,
1511 1599572549000000000 / unit_in_nanos,
1512 40,
1513 1234,
1514 0
1515 ]
1516 );
1517
1518 let col3 = batches[0].column(2).as_primitive::<T>();
1519 assert_eq!(col3.null_count(), 0);
1520 assert_eq!(
1521 col3.values(),
1522 &[
1523 38,
1524 123,
1525 854702816123000000 / unit_in_nanos,
1526 1599572549190855000 / unit_in_nanos,
1527 854702816123000000 / unit_in_nanos,
1528 854738816123000000 / unit_in_nanos
1529 ]
1530 );
1531
1532 let col4 = batches[0].column(3).as_primitive::<T>();
1533
1534 assert_eq!(col4.null_count(), 0);
1535 assert_eq!(
1536 col4.values(),
1537 &[
1538 854674016123000000 / unit_in_nanos,
1539 123,
1540 854702816123000000 / unit_in_nanos,
1541 854720816123000000 / unit_in_nanos,
1542 854674016000000000 / unit_in_nanos,
1543 854640000000000000 / unit_in_nanos
1544 ]
1545 );
1546 }
1547
1548 #[test]
1549 fn test_timestamps() {
1550 test_timestamp::<TimestampSecondType>();
1551 test_timestamp::<TimestampMillisecondType>();
1552 test_timestamp::<TimestampMicrosecondType>();
1553 test_timestamp::<TimestampNanosecondType>();
1554 }
1555
1556 fn test_time<T: ArrowTemporalType>() {
1557 let buf = r#"
1558 {"a": 1, "b": "09:26:56.123 AM", "c": 38.30}
1559 {"a": 2, "b": "23:59:59", "c": 123.456}
1560
1561 {"b": 1337, "b": "6:00 pm", "c": "09:26:56.123"}
1562 {"b": 40, "c": "13:42:29.190855"}
1563 {"b": 1234, "a": null, "c": "09:26:56.123"}
1564 {"c": "14:26:56.123"}
1565 "#;
1566
1567 let unit = match T::DATA_TYPE {
1568 DataType::Time32(unit) | DataType::Time64(unit) => unit,
1569 _ => unreachable!(),
1570 };
1571
1572 let unit_in_nanos = match unit {
1573 TimeUnit::Second => 1_000_000_000,
1574 TimeUnit::Millisecond => 1_000_000,
1575 TimeUnit::Microsecond => 1_000,
1576 TimeUnit::Nanosecond => 1,
1577 };
1578
1579 let schema = Arc::new(Schema::new(vec![
1580 Field::new("a", T::DATA_TYPE, true),
1581 Field::new("b", T::DATA_TYPE, true),
1582 Field::new("c", T::DATA_TYPE, true),
1583 ]));
1584
1585 let batches = do_read(buf, 1024, true, false, schema);
1586 assert_eq!(batches.len(), 1);
1587
1588 let col1 = batches[0].column(0).as_primitive::<T>();
1589 assert_eq!(col1.null_count(), 4);
1590 assert!(col1.is_null(2));
1591 assert!(col1.is_null(3));
1592 assert!(col1.is_null(4));
1593 assert!(col1.is_null(5));
1594 assert_eq!(col1.values(), &[1, 2, 0, 0, 0, 0].map(T::Native::usize_as));
1595
1596 let col2 = batches[0].column(1).as_primitive::<T>();
1597 assert_eq!(col2.null_count(), 1);
1598 assert!(col2.is_null(5));
1599 assert_eq!(
1600 col2.values(),
1601 &[
1602 34016123000000 / unit_in_nanos,
1603 86399000000000 / unit_in_nanos,
1604 64800000000000 / unit_in_nanos,
1605 40,
1606 1234,
1607 0
1608 ]
1609 .map(T::Native::usize_as)
1610 );
1611
1612 let col3 = batches[0].column(2).as_primitive::<T>();
1613 assert_eq!(col3.null_count(), 0);
1614 assert_eq!(
1615 col3.values(),
1616 &[
1617 38,
1618 123,
1619 34016123000000 / unit_in_nanos,
1620 49349190855000 / unit_in_nanos,
1621 34016123000000 / unit_in_nanos,
1622 52016123000000 / unit_in_nanos
1623 ]
1624 .map(T::Native::usize_as)
1625 );
1626 }
1627
1628 #[test]
1629 fn test_times() {
1630 test_time::<Time32MillisecondType>();
1631 test_time::<Time32SecondType>();
1632 test_time::<Time64MicrosecondType>();
1633 test_time::<Time64NanosecondType>();
1634 }
1635
1636 fn test_duration<T: ArrowTemporalType>() {
1637 let buf = r#"
1638 {"a": 1, "b": "2"}
1639 {"a": 3, "b": null}
1640 "#;
1641
1642 let schema = Arc::new(Schema::new(vec![
1643 Field::new("a", T::DATA_TYPE, true),
1644 Field::new("b", T::DATA_TYPE, true),
1645 ]));
1646
1647 let batches = do_read(buf, 1024, true, false, schema);
1648 assert_eq!(batches.len(), 1);
1649
1650 let col_a = batches[0].column_by_name("a").unwrap().as_primitive::<T>();
1651 assert_eq!(col_a.null_count(), 0);
1652 assert_eq!(col_a.values(), &[1, 3].map(T::Native::usize_as));
1653
1654 let col2 = batches[0].column_by_name("b").unwrap().as_primitive::<T>();
1655 assert_eq!(col2.null_count(), 1);
1656 assert_eq!(col2.values(), &[2, 0].map(T::Native::usize_as));
1657 }
1658
1659 #[test]
1660 fn test_durations() {
1661 test_duration::<DurationNanosecondType>();
1662 test_duration::<DurationMicrosecondType>();
1663 test_duration::<DurationMillisecondType>();
1664 test_duration::<DurationSecondType>();
1665 }
1666
1667 #[test]
1668 fn test_delta_checkpoint() {
1669 let json = "{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}";
1670 let schema = Arc::new(Schema::new(vec![
1671 Field::new_struct(
1672 "protocol",
1673 vec![
1674 Field::new("minReaderVersion", DataType::Int32, true),
1675 Field::new("minWriterVersion", DataType::Int32, true),
1676 ],
1677 true,
1678 ),
1679 Field::new_struct(
1680 "add",
1681 vec![Field::new_map(
1682 "partitionValues",
1683 "key_value",
1684 Field::new("key", DataType::Utf8, false),
1685 Field::new("value", DataType::Utf8, true),
1686 false,
1687 false,
1688 )],
1689 true,
1690 ),
1691 ]));
1692
1693 let batches = do_read(json, 1024, true, false, schema);
1694 assert_eq!(batches.len(), 1);
1695
1696 let s: StructArray = batches.into_iter().next().unwrap().into();
1697 let opts = FormatOptions::default().with_null("null");
1698 let formatter = ArrayFormatter::try_new(&s, &opts).unwrap();
1699 assert_eq!(
1700 formatter.value(0).to_string(),
1701 "{protocol: {minReaderVersion: 1, minWriterVersion: 2}, add: null}"
1702 );
1703 }
1704
1705 #[test]
1706 fn struct_nullability() {
1707 let do_test = |child: DataType| {
1708 let non_null = r#"{"foo": {}}"#;
1710 let schema = Arc::new(Schema::new(vec![Field::new_struct(
1711 "foo",
1712 vec![Field::new("bar", child, false)],
1713 true,
1714 )]));
1715 let mut reader = ReaderBuilder::new(schema.clone())
1716 .build(Cursor::new(non_null.as_bytes()))
1717 .unwrap();
1718 assert!(reader.next().unwrap().is_err()); let null = r#"{"foo": {bar: null}}"#;
1721 let mut reader = ReaderBuilder::new(schema.clone())
1722 .build(Cursor::new(null.as_bytes()))
1723 .unwrap();
1724 assert!(reader.next().unwrap().is_err()); let null = r#"{"foo": null}"#;
1728 let mut reader = ReaderBuilder::new(schema)
1729 .build(Cursor::new(null.as_bytes()))
1730 .unwrap();
1731 let batch = reader.next().unwrap().unwrap();
1732 assert_eq!(batch.num_columns(), 1);
1733 let foo = batch.column(0).as_struct();
1734 assert_eq!(foo.len(), 1);
1735 assert!(foo.is_null(0));
1736 assert_eq!(foo.num_columns(), 1);
1737
1738 let bar = foo.column(0);
1739 assert_eq!(bar.len(), 1);
1740 assert!(bar.is_null(0));
1742 };
1743
1744 do_test(DataType::Boolean);
1745 do_test(DataType::Int32);
1746 do_test(DataType::Utf8);
1747 do_test(DataType::Decimal128(2, 1));
1748 do_test(DataType::Timestamp(
1749 TimeUnit::Microsecond,
1750 Some("+00:00".into()),
1751 ));
1752 }
1753
1754 #[test]
1755 fn test_truncation() {
1756 let buf = r#"
1757 {"i64": 9223372036854775807, "u64": 18446744073709551615 }
1758 {"i64": "9223372036854775807", "u64": "18446744073709551615" }
1759 {"i64": -9223372036854775808, "u64": 0 }
1760 {"i64": "-9223372036854775808", "u64": 0 }
1761 "#;
1762
1763 let schema = Arc::new(Schema::new(vec![
1764 Field::new("i64", DataType::Int64, true),
1765 Field::new("u64", DataType::UInt64, true),
1766 ]));
1767
1768 let batches = do_read(buf, 1024, true, false, schema);
1769 assert_eq!(batches.len(), 1);
1770
1771 let i64 = batches[0].column(0).as_primitive::<Int64Type>();
1772 assert_eq!(i64.values(), &[i64::MAX, i64::MAX, i64::MIN, i64::MIN]);
1773
1774 let u64 = batches[0].column(1).as_primitive::<UInt64Type>();
1775 assert_eq!(u64.values(), &[u64::MAX, u64::MAX, u64::MIN, u64::MIN]);
1776 }
1777
1778 #[test]
1779 fn test_timestamp_truncation() {
1780 let buf = r#"
1781 {"time": 9223372036854775807 }
1782 {"time": -9223372036854775808 }
1783 {"time": 9e5 }
1784 "#;
1785
1786 let schema = Arc::new(Schema::new(vec![Field::new(
1787 "time",
1788 DataType::Timestamp(TimeUnit::Nanosecond, None),
1789 true,
1790 )]));
1791
1792 let batches = do_read(buf, 1024, true, false, schema);
1793 assert_eq!(batches.len(), 1);
1794
1795 let i64 = batches[0]
1796 .column(0)
1797 .as_primitive::<TimestampNanosecondType>();
1798 assert_eq!(i64.values(), &[i64::MAX, i64::MIN, 900000]);
1799 }
1800
1801 #[test]
1802 fn test_strict_mode_no_missing_columns_in_schema() {
1803 let buf = r#"
1804 {"a": 1, "b": "2", "c": true}
1805 {"a": 2E0, "b": "4", "c": false}
1806 "#;
1807
1808 let schema = Arc::new(Schema::new(vec![
1809 Field::new("a", DataType::Int16, false),
1810 Field::new("b", DataType::Utf8, false),
1811 Field::new("c", DataType::Boolean, false),
1812 ]));
1813
1814 let batches = do_read(buf, 1024, true, true, schema);
1815 assert_eq!(batches.len(), 1);
1816
1817 let buf = r#"
1818 {"a": 1, "b": "2", "c": {"a": true, "b": 1}}
1819 {"a": 2E0, "b": "4", "c": {"a": false, "b": 2}}
1820 "#;
1821
1822 let schema = Arc::new(Schema::new(vec![
1823 Field::new("a", DataType::Int16, false),
1824 Field::new("b", DataType::Utf8, false),
1825 Field::new_struct(
1826 "c",
1827 vec![
1828 Field::new("a", DataType::Boolean, false),
1829 Field::new("b", DataType::Int16, false),
1830 ],
1831 false,
1832 ),
1833 ]));
1834
1835 let batches = do_read(buf, 1024, true, true, schema);
1836 assert_eq!(batches.len(), 1);
1837 }
1838
1839 #[test]
1840 fn test_strict_mode_missing_columns_in_schema() {
1841 let buf = r#"
1842 {"a": 1, "b": "2", "c": true}
1843 {"a": 2E0, "b": "4", "c": false}
1844 "#;
1845
1846 let schema = Arc::new(Schema::new(vec![
1847 Field::new("a", DataType::Int16, true),
1848 Field::new("c", DataType::Boolean, true),
1849 ]));
1850
1851 let err = ReaderBuilder::new(schema)
1852 .with_batch_size(1024)
1853 .with_strict_mode(true)
1854 .build(Cursor::new(buf.as_bytes()))
1855 .unwrap()
1856 .read()
1857 .unwrap_err();
1858
1859 assert_eq!(
1860 err.to_string(),
1861 "Json error: column 'b' missing from schema"
1862 );
1863
1864 let buf = r#"
1865 {"a": 1, "b": "2", "c": {"a": true, "b": 1}}
1866 {"a": 2E0, "b": "4", "c": {"a": false, "b": 2}}
1867 "#;
1868
1869 let schema = Arc::new(Schema::new(vec![
1870 Field::new("a", DataType::Int16, false),
1871 Field::new("b", DataType::Utf8, false),
1872 Field::new_struct("c", vec![Field::new("a", DataType::Boolean, false)], false),
1873 ]));
1874
1875 let err = ReaderBuilder::new(schema)
1876 .with_batch_size(1024)
1877 .with_strict_mode(true)
1878 .build(Cursor::new(buf.as_bytes()))
1879 .unwrap()
1880 .read()
1881 .unwrap_err();
1882
1883 assert_eq!(
1884 err.to_string(),
1885 "Json error: whilst decoding field 'c': column 'b' missing from schema"
1886 );
1887 }
1888
1889 fn read_file(path: &str, schema: Option<Schema>) -> Reader<BufReader<File>> {
1890 let file = File::open(path).unwrap();
1891 let mut reader = BufReader::new(file);
1892 let schema = schema.unwrap_or_else(|| {
1893 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
1894 reader.rewind().unwrap();
1895 schema
1896 });
1897 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(64);
1898 builder.build(reader).unwrap()
1899 }
1900
1901 #[test]
1902 fn test_json_basic() {
1903 let mut reader = read_file("test/data/basic.json", None);
1904 let batch = reader.next().unwrap().unwrap();
1905
1906 assert_eq!(8, batch.num_columns());
1907 assert_eq!(12, batch.num_rows());
1908
1909 let schema = reader.schema();
1910 let batch_schema = batch.schema();
1911 assert_eq!(schema, batch_schema);
1912
1913 let a = schema.column_with_name("a").unwrap();
1914 assert_eq!(0, a.0);
1915 assert_eq!(&DataType::Int64, a.1.data_type());
1916 let b = schema.column_with_name("b").unwrap();
1917 assert_eq!(1, b.0);
1918 assert_eq!(&DataType::Float64, b.1.data_type());
1919 let c = schema.column_with_name("c").unwrap();
1920 assert_eq!(2, c.0);
1921 assert_eq!(&DataType::Boolean, c.1.data_type());
1922 let d = schema.column_with_name("d").unwrap();
1923 assert_eq!(3, d.0);
1924 assert_eq!(&DataType::Utf8, d.1.data_type());
1925
1926 let aa = batch.column(a.0).as_primitive::<Int64Type>();
1927 assert_eq!(1, aa.value(0));
1928 assert_eq!(-10, aa.value(1));
1929 let bb = batch.column(b.0).as_primitive::<Float64Type>();
1930 assert_eq!(2.0, bb.value(0));
1931 assert_eq!(-3.5, bb.value(1));
1932 let cc = batch.column(c.0).as_boolean();
1933 assert!(!cc.value(0));
1934 assert!(cc.value(10));
1935 let dd = batch.column(d.0).as_string::<i32>();
1936 assert_eq!("4", dd.value(0));
1937 assert_eq!("text", dd.value(8));
1938 }
1939
1940 #[test]
1941 fn test_json_empty_projection() {
1942 let mut reader = read_file("test/data/basic.json", Some(Schema::empty()));
1943 let batch = reader.next().unwrap().unwrap();
1944
1945 assert_eq!(0, batch.num_columns());
1946 assert_eq!(12, batch.num_rows());
1947 }
1948
1949 #[test]
1950 fn test_json_basic_with_nulls() {
1951 let mut reader = read_file("test/data/basic_nulls.json", None);
1952 let batch = reader.next().unwrap().unwrap();
1953
1954 assert_eq!(4, batch.num_columns());
1955 assert_eq!(12, batch.num_rows());
1956
1957 let schema = reader.schema();
1958 let batch_schema = batch.schema();
1959 assert_eq!(schema, batch_schema);
1960
1961 let a = schema.column_with_name("a").unwrap();
1962 assert_eq!(&DataType::Int64, a.1.data_type());
1963 let b = schema.column_with_name("b").unwrap();
1964 assert_eq!(&DataType::Float64, b.1.data_type());
1965 let c = schema.column_with_name("c").unwrap();
1966 assert_eq!(&DataType::Boolean, c.1.data_type());
1967 let d = schema.column_with_name("d").unwrap();
1968 assert_eq!(&DataType::Utf8, d.1.data_type());
1969
1970 let aa = batch.column(a.0).as_primitive::<Int64Type>();
1971 assert!(aa.is_valid(0));
1972 assert!(!aa.is_valid(1));
1973 assert!(!aa.is_valid(11));
1974 let bb = batch.column(b.0).as_primitive::<Float64Type>();
1975 assert!(bb.is_valid(0));
1976 assert!(!bb.is_valid(2));
1977 assert!(!bb.is_valid(11));
1978 let cc = batch.column(c.0).as_boolean();
1979 assert!(cc.is_valid(0));
1980 assert!(!cc.is_valid(4));
1981 assert!(!cc.is_valid(11));
1982 let dd = batch.column(d.0).as_string::<i32>();
1983 assert!(!dd.is_valid(0));
1984 assert!(dd.is_valid(1));
1985 assert!(!dd.is_valid(4));
1986 assert!(!dd.is_valid(11));
1987 }
1988
1989 #[test]
1990 fn test_json_basic_schema() {
1991 let schema = Schema::new(vec![
1992 Field::new("a", DataType::Int64, true),
1993 Field::new("b", DataType::Float32, false),
1994 Field::new("c", DataType::Boolean, false),
1995 Field::new("d", DataType::Utf8, false),
1996 ]);
1997
1998 let mut reader = read_file("test/data/basic.json", Some(schema.clone()));
1999 let reader_schema = reader.schema();
2000 assert_eq!(reader_schema.as_ref(), &schema);
2001 let batch = reader.next().unwrap().unwrap();
2002
2003 assert_eq!(4, batch.num_columns());
2004 assert_eq!(12, batch.num_rows());
2005
2006 let schema = batch.schema();
2007
2008 let a = schema.column_with_name("a").unwrap();
2009 assert_eq!(&DataType::Int64, a.1.data_type());
2010 let b = schema.column_with_name("b").unwrap();
2011 assert_eq!(&DataType::Float32, b.1.data_type());
2012 let c = schema.column_with_name("c").unwrap();
2013 assert_eq!(&DataType::Boolean, c.1.data_type());
2014 let d = schema.column_with_name("d").unwrap();
2015 assert_eq!(&DataType::Utf8, d.1.data_type());
2016
2017 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2018 assert_eq!(1, aa.value(0));
2019 assert_eq!(100000000000000, aa.value(11));
2020 let bb = batch.column(b.0).as_primitive::<Float32Type>();
2021 assert_eq!(2.0, bb.value(0));
2022 assert_eq!(-3.5, bb.value(1));
2023 }
2024
2025 #[test]
2026 fn test_json_basic_schema_projection() {
2027 let schema = Schema::new(vec![
2028 Field::new("a", DataType::Int64, true),
2029 Field::new("c", DataType::Boolean, false),
2030 ]);
2031
2032 let mut reader = read_file("test/data/basic.json", Some(schema.clone()));
2033 let batch = reader.next().unwrap().unwrap();
2034
2035 assert_eq!(2, batch.num_columns());
2036 assert_eq!(2, batch.schema().fields().len());
2037 assert_eq!(12, batch.num_rows());
2038
2039 assert_eq!(batch.schema().as_ref(), &schema);
2040
2041 let a = schema.column_with_name("a").unwrap();
2042 assert_eq!(0, a.0);
2043 assert_eq!(&DataType::Int64, a.1.data_type());
2044 let c = schema.column_with_name("c").unwrap();
2045 assert_eq!(1, c.0);
2046 assert_eq!(&DataType::Boolean, c.1.data_type());
2047 }
2048
2049 #[test]
2050 fn test_json_arrays() {
2051 let mut reader = read_file("test/data/arrays.json", None);
2052 let batch = reader.next().unwrap().unwrap();
2053
2054 assert_eq!(4, batch.num_columns());
2055 assert_eq!(3, batch.num_rows());
2056
2057 let schema = batch.schema();
2058
2059 let a = schema.column_with_name("a").unwrap();
2060 assert_eq!(&DataType::Int64, a.1.data_type());
2061 let b = schema.column_with_name("b").unwrap();
2062 assert_eq!(
2063 &DataType::List(Arc::new(Field::new_list_field(DataType::Float64, true))),
2064 b.1.data_type()
2065 );
2066 let c = schema.column_with_name("c").unwrap();
2067 assert_eq!(
2068 &DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))),
2069 c.1.data_type()
2070 );
2071 let d = schema.column_with_name("d").unwrap();
2072 assert_eq!(&DataType::Utf8, d.1.data_type());
2073
2074 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2075 assert_eq!(1, aa.value(0));
2076 assert_eq!(-10, aa.value(1));
2077 assert_eq!(1627668684594000000, aa.value(2));
2078 let bb = batch.column(b.0).as_list::<i32>();
2079 let bb = bb.values().as_primitive::<Float64Type>();
2080 assert_eq!(9, bb.len());
2081 assert_eq!(2.0, bb.value(0));
2082 assert_eq!(-6.1, bb.value(5));
2083 assert!(!bb.is_valid(7));
2084
2085 let cc = batch
2086 .column(c.0)
2087 .as_any()
2088 .downcast_ref::<ListArray>()
2089 .unwrap();
2090 let cc = cc.values().as_boolean();
2091 assert_eq!(6, cc.len());
2092 assert!(!cc.value(0));
2093 assert!(!cc.value(4));
2094 assert!(!cc.is_valid(5));
2095 }
2096
2097 #[test]
2098 fn test_empty_json_arrays() {
2099 let json_content = r#"
2100 {"items": []}
2101 {"items": null}
2102 {}
2103 "#;
2104
2105 let schema = Arc::new(Schema::new(vec![Field::new(
2106 "items",
2107 DataType::List(FieldRef::new(Field::new_list_field(DataType::Null, true))),
2108 true,
2109 )]));
2110
2111 let batches = do_read(json_content, 1024, false, false, schema);
2112 assert_eq!(batches.len(), 1);
2113
2114 let col1 = batches[0].column(0).as_list::<i32>();
2115 assert_eq!(col1.null_count(), 2);
2116 assert!(col1.value(0).is_empty());
2117 assert_eq!(col1.value(0).data_type(), &DataType::Null);
2118 assert!(col1.is_null(1));
2119 assert!(col1.is_null(2));
2120 }
2121
2122 #[test]
2123 fn test_nested_empty_json_arrays() {
2124 let json_content = r#"
2125 {"items": [[],[]]}
2126 {"items": [[null, null],[null]]}
2127 "#;
2128
2129 let schema = Arc::new(Schema::new(vec![Field::new(
2130 "items",
2131 DataType::List(FieldRef::new(Field::new_list_field(
2132 DataType::List(FieldRef::new(Field::new_list_field(DataType::Null, true))),
2133 true,
2134 ))),
2135 true,
2136 )]));
2137
2138 let batches = do_read(json_content, 1024, false, false, schema);
2139 assert_eq!(batches.len(), 1);
2140
2141 let col1 = batches[0].column(0).as_list::<i32>();
2142 assert_eq!(col1.null_count(), 0);
2143 assert_eq!(col1.value(0).len(), 2);
2144 assert!(col1.value(0).as_list::<i32>().value(0).is_empty());
2145 assert!(col1.value(0).as_list::<i32>().value(1).is_empty());
2146
2147 assert_eq!(col1.value(1).len(), 2);
2148 assert_eq!(col1.value(1).as_list::<i32>().value(0).len(), 2);
2149 assert_eq!(col1.value(1).as_list::<i32>().value(1).len(), 1);
2150 }
2151
2152 #[test]
2153 fn test_nested_list_json_arrays() {
2154 let c_field = Field::new_struct("c", vec![Field::new("d", DataType::Utf8, true)], true);
2155 let a_struct_field = Field::new_struct(
2156 "a",
2157 vec![Field::new("b", DataType::Boolean, true), c_field.clone()],
2158 true,
2159 );
2160 let a_field = Field::new("a", DataType::List(Arc::new(a_struct_field.clone())), true);
2161 let schema = Arc::new(Schema::new(vec![a_field.clone()]));
2162 let builder = ReaderBuilder::new(schema).with_batch_size(64);
2163 let json_content = r#"
2164 {"a": [{"b": true, "c": {"d": "a_text"}}, {"b": false, "c": {"d": "b_text"}}]}
2165 {"a": [{"b": false, "c": null}]}
2166 {"a": [{"b": true, "c": {"d": "c_text"}}, {"b": null, "c": {"d": "d_text"}}, {"b": true, "c": {"d": null}}]}
2167 {"a": null}
2168 {"a": []}
2169 {"a": [null]}
2170 "#;
2171 let mut reader = builder.build(Cursor::new(json_content)).unwrap();
2172
2173 let d = StringArray::from(vec![
2175 Some("a_text"),
2176 Some("b_text"),
2177 None,
2178 Some("c_text"),
2179 Some("d_text"),
2180 None,
2181 None,
2182 ]);
2183 let c = StructArray::new(
2184 vec![Field::new("d", DataType::Utf8, true)].into(),
2185 vec![Arc::new(d.clone()) as ArrayRef],
2186 Some(NullBuffer::from(vec![
2187 true, true, false, true, true, true, false,
2188 ])),
2189 );
2190 let b = BooleanArray::from(vec![
2191 Some(true),
2192 Some(false),
2193 Some(false),
2194 Some(true),
2195 None,
2196 Some(true),
2197 None,
2198 ]);
2199 let a = StructArray::new(
2200 vec![Field::new("b", DataType::Boolean, true), c_field.clone()].into(),
2201 vec![
2202 Arc::new(b.clone()) as ArrayRef,
2203 Arc::new(c.clone()) as ArrayRef,
2204 ],
2205 Some(NullBuffer::from(vec![
2206 true, true, true, true, true, true, false,
2207 ])),
2208 );
2209 let a_list = ListArray::new(
2210 Arc::new(a_struct_field.clone()),
2211 OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 2, 3, 6, 6, 6, 7])),
2212 Arc::new(a),
2213 Some(NullBuffer::from(vec![true, true, true, false, true, true])),
2214 );
2215
2216 let batch = reader.next().unwrap().unwrap();
2218 let read = batch.column(0);
2219 assert_eq!(read.len(), 6);
2220 let read: &ListArray = read.as_list::<i32>();
2222 let expected = &a_list;
2223 assert_eq!(read.value_offsets(), &[0, 2, 3, 6, 6, 6, 7]);
2224 assert_eq!(read.nulls(), expected.nulls());
2226 let struct_array = read.values().as_struct();
2228 let expected_struct_array = expected.values().as_struct();
2229
2230 assert_eq!(7, struct_array.len());
2231 assert_eq!(1, struct_array.null_count());
2232 assert_eq!(7, expected_struct_array.len());
2233 assert_eq!(1, expected_struct_array.null_count());
2234 assert_eq!(struct_array.nulls(), expected_struct_array.nulls());
2236 let read_b = struct_array.column(0);
2238 assert_eq!(read_b.as_ref(), &b);
2239 let read_c = struct_array.column(1);
2240 assert_eq!(read_c.as_struct(), &c);
2241 let read_c = read_c.as_struct();
2242 let read_d = read_c.column(0);
2243 assert_eq!(read_d.as_ref(), &d);
2244
2245 assert_eq!(read, expected);
2246 }
2247
2248 fn assert_read_list_view<O: OffsetSizeTrait>() {
2249 let field = Arc::new(Field::new("item", DataType::Int32, true));
2250 let data_type = GenericListViewArray::<O>::DATA_TYPE_CONSTRUCTOR(field.clone());
2251 let schema = Arc::new(Schema::new(vec![Field::new("lv", data_type, true)]));
2252
2253 let buf = r#"
2254 {"lv": [1, 2, 3]}
2255 {"lv": [4, null]}
2256 {"lv": null}
2257 {"lv": [6]}
2258 {"lv": []}
2259 "#;
2260
2261 let batches = do_read(buf, 1024, false, false, schema);
2262 assert_eq!(batches.len(), 1);
2263 let batch = &batches[0];
2264 let col = batch.column(0);
2265 let list_view = col
2266 .as_any()
2267 .downcast_ref::<GenericListViewArray<O>>()
2268 .unwrap();
2269
2270 assert_eq!(list_view.len(), 5);
2271
2272 let expected_offsets: Vec<O> = vec![0, 3, 5, 5, 6]
2274 .into_iter()
2275 .map(|v| O::usize_as(v))
2276 .collect();
2277 let expected_sizes: Vec<O> = vec![3, 2, 0, 1, 0]
2278 .into_iter()
2279 .map(|v| O::usize_as(v))
2280 .collect();
2281 assert_eq!(list_view.value_offsets(), &expected_offsets);
2282 assert_eq!(list_view.value_sizes(), &expected_sizes);
2283
2284 assert!(list_view.is_valid(0));
2286 let vals = list_view.value(0);
2287 let ints = vals.as_primitive::<Int32Type>();
2288 assert_eq!(ints.values(), &[1, 2, 3]);
2289
2290 assert!(list_view.is_valid(1));
2292 let vals = list_view.value(1);
2293 let ints = vals.as_primitive::<Int32Type>();
2294 assert_eq!(ints.len(), 2);
2295 assert_eq!(ints.value(0), 4);
2296 assert!(ints.is_null(1));
2297
2298 assert!(list_view.is_null(2));
2300
2301 assert!(list_view.is_valid(3));
2303 let vals = list_view.value(3);
2304 let ints = vals.as_primitive::<Int32Type>();
2305 assert_eq!(ints.values(), &[6]);
2306
2307 assert!(list_view.is_valid(4));
2309 let vals = list_view.value(4);
2310 assert_eq!(vals.len(), 0);
2311 }
2312
2313 #[test]
2314 fn test_read_list_view() {
2315 assert_read_list_view::<i32>();
2316 assert_read_list_view::<i64>();
2317 }
2318
2319 #[test]
2320 fn test_fixed_size_list() {
2321 let buf = r#"
2322 {"a": [1, 2, 3]}
2323 {"a": [4, 5, 6]}
2324 {"a": [7, 8, 9]}
2325 "#;
2326
2327 let field = Field::new_list_field(DataType::Int32, true);
2328 let schema = Arc::new(Schema::new(vec![Field::new(
2329 "a",
2330 DataType::FixedSizeList(Arc::new(field), 3),
2331 false,
2332 )]));
2333
2334 let batches = do_read(buf, 1024, false, false, schema);
2335 assert_eq!(batches.len(), 1);
2336
2337 let col = batches[0].column(0).as_fixed_size_list();
2338 assert_eq!(col.len(), 3);
2339 assert_eq!(col.value_length(), 3);
2340
2341 let values = col.values().as_primitive::<Int32Type>();
2342 assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
2343 }
2344
2345 #[test]
2346 fn test_fixed_size_list_nullable() {
2347 let buf = r#"
2348 {"a": [1, 2]}
2349 {"a": null}
2350 {"a": [3, null]}
2351 "#;
2352
2353 let field = Field::new_list_field(DataType::Int32, true);
2354 let schema = Arc::new(Schema::new(vec![Field::new(
2355 "a",
2356 DataType::FixedSizeList(Arc::new(field), 2),
2357 true,
2358 )]));
2359
2360 let batches = do_read(buf, 1024, false, false, schema);
2361 assert_eq!(batches.len(), 1);
2362
2363 let col = batches[0].column(0).as_fixed_size_list();
2364 assert_eq!(col.len(), 3);
2365 assert!(col.is_valid(0));
2366 assert!(col.is_null(1));
2367 assert!(col.is_valid(2));
2368
2369 let values = col.values().as_primitive::<Int32Type>();
2370 assert_eq!(values.value(0), 1);
2371 assert_eq!(values.value(1), 2);
2372 assert_eq!(values.value(4), 3);
2373 assert!(values.is_null(5));
2374 }
2375
2376 #[test]
2377 fn test_fixed_size_list_zero_size_non_nullable() {
2378 let buf = r#"
2379 {"a": []}
2380 {"a": []}
2381 {"a": []}
2382 "#;
2383
2384 let field = Field::new_list_field(DataType::Int32, true);
2385 let schema = Arc::new(Schema::new(vec![Field::new(
2386 "a",
2387 DataType::FixedSizeList(Arc::new(field), 0),
2388 false,
2389 )]));
2390
2391 let batches = do_read(buf, 1024, false, false, schema);
2392 assert_eq!(batches.len(), 1);
2393
2394 let col = batches[0].column(0).as_fixed_size_list();
2395 assert_eq!(col.len(), 3);
2396 assert_eq!(col.value_length(), 0);
2397
2398 let values = col.values().as_primitive::<Int32Type>();
2399 assert!(values.values().is_empty());
2400 }
2401
2402 #[test]
2403 fn test_fixed_size_list_wrong_size() {
2404 let buf = r#"{"a": [1, 2, 3]}"#;
2405
2406 let field = Field::new_list_field(DataType::Int32, true);
2407 let schema = Arc::new(Schema::new(vec![Field::new(
2408 "a",
2409 DataType::FixedSizeList(Arc::new(field), 2),
2410 false,
2411 )]));
2412
2413 let err = ReaderBuilder::new(schema)
2414 .build(Cursor::new(buf.as_bytes()))
2415 .unwrap()
2416 .next()
2417 .unwrap()
2418 .unwrap_err();
2419
2420 assert!(err.to_string().contains("expected 2 but got 3"), "{}", err);
2421 }
2422
2423 #[test]
2424 fn test_fixed_size_list_nested() {
2425 let buf = r#"
2426 {"a": [[1, 2], [3, 4]]}
2427 {"a": [[5, 6], [7, 8]]}
2428 "#;
2429
2430 let inner_field = Field::new_list_field(DataType::Int32, true);
2431 let inner_type = DataType::FixedSizeList(Arc::new(inner_field), 2);
2432 let outer_field = Arc::new(Field::new_list_field(inner_type.clone(), true));
2433 let schema = Arc::new(Schema::new(vec![Field::new(
2434 "a",
2435 DataType::FixedSizeList(outer_field, 2),
2436 false,
2437 )]));
2438
2439 let batches = do_read(buf, 1024, false, false, schema);
2440 assert_eq!(batches.len(), 1);
2441
2442 let col = batches[0].column(0).as_fixed_size_list();
2443 assert_eq!(col.len(), 2);
2444 assert_eq!(col.value_length(), 2);
2445
2446 let inner = col.values().as_fixed_size_list();
2447 assert_eq!(inner.len(), 4);
2448 assert_eq!(inner.value_length(), 2);
2449
2450 let values = inner.values().as_primitive::<Int32Type>();
2451 assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6, 7, 8]);
2452 }
2453
2454 #[test]
2455 fn test_fixed_size_list_ignore_type_conflicts() {
2456 let field = Field::new("item", DataType::Int32, true);
2457 let schema = Arc::new(Schema::new(vec![Field::new(
2458 "a",
2459 DataType::FixedSizeList(Arc::new(field), 2),
2460 true,
2461 )]));
2462
2463 let json = vec![
2464 json!({"a": [1, 2]}),
2465 json!({"a": "not a list"}),
2466 json!({"a": 42}),
2467 json!({"a": [6, 7]}),
2468 ];
2469
2470 let mut decoder = ReaderBuilder::new(schema)
2471 .with_ignore_type_conflicts(true)
2472 .build_decoder()
2473 .unwrap();
2474 decoder.serialize(&json).unwrap();
2475 let batch = decoder.flush().unwrap().unwrap();
2476
2477 let col = batch.column(0).as_fixed_size_list();
2478 assert_eq!(col.len(), 4);
2479 assert!(col.is_valid(0));
2480 assert!(col.is_null(1)); assert!(col.is_null(2)); assert!(col.is_valid(3));
2483
2484 let values = col.values().as_primitive::<Int32Type>();
2485 assert_eq!(values.value(0), 1);
2486 assert_eq!(values.value(1), 2);
2487 assert_eq!(values.value(6), 6);
2488 assert_eq!(values.value(7), 7);
2489 }
2490
2491 #[test]
2492 fn test_skip_empty_lines() {
2493 let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
2494 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(64);
2495 let json_content = "
2496 {\"a\": 1}
2497 {\"a\": 2}
2498 {\"a\": 3}";
2499 let mut reader = builder.build(Cursor::new(json_content)).unwrap();
2500 let batch = reader.next().unwrap().unwrap();
2501
2502 assert_eq!(1, batch.num_columns());
2503 assert_eq!(3, batch.num_rows());
2504
2505 let schema = reader.schema();
2506 let c = schema.column_with_name("a").unwrap();
2507 assert_eq!(&DataType::Int64, c.1.data_type());
2508 }
2509
2510 #[test]
2511 fn test_with_multiple_batches() {
2512 let file = File::open("test/data/basic_nulls.json").unwrap();
2513 let mut reader = BufReader::new(file);
2514 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
2515 reader.rewind().unwrap();
2516
2517 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(5);
2518 let mut reader = builder.build(reader).unwrap();
2519
2520 let mut num_records = Vec::new();
2521 while let Some(rb) = reader.next().transpose().unwrap() {
2522 num_records.push(rb.num_rows());
2523 }
2524
2525 assert_eq!(vec![5, 5, 2], num_records);
2526 }
2527
2528 #[test]
2529 fn test_timestamp_from_json_seconds() {
2530 let schema = Schema::new(vec![Field::new(
2531 "a",
2532 DataType::Timestamp(TimeUnit::Second, None),
2533 true,
2534 )]);
2535
2536 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
2537 let batch = reader.next().unwrap().unwrap();
2538
2539 assert_eq!(1, batch.num_columns());
2540 assert_eq!(12, batch.num_rows());
2541
2542 let schema = reader.schema();
2543 let batch_schema = batch.schema();
2544 assert_eq!(schema, batch_schema);
2545
2546 let a = schema.column_with_name("a").unwrap();
2547 assert_eq!(
2548 &DataType::Timestamp(TimeUnit::Second, None),
2549 a.1.data_type()
2550 );
2551
2552 let aa = batch.column(a.0).as_primitive::<TimestampSecondType>();
2553 assert!(aa.is_valid(0));
2554 assert!(!aa.is_valid(1));
2555 assert!(!aa.is_valid(2));
2556 assert_eq!(1, aa.value(0));
2557 assert_eq!(1, aa.value(3));
2558 assert_eq!(5, aa.value(7));
2559 }
2560
2561 #[test]
2562 fn test_timestamp_from_json_milliseconds() {
2563 let schema = Schema::new(vec![Field::new(
2564 "a",
2565 DataType::Timestamp(TimeUnit::Millisecond, None),
2566 true,
2567 )]);
2568
2569 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
2570 let batch = reader.next().unwrap().unwrap();
2571
2572 assert_eq!(1, batch.num_columns());
2573 assert_eq!(12, batch.num_rows());
2574
2575 let schema = reader.schema();
2576 let batch_schema = batch.schema();
2577 assert_eq!(schema, batch_schema);
2578
2579 let a = schema.column_with_name("a").unwrap();
2580 assert_eq!(
2581 &DataType::Timestamp(TimeUnit::Millisecond, None),
2582 a.1.data_type()
2583 );
2584
2585 let aa = batch.column(a.0).as_primitive::<TimestampMillisecondType>();
2586 assert!(aa.is_valid(0));
2587 assert!(!aa.is_valid(1));
2588 assert!(!aa.is_valid(2));
2589 assert_eq!(1, aa.value(0));
2590 assert_eq!(1, aa.value(3));
2591 assert_eq!(5, aa.value(7));
2592 }
2593
2594 #[test]
2595 fn test_date_from_json_milliseconds() {
2596 let schema = Schema::new(vec![Field::new("a", DataType::Date64, true)]);
2597
2598 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
2599 let batch = reader.next().unwrap().unwrap();
2600
2601 assert_eq!(1, batch.num_columns());
2602 assert_eq!(12, batch.num_rows());
2603
2604 let schema = reader.schema();
2605 let batch_schema = batch.schema();
2606 assert_eq!(schema, batch_schema);
2607
2608 let a = schema.column_with_name("a").unwrap();
2609 assert_eq!(&DataType::Date64, a.1.data_type());
2610
2611 let aa = batch.column(a.0).as_primitive::<Date64Type>();
2612 assert!(aa.is_valid(0));
2613 assert!(!aa.is_valid(1));
2614 assert!(!aa.is_valid(2));
2615 assert_eq!(1, aa.value(0));
2616 assert_eq!(1, aa.value(3));
2617 assert_eq!(5, aa.value(7));
2618 }
2619
2620 #[test]
2621 fn test_time_from_json_nanoseconds() {
2622 let schema = Schema::new(vec![Field::new(
2623 "a",
2624 DataType::Time64(TimeUnit::Nanosecond),
2625 true,
2626 )]);
2627
2628 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
2629 let batch = reader.next().unwrap().unwrap();
2630
2631 assert_eq!(1, batch.num_columns());
2632 assert_eq!(12, batch.num_rows());
2633
2634 let schema = reader.schema();
2635 let batch_schema = batch.schema();
2636 assert_eq!(schema, batch_schema);
2637
2638 let a = schema.column_with_name("a").unwrap();
2639 assert_eq!(&DataType::Time64(TimeUnit::Nanosecond), a.1.data_type());
2640
2641 let aa = batch.column(a.0).as_primitive::<Time64NanosecondType>();
2642 assert!(aa.is_valid(0));
2643 assert!(!aa.is_valid(1));
2644 assert!(!aa.is_valid(2));
2645 assert_eq!(1, aa.value(0));
2646 assert_eq!(1, aa.value(3));
2647 assert_eq!(5, aa.value(7));
2648 }
2649
2650 #[test]
2651 fn test_json_iterator() {
2652 let file = File::open("test/data/basic.json").unwrap();
2653 let mut reader = BufReader::new(file);
2654 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
2655 reader.rewind().unwrap();
2656
2657 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(5);
2658 let reader = builder.build(reader).unwrap();
2659 let schema = reader.schema();
2660 let (col_a_index, _) = schema.column_with_name("a").unwrap();
2661
2662 let mut sum_num_rows = 0;
2663 let mut num_batches = 0;
2664 let mut sum_a = 0;
2665 for batch in reader {
2666 let batch = batch.unwrap();
2667 assert_eq!(8, batch.num_columns());
2668 sum_num_rows += batch.num_rows();
2669 num_batches += 1;
2670 let batch_schema = batch.schema();
2671 assert_eq!(schema, batch_schema);
2672 let a_array = batch.column(col_a_index).as_primitive::<Int64Type>();
2673 sum_a += (0..a_array.len()).map(|i| a_array.value(i)).sum::<i64>();
2674 }
2675 assert_eq!(12, sum_num_rows);
2676 assert_eq!(3, num_batches);
2677 assert_eq!(100000000000011, sum_a);
2678 }
2679
2680 #[test]
2681 fn test_decoder_error() {
2682 let schema = Arc::new(Schema::new(vec![Field::new_struct(
2683 "a",
2684 vec![Field::new("child", DataType::Int32, false)],
2685 true,
2686 )]));
2687
2688 let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder().unwrap();
2689 let _ = decoder.decode(r#"{"a": { "child":"#.as_bytes()).unwrap();
2690 assert!(decoder.tape_decoder.has_partial_row());
2691 assert_eq!(decoder.tape_decoder.num_buffered_rows(), 1);
2692 let _ = decoder.flush().unwrap_err();
2693 assert!(decoder.tape_decoder.has_partial_row());
2694 assert_eq!(decoder.tape_decoder.num_buffered_rows(), 1);
2695
2696 let parse_err = |s: &str| {
2697 ReaderBuilder::new(schema.clone())
2698 .build(Cursor::new(s.as_bytes()))
2699 .unwrap()
2700 .next()
2701 .unwrap()
2702 .unwrap_err()
2703 .to_string()
2704 };
2705
2706 let err = parse_err(r#"{"a": 123}"#);
2707 assert_eq!(
2708 err,
2709 "Json error: whilst decoding field 'a': expected { got 123"
2710 );
2711
2712 let err = parse_err(r#"{"a": ["bar"]}"#);
2713 assert_eq!(
2714 err,
2715 r#"Json error: whilst decoding field 'a': expected { got ["bar"]"#
2716 );
2717
2718 let err = parse_err(r#"{"a": []}"#);
2719 assert_eq!(
2720 err,
2721 "Json error: whilst decoding field 'a': expected { got []"
2722 );
2723
2724 let err = parse_err(r#"{"a": [{"child": 234}]}"#);
2725 assert_eq!(
2726 err,
2727 r#"Json error: whilst decoding field 'a': expected { got [{"child": 234}]"#
2728 );
2729
2730 let err = parse_err(r#"{"a": [{"child": {"foo": [{"foo": ["bar"]}]}}]}"#);
2731 assert_eq!(
2732 err,
2733 r#"Json error: whilst decoding field 'a': expected { got [{"child": {"foo": [{"foo": ["bar"]}]}}]"#
2734 );
2735
2736 let err = parse_err(r#"{"a": true}"#);
2737 assert_eq!(
2738 err,
2739 "Json error: whilst decoding field 'a': expected { got true"
2740 );
2741
2742 let err = parse_err(r#"{"a": false}"#);
2743 assert_eq!(
2744 err,
2745 "Json error: whilst decoding field 'a': expected { got false"
2746 );
2747
2748 let err = parse_err(r#"{"a": "foo"}"#);
2749 assert_eq!(
2750 err,
2751 "Json error: whilst decoding field 'a': expected { got \"foo\""
2752 );
2753
2754 let err = parse_err(r#"{"a": {"child": false}}"#);
2755 assert_eq!(
2756 err,
2757 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got false"
2758 );
2759
2760 let err = parse_err(r#"{"a": {"child": []}}"#);
2761 assert_eq!(
2762 err,
2763 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got []"
2764 );
2765
2766 let err = parse_err(r#"{"a": {"child": [123]}}"#);
2767 assert_eq!(
2768 err,
2769 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got [123]"
2770 );
2771
2772 let err = parse_err(r#"{"a": {"child": [123, 3465346]}}"#);
2773 assert_eq!(
2774 err,
2775 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got [123, 3465346]"
2776 );
2777 }
2778
2779 #[test]
2780 fn test_serialize_timestamp() {
2781 let json = vec![
2782 json!({"timestamp": 1681319393}),
2783 json!({"timestamp": "1970-01-01T00:00:00+02:00"}),
2784 ];
2785 let schema = Schema::new(vec![Field::new(
2786 "timestamp",
2787 DataType::Timestamp(TimeUnit::Second, None),
2788 true,
2789 )]);
2790 let mut decoder = ReaderBuilder::new(Arc::new(schema))
2791 .build_decoder()
2792 .unwrap();
2793 decoder.serialize(&json).unwrap();
2794 let batch = decoder.flush().unwrap().unwrap();
2795 assert_eq!(batch.num_rows(), 2);
2796 assert_eq!(batch.num_columns(), 1);
2797 let values = batch.column(0).as_primitive::<TimestampSecondType>();
2798 assert_eq!(values.values(), &[1681319393, -7200]);
2799 }
2800
2801 #[test]
2802 fn test_serialize_decimal() {
2803 let json = vec![
2804 json!({"decimal": 1.234}),
2805 json!({"decimal": "1.234"}),
2806 json!({"decimal": 1234}),
2807 json!({"decimal": "1234"}),
2808 ];
2809 let schema = Schema::new(vec![Field::new(
2810 "decimal",
2811 DataType::Decimal128(10, 3),
2812 true,
2813 )]);
2814 let mut decoder = ReaderBuilder::new(Arc::new(schema))
2815 .build_decoder()
2816 .unwrap();
2817 decoder.serialize(&json).unwrap();
2818 let batch = decoder.flush().unwrap().unwrap();
2819 assert_eq!(batch.num_rows(), 4);
2820 assert_eq!(batch.num_columns(), 1);
2821 let values = batch.column(0).as_primitive::<Decimal128Type>();
2822 assert_eq!(values.values(), &[1234, 1234, 1234000, 1234000]);
2823 }
2824
2825 #[test]
2826 fn test_serde_field() {
2827 let field = Field::new("int", DataType::Int32, true);
2828 let mut decoder = ReaderBuilder::new_with_field(field)
2829 .build_decoder()
2830 .unwrap();
2831 decoder.serialize(&[1_i32, 2, 3, 4]).unwrap();
2832 let b = decoder.flush().unwrap().unwrap();
2833 let values = b.column(0).as_primitive::<Int32Type>().values();
2834 assert_eq!(values, &[1, 2, 3, 4]);
2835 }
2836
2837 #[test]
2838 fn test_serde_large_numbers() {
2839 let field = Field::new("int", DataType::Int64, true);
2840 let mut decoder = ReaderBuilder::new_with_field(field)
2841 .build_decoder()
2842 .unwrap();
2843
2844 decoder.serialize(&[1699148028689_u64, 2, 3, 4]).unwrap();
2845 let b = decoder.flush().unwrap().unwrap();
2846 let values = b.column(0).as_primitive::<Int64Type>().values();
2847 assert_eq!(values, &[1699148028689, 2, 3, 4]);
2848
2849 let field = Field::new(
2850 "int",
2851 DataType::Timestamp(TimeUnit::Microsecond, None),
2852 true,
2853 );
2854 let mut decoder = ReaderBuilder::new_with_field(field)
2855 .build_decoder()
2856 .unwrap();
2857
2858 decoder.serialize(&[1699148028689_u64, 2, 3, 4]).unwrap();
2859 let b = decoder.flush().unwrap().unwrap();
2860 let values = b
2861 .column(0)
2862 .as_primitive::<TimestampMicrosecondType>()
2863 .values();
2864 assert_eq!(values, &[1699148028689, 2, 3, 4]);
2865 }
2866
2867 #[test]
2868 fn test_coercing_primitive_into_string_decoder() {
2869 let buf = &format!(
2870 r#"[{{"a": 1, "b": "A", "c": "T"}}, {{"a": 2, "b": "BB", "c": "F"}}, {{"a": {}, "b": 123, "c": false}}, {{"a": {}, "b": 789, "c": true}}]"#,
2871 (i32::MAX as i64 + 10),
2872 i64::MAX - 10
2873 );
2874 let schema = Schema::new(vec![
2875 Field::new("a", DataType::Float64, true),
2876 Field::new("b", DataType::Utf8, true),
2877 Field::new("c", DataType::Utf8, true),
2878 ]);
2879 let json_array: Vec<serde_json::Value> = serde_json::from_str(buf).unwrap();
2880 let schema_ref = Arc::new(schema);
2881
2882 let reader = ReaderBuilder::new(schema_ref.clone()).with_coerce_primitive(true);
2884 let mut decoder = reader.build_decoder().unwrap();
2885 decoder.serialize(json_array.as_slice()).unwrap();
2886 let batch = decoder.flush().unwrap().unwrap();
2887 assert_eq!(
2888 batch,
2889 RecordBatch::try_new(
2890 schema_ref,
2891 vec![
2892 Arc::new(Float64Array::from(vec![
2893 1.0,
2894 2.0,
2895 (i32::MAX as i64 + 10) as f64,
2896 (i64::MAX - 10) as f64
2897 ])),
2898 Arc::new(StringArray::from(vec!["A", "BB", "123", "789"])),
2899 Arc::new(StringArray::from(vec!["T", "F", "false", "true"])),
2900 ]
2901 )
2902 .unwrap()
2903 );
2904 }
2905
2906 fn _parse_structs(
2911 row: &str,
2912 struct_mode: StructMode,
2913 fields: Fields,
2914 as_struct: bool,
2915 ) -> Result<RecordBatch, ArrowError> {
2916 let builder = if as_struct {
2917 ReaderBuilder::new_with_field(Field::new("r", DataType::Struct(fields), true))
2918 } else {
2919 ReaderBuilder::new(Arc::new(Schema::new(fields)))
2920 };
2921 builder
2922 .with_struct_mode(struct_mode)
2923 .build(Cursor::new(row.as_bytes()))
2924 .unwrap()
2925 .next()
2926 .unwrap()
2927 }
2928
2929 #[test]
2930 fn test_struct_decoding_list_length() {
2931 use arrow_array::array;
2932
2933 let row = "[1, 2]";
2934
2935 let mut fields = vec![Field::new("a", DataType::Int32, true)];
2936 let too_few_fields = Fields::from(fields.clone());
2937 fields.push(Field::new("b", DataType::Int32, true));
2938 let correct_fields = Fields::from(fields.clone());
2939 fields.push(Field::new("c", DataType::Int32, true));
2940 let too_many_fields = Fields::from(fields.clone());
2941
2942 let parse = |fields: Fields, as_struct: bool| {
2943 _parse_structs(row, StructMode::ListOnly, fields, as_struct)
2944 };
2945
2946 let expected_row = StructArray::new(
2947 correct_fields.clone(),
2948 vec![
2949 Arc::new(array::Int32Array::from(vec![1])),
2950 Arc::new(array::Int32Array::from(vec![2])),
2951 ],
2952 None,
2953 );
2954 let row_field = Field::new("r", DataType::Struct(correct_fields.clone()), true);
2955
2956 assert_eq!(
2957 parse(too_few_fields.clone(), true).unwrap_err().to_string(),
2958 "Json error: found extra columns for 1 fields".to_string()
2959 );
2960 assert_eq!(
2961 parse(too_few_fields, false).unwrap_err().to_string(),
2962 "Json error: found extra columns for 1 fields".to_string()
2963 );
2964 assert_eq!(
2965 parse(correct_fields.clone(), true).unwrap(),
2966 RecordBatch::try_new(
2967 Arc::new(Schema::new(vec![row_field])),
2968 vec![Arc::new(expected_row.clone())]
2969 )
2970 .unwrap()
2971 );
2972 assert_eq!(
2973 parse(correct_fields, false).unwrap(),
2974 RecordBatch::from(expected_row)
2975 );
2976 assert_eq!(
2977 parse(too_many_fields.clone(), true)
2978 .unwrap_err()
2979 .to_string(),
2980 "Json error: found 2 columns for 3 fields".to_string()
2981 );
2982 assert_eq!(
2983 parse(too_many_fields, false).unwrap_err().to_string(),
2984 "Json error: found 2 columns for 3 fields".to_string()
2985 );
2986 }
2987
2988 #[test]
2989 fn test_struct_decoding() {
2990 use arrow_array::builder;
2991
2992 let nested_object_json = r#"{"a": {"b": [1, 2], "c": {"d": 3}}}"#;
2993 let nested_list_json = r#"[[[1, 2], {"d": 3}]]"#;
2994 let nested_mixed_json = r#"{"a": [[1, 2], {"d": 3}]}"#;
2995
2996 let struct_fields = Fields::from(vec![
2997 Field::new("b", DataType::new_list(DataType::Int32, true), true),
2998 Field::new_map(
2999 "c",
3000 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3001 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3002 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
3003 false,
3004 false,
3005 ),
3006 ]);
3007
3008 let list_array =
3009 ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(1), Some(2)])]);
3010
3011 let map_array = {
3012 let mut map_builder = builder::MapBuilder::new(
3013 None,
3014 builder::StringBuilder::new(),
3015 builder::Int32Builder::new(),
3016 );
3017 map_builder.keys().append_value("d");
3018 map_builder.values().append_value(3);
3019 map_builder.append(true).unwrap();
3020 map_builder.finish()
3021 };
3022
3023 let struct_array = StructArray::new(
3024 struct_fields.clone(),
3025 vec![Arc::new(list_array), Arc::new(map_array)],
3026 None,
3027 );
3028
3029 let fields = Fields::from(vec![Field::new("a", DataType::Struct(struct_fields), true)]);
3030 let schema = Arc::new(Schema::new(fields.clone()));
3031 let expected = RecordBatch::try_new(schema.clone(), vec![Arc::new(struct_array)]).unwrap();
3032
3033 let parse = |row: &str, struct_mode: StructMode| {
3034 _parse_structs(row, struct_mode, fields.clone(), false)
3035 };
3036
3037 assert_eq!(
3038 parse(nested_object_json, StructMode::ObjectOnly).unwrap(),
3039 expected
3040 );
3041 assert_eq!(
3042 parse(nested_list_json, StructMode::ObjectOnly)
3043 .unwrap_err()
3044 .to_string(),
3045 "Json error: expected { got [[[1, 2], {\"d\": 3}]]".to_owned()
3046 );
3047 assert_eq!(
3048 parse(nested_mixed_json, StructMode::ObjectOnly)
3049 .unwrap_err()
3050 .to_string(),
3051 "Json error: whilst decoding field 'a': expected { got [[1, 2], {\"d\": 3}]".to_owned()
3052 );
3053
3054 assert_eq!(
3055 parse(nested_list_json, StructMode::ListOnly).unwrap(),
3056 expected
3057 );
3058 assert_eq!(
3059 parse(nested_object_json, StructMode::ListOnly)
3060 .unwrap_err()
3061 .to_string(),
3062 "Json error: expected [ got {\"a\": {\"b\": [1, 2]\"c\": {\"d\": 3}}}".to_owned()
3063 );
3064 assert_eq!(
3065 parse(nested_mixed_json, StructMode::ListOnly)
3066 .unwrap_err()
3067 .to_string(),
3068 "Json error: expected [ got {\"a\": [[1, 2], {\"d\": 3}]}".to_owned()
3069 );
3070 }
3071
3072 #[test]
3078 fn test_struct_decoding_empty_list() {
3079 let int_field = Field::new("a", DataType::Int32, true);
3080 let struct_field = Field::new(
3081 "r",
3082 DataType::Struct(Fields::from(vec![int_field.clone()])),
3083 true,
3084 );
3085
3086 let parse = |row: &str, as_struct: bool, field: Field| {
3087 _parse_structs(
3088 row,
3089 StructMode::ListOnly,
3090 Fields::from(vec![field]),
3091 as_struct,
3092 )
3093 };
3094
3095 assert_eq!(
3097 parse("[]", true, struct_field.clone())
3098 .unwrap_err()
3099 .to_string(),
3100 "Json error: found 0 columns for 1 fields".to_owned()
3101 );
3102 assert_eq!(
3103 parse("[]", false, int_field.clone())
3104 .unwrap_err()
3105 .to_string(),
3106 "Json error: found 0 columns for 1 fields".to_owned()
3107 );
3108 assert_eq!(
3109 parse("[]", false, struct_field.clone())
3110 .unwrap_err()
3111 .to_string(),
3112 "Json error: found 0 columns for 1 fields".to_owned()
3113 );
3114 assert_eq!(
3115 parse("[[]]", false, struct_field.clone())
3116 .unwrap_err()
3117 .to_string(),
3118 "Json error: whilst decoding field 'r': found 0 columns for 1 fields".to_owned()
3119 );
3120 }
3121
3122 #[test]
3123 fn test_decode_list_struct_with_wrong_types() {
3124 let int_field = Field::new("a", DataType::Int32, true);
3125 let struct_field = Field::new(
3126 "r",
3127 DataType::Struct(Fields::from(vec![int_field.clone()])),
3128 true,
3129 );
3130
3131 let parse = |row: &str, as_struct: bool, field: Field| {
3132 _parse_structs(
3133 row,
3134 StructMode::ListOnly,
3135 Fields::from(vec![field]),
3136 as_struct,
3137 )
3138 };
3139
3140 assert_eq!(
3142 parse(r#"[["a"]]"#, false, struct_field.clone())
3143 .unwrap_err()
3144 .to_string(),
3145 "Json error: whilst decoding field 'r': whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3146 );
3147 assert_eq!(
3148 parse(r#"[["a"]]"#, true, struct_field.clone())
3149 .unwrap_err()
3150 .to_string(),
3151 "Json error: whilst decoding field 'r': whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3152 );
3153 assert_eq!(
3154 parse(r#"["a"]"#, true, int_field.clone())
3155 .unwrap_err()
3156 .to_string(),
3157 "Json error: whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3158 );
3159 assert_eq!(
3160 parse(r#"["a"]"#, false, int_field.clone())
3161 .unwrap_err()
3162 .to_string(),
3163 "Json error: whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3164 );
3165 }
3166
3167 #[test]
3168 fn test_type_conflict_nulls() {
3169 let schema = Schema::new(vec![
3170 Field::new("null", DataType::Null, true),
3171 Field::new("bool", DataType::Boolean, true),
3172 Field::new("primitive", DataType::Int32, true),
3173 Field::new("numeric", DataType::Decimal128(10, 3), true),
3174 Field::new("string", DataType::Utf8, true),
3175 Field::new("string_view", DataType::Utf8View, true),
3176 Field::new(
3177 "timestamp",
3178 DataType::Timestamp(TimeUnit::Second, None),
3179 true,
3180 ),
3181 Field::new(
3182 "array",
3183 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
3184 true,
3185 ),
3186 Field::new(
3187 "map",
3188 DataType::Map(
3189 Arc::new(Field::new(
3190 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3191 DataType::Struct(Fields::from(vec![
3192 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3193 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3194 ])),
3195 false, )),
3197 false, ),
3199 true, ),
3201 Field::new(
3202 "struct",
3203 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
3204 true,
3205 ),
3206 ]);
3207
3208 let json_values = vec![
3210 json!(null),
3211 json!(true),
3212 json!(42),
3213 json!(1.234),
3214 json!("hi"),
3215 json!("ho"),
3216 json!("1970-01-01T00:00:00+02:00"),
3217 json!([1, "ho", 3]),
3218 json!({"k": "value"}),
3219 json!({"a": 1}),
3220 ];
3221
3222 let json: Vec<_> = (0..json_values.len())
3224 .map(|i| {
3225 let pairs = json_values[i..]
3226 .iter()
3227 .chain(json_values[..i].iter())
3228 .zip(&schema.fields)
3229 .map(|(v, f)| (f.name().to_string(), v.clone()))
3230 .collect();
3231 serde_json::Value::Object(pairs)
3232 })
3233 .collect();
3234 let mut decoder = ReaderBuilder::new(Arc::new(schema))
3235 .with_ignore_type_conflicts(true)
3236 .with_coerce_primitive(true)
3237 .build_decoder()
3238 .unwrap();
3239 decoder.serialize(&json).unwrap();
3240 let batch = decoder.flush().unwrap().unwrap();
3241 assert_eq!(batch.num_rows(), 10);
3242 assert_eq!(batch.num_columns(), 10);
3243
3244 let _ = batch
3246 .column(0)
3247 .as_any()
3248 .downcast_ref::<NullArray>()
3249 .unwrap();
3250
3251 assert!(
3252 batch
3253 .column(1)
3254 .as_any()
3255 .downcast_ref::<BooleanArray>()
3256 .unwrap()
3257 .iter()
3258 .eq([
3259 Some(true),
3260 None,
3261 None,
3262 None,
3263 None,
3264 None,
3265 None,
3266 None,
3267 None,
3268 None
3269 ])
3270 );
3271
3272 assert!(batch.column(2).as_primitive::<Int32Type>().iter().eq([
3273 Some(42),
3274 Some(1),
3275 None,
3276 None,
3277 None,
3278 None,
3279 None,
3280 None,
3281 None,
3282 None
3283 ]));
3284
3285 assert!(batch.column(3).as_primitive::<Decimal128Type>().iter().eq([
3286 Some(1234),
3287 None,
3288 None,
3289 None,
3290 None,
3291 None,
3292 None,
3293 None,
3294 None,
3295 Some(42000)
3296 ]));
3297
3298 assert!(
3299 batch
3300 .column(4)
3301 .as_any()
3302 .downcast_ref::<StringArray>()
3303 .unwrap()
3304 .iter()
3305 .eq([
3306 Some("hi"),
3307 Some("ho"),
3308 Some("1970-01-01T00:00:00+02:00"),
3309 None,
3310 None,
3311 None,
3312 None,
3313 Some("true"),
3314 Some("42"),
3315 Some("1.234"),
3316 ])
3317 );
3318
3319 assert!(
3320 batch
3321 .column(5)
3322 .as_any()
3323 .downcast_ref::<StringViewArray>()
3324 .unwrap()
3325 .iter()
3326 .eq([
3327 Some("ho"),
3328 Some("1970-01-01T00:00:00+02:00"),
3329 None,
3330 None,
3331 None,
3332 None,
3333 Some("true"),
3334 Some("42"),
3335 Some("1.234"),
3336 Some("hi"),
3337 ])
3338 );
3339
3340 assert!(
3341 batch
3342 .column(6)
3343 .as_primitive::<TimestampSecondType>()
3344 .iter()
3345 .eq([
3346 Some(-7200),
3347 None,
3348 None,
3349 None,
3350 None,
3351 None,
3352 Some(42),
3353 None,
3354 None,
3355 None,
3356 ])
3357 );
3358
3359 let arrays = batch
3360 .column(7)
3361 .as_any()
3362 .downcast_ref::<ListArray>()
3363 .unwrap();
3364 assert_eq!(
3365 arrays.nulls(),
3366 Some(&NullBuffer::from(
3367 &[
3368 true, false, false, false, false, false, false, false, false, false
3369 ][..]
3370 ))
3371 );
3372 assert_eq!(arrays.offsets()[1], 3);
3373 let array_values = arrays
3374 .values()
3375 .as_any()
3376 .downcast_ref::<Int32Array>()
3377 .unwrap();
3378 assert!(array_values.iter().eq([Some(1), None, Some(3)]));
3379
3380 let maps = batch.column(8).as_any().downcast_ref::<MapArray>().unwrap();
3381 assert_eq!(
3382 maps.nulls(),
3383 Some(&NullBuffer::from(
3384 &[
3386 true, true, false, false, false, false, false, false, false, false
3387 ][..]
3388 ))
3389 );
3390 let map_keys = maps.keys().as_any().downcast_ref::<StringArray>().unwrap();
3391 assert!(map_keys.iter().eq([Some("k"), Some("a")]));
3392 let map_values = maps
3393 .values()
3394 .as_any()
3395 .downcast_ref::<StringArray>()
3396 .unwrap();
3397 assert!(map_values.iter().eq([Some("value"), Some("1")]));
3398
3399 let structs = batch
3400 .column(9)
3401 .as_any()
3402 .downcast_ref::<StructArray>()
3403 .unwrap();
3404 assert_eq!(
3405 structs.nulls(),
3406 Some(&NullBuffer::from(
3407 &[
3409 true, false, false, false, false, false, false, false, false, true
3410 ][..]
3411 ))
3412 );
3413 let struct_fields = structs
3414 .column(0)
3415 .as_any()
3416 .downcast_ref::<Int32Array>()
3417 .unwrap();
3418 assert!(struct_fields.slice(0, 2).iter().eq([Some(1), None]));
3419 }
3420
3421 #[test]
3422 fn test_type_conflict_non_nullable() {
3423 let fields = [
3424 Field::new("bool", DataType::Boolean, false),
3425 Field::new("primitive", DataType::Int32, false),
3426 Field::new("numeric", DataType::Decimal128(10, 3), false),
3427 Field::new("string", DataType::Utf8, false),
3428 Field::new("string_view", DataType::Utf8View, false),
3429 Field::new(
3430 "timestamp",
3431 DataType::Timestamp(TimeUnit::Second, None),
3432 false,
3433 ),
3434 Field::new(
3435 "array",
3436 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
3437 false,
3438 ),
3439 Field::new(
3440 "fixed_size_list",
3441 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 2),
3442 false,
3443 ),
3444 Field::new(
3445 "map",
3446 DataType::Map(
3447 Arc::new(Field::new(
3448 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3449 DataType::Struct(Fields::from(vec![
3450 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3451 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3452 ])),
3453 false, )),
3455 false, ),
3457 false, ),
3459 Field::new(
3460 "struct",
3461 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
3462 false,
3463 ),
3464 ];
3465
3466 let json_values = vec![json!(true), json!({"a": 1})];
3468
3469 for field in fields {
3470 let mut decoder = ReaderBuilder::new_with_field(field)
3471 .with_ignore_type_conflicts(true)
3472 .build_decoder()
3473 .unwrap();
3474 decoder.serialize(&json_values).unwrap();
3475 decoder
3476 .flush()
3477 .expect_err("type conflict on non-nullable type");
3478 }
3479 }
3480
3481 #[test]
3482 fn test_ignore_type_conflicts_disabled() {
3483 let fields = [
3484 Field::new("null", DataType::Null, true),
3485 Field::new("bool", DataType::Boolean, true),
3486 Field::new("primitive", DataType::Int32, true),
3487 Field::new("numeric", DataType::Decimal128(10, 3), true),
3488 Field::new("string", DataType::Utf8, true),
3489 Field::new("string_view", DataType::Utf8View, true),
3490 Field::new(
3491 "timestamp",
3492 DataType::Timestamp(TimeUnit::Second, None),
3493 true,
3494 ),
3495 Field::new(
3496 "array",
3497 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
3498 true,
3499 ),
3500 Field::new(
3501 "fixed_size_list",
3502 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 2),
3503 true,
3504 ),
3505 Field::new(
3506 "map",
3507 DataType::Map(
3508 Arc::new(Field::new(
3509 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3510 DataType::Struct(Fields::from(vec![
3511 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3512 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3513 ])),
3514 false, )),
3516 false, ),
3518 true, ),
3520 Field::new(
3521 "struct",
3522 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
3523 true,
3524 ),
3525 ];
3526
3527 let json_values = vec![json!(true), json!({"a": 1})];
3529
3530 for field in fields {
3531 let mut decoder = ReaderBuilder::new_with_field(field)
3532 .build_decoder()
3533 .unwrap();
3534 decoder.serialize(&json_values).unwrap();
3535 decoder
3536 .flush()
3537 .expect_err("type conflict on non-nullable type");
3538 }
3539 }
3540
3541 #[test]
3542 fn test_read_run_end_encoded() {
3543 let buf = r#"
3544 {"a": "x"}
3545 {"a": "x"}
3546 {"a": "y"}
3547 {"a": "y"}
3548 {"a": "y"}
3549 "#;
3550
3551 let ree_type = DataType::RunEndEncoded(
3552 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3553 Arc::new(Field::new("values", DataType::Utf8, true)),
3554 );
3555 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
3556 let batches = do_read(buf, 1024, false, false, schema);
3557 assert_eq!(batches.len(), 1);
3558
3559 let col = batches[0].column(0);
3560 let run_array = col.as_run::<arrow_array::types::Int32Type>();
3561
3562 assert_eq!(run_array.len(), 5);
3564 assert_eq!(run_array.run_ends().values(), &[2, 5]);
3565
3566 let values = run_array.values().as_string::<i32>();
3567 assert_eq!(values.len(), 2);
3568 assert_eq!(values.value(0), "x");
3569 assert_eq!(values.value(1), "y");
3570 }
3571
3572 #[test]
3573 fn test_read_run_end_encoded_consecutive_nulls() {
3574 let buf = r#"
3575 {"a": "x"}
3576 {}
3577 {}
3578 {}
3579 {"a": "y"}
3580 "#;
3581
3582 let ree_type = DataType::RunEndEncoded(
3583 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3584 Arc::new(Field::new("values", DataType::Utf8, true)),
3585 );
3586 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
3587 let batches = do_read(buf, 1024, false, false, schema);
3588 assert_eq!(batches.len(), 1);
3589
3590 let col = batches[0].column(0);
3591 let run_array = col.as_run::<arrow_array::types::Int32Type>();
3592
3593 assert_eq!(run_array.len(), 5);
3595 assert_eq!(run_array.run_ends().values(), &[1, 4, 5]);
3596
3597 let values = run_array.values().as_string::<i32>();
3598 assert_eq!(values.len(), 3);
3599 assert_eq!(values.value(0), "x");
3600 assert!(values.is_null(1));
3601 assert_eq!(values.value(2), "y");
3602 }
3603
3604 #[test]
3605 fn test_read_run_end_encoded_all_unique() {
3606 let buf = r#"
3607 {"a": 1}
3608 {"a": 2}
3609 {"a": 3}
3610 "#;
3611
3612 let ree_type = DataType::RunEndEncoded(
3613 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3614 Arc::new(Field::new("values", DataType::Int32, true)),
3615 );
3616 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
3617 let batches = do_read(buf, 1024, false, false, schema);
3618 assert_eq!(batches.len(), 1);
3619
3620 let col = batches[0].column(0);
3621 let run_array = col.as_run::<arrow_array::types::Int32Type>();
3622
3623 assert_eq!(run_array.len(), 3);
3625 assert_eq!(run_array.run_ends().values(), &[1, 2, 3]);
3626 }
3627
3628 #[test]
3629 fn test_read_run_end_encoded_int16_run_ends() {
3630 let buf = r#"
3631 {"a": "x"}
3632 {"a": "x"}
3633 {"a": "y"}
3634 "#;
3635
3636 let ree_type = DataType::RunEndEncoded(
3637 Arc::new(Field::new("run_ends", DataType::Int16, false)),
3638 Arc::new(Field::new("values", DataType::Utf8, true)),
3639 );
3640 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
3641 let batches = do_read(buf, 1024, false, false, schema);
3642 assert_eq!(batches.len(), 1);
3643
3644 let col = batches[0].column(0);
3645 let run_array = col.as_run::<arrow_array::types::Int16Type>();
3646
3647 assert_eq!(run_array.len(), 3);
3648 assert_eq!(run_array.run_ends().values(), &[2i16, 3]);
3649 }
3650
3651 #[test]
3652 fn test_read_nested_run_end_encoded() {
3653 let buf = r#"
3654 {"a": "x"}
3655 {"a": "x"}
3656 {"a": "y"}
3657 "#;
3658
3659 let inner_type = DataType::RunEndEncoded(
3662 Arc::new(Field::new("run_ends", DataType::Int64, false)),
3663 Arc::new(Field::new("values", DataType::Utf8, true)),
3664 );
3665 let outer_type = DataType::RunEndEncoded(
3666 Arc::new(Field::new("run_ends", DataType::Int64, false)),
3667 Arc::new(Field::new("values", inner_type, true)),
3668 );
3669 let schema = Arc::new(Schema::new(vec![Field::new("a", outer_type, true)]));
3670 let batches = do_read(buf, 1024, false, false, schema);
3671 assert_eq!(batches.len(), 1);
3672
3673 let col = batches[0].column(0);
3674 let outer = col.as_run::<arrow_array::types::Int64Type>();
3675 assert_eq!(outer.len(), 3);
3677 assert_eq!(outer.run_ends().values(), &[2, 3]);
3678
3679 let nested = outer.values().as_run::<arrow_array::types::Int64Type>();
3680 assert_eq!(nested.len(), 2);
3682 assert_eq!(nested.run_ends().values(), &[1, 2]);
3683
3684 let nested_values = nested.values().as_string::<i32>();
3685 assert_eq!(nested_values.len(), 2);
3686 assert_eq!(nested_values.value(0), "x");
3687 assert_eq!(nested_values.value(1), "y");
3688 }
3689}