1use std::io::BufRead;
145use std::sync::Arc;
146
147use arrow_array::cast::AsArray;
148use arrow_array::timezone::Tz;
149use arrow_array::types::*;
150use arrow_array::{ArrayRef, RecordBatch, RecordBatchReader, downcast_integer};
151use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema, SchemaRef, TimeUnit};
152use chrono::Utc;
153use serde_core::Serialize;
154
155use crate::StructMode;
156use crate::reader::binary_array::{
157 BinaryArrayDecoder, BinaryViewDecoder, FixedSizeBinaryArrayDecoder,
158};
159use crate::reader::boolean_array::BooleanArrayDecoder;
160use crate::reader::decimal_array::DecimalArrayDecoder;
161use crate::reader::list_array::{
162 FixedSizeListArrayDecoder, ListArrayDecoder, ListViewArrayDecoder,
163};
164use crate::reader::map_array::MapArrayDecoder;
165use crate::reader::null_array::NullArrayDecoder;
166use crate::reader::primitive_array::PrimitiveArrayDecoder;
167use crate::reader::run_end_array::RunEndEncodedArrayDecoder;
168use crate::reader::string_array::StringArrayDecoder;
169use crate::reader::string_view_array::StringViewArrayDecoder;
170use crate::reader::struct_array::StructArrayDecoder;
171use crate::reader::tape::{TapeDecoder, TapeDecoderOptions};
172use crate::reader::timestamp_array::TimestampArrayDecoder;
173
174pub use schema::*;
175pub use tape::{Tape, TapeElement};
178pub use value_iter::ValueIter;
179
180mod binary_array;
181mod boolean_array;
182mod decimal_array;
183mod list_array;
184mod map_array;
185mod null_array;
186mod primitive_array;
187mod run_end_array;
188mod schema;
189mod serializer;
190mod string_array;
191mod string_view_array;
192mod struct_array;
193mod tape;
194mod timestamp_array;
195mod value_iter;
196
197#[derive(Clone)]
199pub struct ReaderBuilder {
200 batch_size: usize,
201 coerce_primitive: bool,
202 strict_mode: bool,
203 ignore_type_conflicts: bool,
204 is_field: bool,
205 struct_mode: StructMode,
206 flatten_top_level_arrays: bool,
207 decoder_factory: Option<Arc<dyn DecoderFactory>>,
208 schema: SchemaRef,
209}
210
211impl ReaderBuilder {
212 pub fn new(schema: SchemaRef) -> Self {
221 Self {
222 batch_size: 1024,
223 coerce_primitive: false,
224 strict_mode: false,
225 ignore_type_conflicts: false,
226 is_field: false,
227 struct_mode: Default::default(),
228 flatten_top_level_arrays: false,
229 decoder_factory: None,
230 schema,
231 }
232 }
233
234 pub fn new_with_field(field: impl Into<FieldRef>) -> Self {
265 Self {
266 batch_size: 1024,
267 coerce_primitive: false,
268 strict_mode: false,
269 ignore_type_conflicts: false,
270 is_field: true,
271 struct_mode: Default::default(),
272 flatten_top_level_arrays: false,
273 decoder_factory: None,
274 schema: Arc::new(Schema::new([field.into()])),
275 }
276 }
277
278 pub fn with_batch_size(self, batch_size: usize) -> Self {
280 Self { batch_size, ..self }
281 }
282
283 pub fn with_coerce_primitive(self, coerce_primitive: bool) -> Self {
286 Self {
287 coerce_primitive,
288 ..self
289 }
290 }
291
292 pub fn with_strict_mode(self, strict_mode: bool) -> Self {
298 Self {
299 strict_mode,
300 ..self
301 }
302 }
303
304 pub fn with_struct_mode(self, struct_mode: StructMode) -> Self {
308 Self {
309 struct_mode,
310 ..self
311 }
312 }
313
314 pub fn with_ignore_type_conflicts(self, ignore_type_conflicts: bool) -> Self {
327 Self {
328 ignore_type_conflicts,
329 ..self
330 }
331 }
332 pub fn with_flatten(self, flatten_top_level_arrays: bool) -> Self {
350 Self {
351 flatten_top_level_arrays,
352 ..self
353 }
354 }
355
356 pub fn with_decoder_factory(self, decoder_factory: Arc<dyn DecoderFactory>) -> Self {
358 Self {
359 decoder_factory: Some(decoder_factory),
360 ..self
361 }
362 }
363
364 pub fn build<R: BufRead>(self, reader: R) -> Result<Reader<R>, ArrowError> {
366 Ok(Reader {
367 reader,
368 decoder: self.build_decoder()?,
369 })
370 }
371
372 pub fn build_decoder(self) -> Result<Decoder, ArrowError> {
374 let (field, nullable) = if self.is_field {
375 let field = self.schema.fields[0].clone();
376 let nullable = field.is_nullable();
377 (field, nullable)
378 } else {
379 let data_type = DataType::Struct(self.schema.fields.clone());
381 (Arc::new(Field::new("", data_type, false)), false)
382 };
383
384 let ctx = DecoderContext {
385 coerce_primitive: self.coerce_primitive,
386 strict_mode: self.strict_mode,
387 struct_mode: self.struct_mode,
388 ignore_type_conflicts: self.ignore_type_conflicts,
389 decoder_factory: self.decoder_factory,
390 };
391 let decoder = ctx.make_decoder(&field, nullable)?;
392
393 let num_fields = self.schema.flattened_fields().len();
394
395 Ok(Decoder {
396 decoder,
397 is_field: self.is_field,
398 tape_decoder: TapeDecoder::new(TapeDecoderOptions {
399 batch_size: self.batch_size,
400 num_fields,
401 flatten_top_level_arrays: self.flatten_top_level_arrays,
402 }),
403 batch_size: self.batch_size,
404 schema: self.schema,
405 })
406 }
407}
408
409pub struct Reader<R> {
413 reader: R,
414 decoder: Decoder,
415}
416
417impl<R> std::fmt::Debug for Reader<R> {
418 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
419 f.debug_struct("Reader")
420 .field("decoder", &self.decoder)
421 .finish()
422 }
423}
424
425impl<R: BufRead> Reader<R> {
426 fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
428 loop {
429 let buf = self.reader.fill_buf()?;
430 if buf.is_empty() {
431 break;
432 }
433 let read = buf.len();
434
435 let decoded = self.decoder.decode(buf)?;
436 self.reader.consume(decoded);
437 if decoded != read {
438 break;
439 }
440 }
441 self.decoder.flush()
442 }
443}
444
445impl<R: BufRead> Iterator for Reader<R> {
446 type Item = Result<RecordBatch, ArrowError>;
447
448 fn next(&mut self) -> Option<Self::Item> {
449 self.read().transpose()
450 }
451}
452
453impl<R: BufRead> RecordBatchReader for Reader<R> {
454 fn schema(&self) -> SchemaRef {
455 self.decoder.schema.clone()
456 }
457}
458
459pub struct Decoder {
500 tape_decoder: TapeDecoder,
501 decoder: Box<dyn ArrayDecoder>,
502 batch_size: usize,
503 is_field: bool,
504 schema: SchemaRef,
505}
506
507impl std::fmt::Debug for Decoder {
508 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
509 f.debug_struct("Decoder")
510 .field("schema", &self.schema)
511 .field("batch_size", &self.batch_size)
512 .finish()
513 }
514}
515
516impl Decoder {
517 pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
526 self.tape_decoder.decode(buf)
527 }
528
529 pub fn serialize<S: Serialize>(&mut self, rows: &[S]) -> Result<(), ArrowError> {
706 self.tape_decoder.serialize(rows)
707 }
708
709 pub fn has_partial_record(&self) -> bool {
711 self.tape_decoder.has_partial_row()
712 }
713
714 pub fn len(&self) -> usize {
716 self.tape_decoder.num_buffered_rows()
717 }
718
719 pub fn is_empty(&self) -> bool {
721 self.len() == 0
722 }
723
724 pub fn flush(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
731 let tape = self.tape_decoder.finish()?;
732
733 if tape.num_rows() == 0 {
734 return Ok(None);
735 }
736
737 let mut next_object = 1;
739 let pos = (0..tape.num_rows())
740 .map(|_| {
741 let next = tape.next(next_object, "row")?;
742 Ok(std::mem::replace(&mut next_object, next))
743 })
744 .collect::<Result<Vec<_>, ArrowError>>()?;
745
746 let decoded = self.decoder.decode(&tape, &pos)?;
747 self.tape_decoder.clear();
748
749 let batch = match self.is_field {
750 true => RecordBatch::try_new(self.schema.clone(), vec![decoded])?,
751 false => {
752 RecordBatch::from(decoded.as_struct().clone()).with_schema(self.schema.clone())?
753 }
754 };
755
756 Ok(Some(batch))
757 }
758}
759
760pub trait ArrayDecoder: Send {
764 fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError>;
770}
771
772pub trait DecoderFactory: std::fmt::Debug + Send + Sync {
957 fn make_default_decoder(
975 &self,
976 _ctx: &DecoderContext,
977 _field: &FieldRef,
978 _is_nullable: bool,
979 ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
980 Ok(None)
981 }
982}
983
984struct CheckedDecoder {
987 inner: Box<dyn ArrayDecoder>,
988 data_type: DataType,
989}
990
991impl ArrayDecoder for CheckedDecoder {
992 fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
993 let array = self.inner.decode(tape, pos)?;
994 if array.data_type() != &self.data_type {
995 return Err(ArrowError::JsonError(format!(
996 "custom decoder returned {} for a field of type {}",
997 array.data_type(),
998 self.data_type
999 )));
1000 }
1001 if array.len() != pos.len() {
1002 return Err(ArrowError::JsonError(format!(
1003 "custom decoder returned {} values for {} rows",
1004 array.len(),
1005 pos.len()
1006 )));
1007 }
1008 Ok(array)
1009 }
1010}
1011
1012pub struct DecoderContext {
1017 coerce_primitive: bool,
1019 strict_mode: bool,
1021 struct_mode: StructMode,
1023 ignore_type_conflicts: bool,
1025 decoder_factory: Option<Arc<dyn DecoderFactory>>,
1027}
1028
1029impl DecoderContext {
1030 pub fn coerce_primitive(&self) -> bool {
1032 self.coerce_primitive
1033 }
1034
1035 pub fn strict_mode(&self) -> bool {
1037 self.strict_mode
1038 }
1039
1040 pub fn struct_mode(&self) -> StructMode {
1042 self.struct_mode
1043 }
1044
1045 pub fn ignore_type_conflicts(&self) -> bool {
1047 self.ignore_type_conflicts
1048 }
1049
1050 pub fn decoder_factory(&self) -> Option<&Arc<dyn DecoderFactory>> {
1052 self.decoder_factory.as_ref()
1053 }
1054
1055 pub fn make_decoder(
1060 &self,
1061 field: &FieldRef,
1062 is_nullable: bool,
1063 ) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
1064 make_decoder(self, field, is_nullable)
1065 }
1066
1067 pub fn make_builtin_decoder(
1073 &self,
1074 field: &FieldRef,
1075 is_nullable: bool,
1076 ) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
1077 make_builtin_decoder(self, field, is_nullable)
1078 }
1079}
1080
1081fn make_decoder(
1082 ctx: &DecoderContext,
1083 field: &FieldRef,
1084 is_nullable: bool,
1085) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
1086 if let Some(factory) = ctx.decoder_factory()
1087 && let Some(decoder) = factory.make_default_decoder(ctx, field, is_nullable)?
1088 {
1089 return Ok(Box::new(CheckedDecoder {
1090 inner: decoder,
1091 data_type: field.data_type().clone(),
1092 }));
1093 }
1094
1095 make_builtin_decoder(ctx, field, is_nullable)
1096}
1097
1098fn make_builtin_decoder(
1099 ctx: &DecoderContext,
1100 field: &FieldRef,
1101 is_nullable: bool,
1102) -> Result<Box<dyn ArrayDecoder>, ArrowError> {
1103 let data_type = field.data_type();
1104
1105 macro_rules! primitive_decoder {
1106 ($t:ty, $data_type:expr) => {
1107 Ok(Box::new(PrimitiveArrayDecoder::<$t>::new(ctx, $data_type)))
1108 };
1109 }
1110 macro_rules! timestamp_decoder {
1111 ($t:ty, $data_type:expr, $tz:expr) => {{
1112 Ok(Box::new(TimestampArrayDecoder::<$t, _>::new(
1113 ctx, $data_type, $tz,
1114 )))
1115 }};
1116 }
1117 macro_rules! decimal_decoder {
1118 ($t:ty, $p:expr, $s:expr) => {
1119 Ok(Box::new(DecimalArrayDecoder::<$t>::new(ctx, $p, $s)))
1120 };
1121 }
1122
1123 downcast_integer! {
1124 *data_type => (primitive_decoder, data_type),
1125 DataType::Null => Ok(Box::new(NullArrayDecoder::new(ctx))),
1126 DataType::Float16 => primitive_decoder!(Float16Type, data_type),
1127 DataType::Float32 => primitive_decoder!(Float32Type, data_type),
1128 DataType::Float64 => primitive_decoder!(Float64Type, data_type),
1129 DataType::Timestamp(TimeUnit::Second, None) => {
1130 timestamp_decoder!(TimestampSecondType, data_type, Utc)
1131 },
1132 DataType::Timestamp(TimeUnit::Millisecond, None) => {
1133 timestamp_decoder!(TimestampMillisecondType, data_type, Utc)
1134 },
1135 DataType::Timestamp(TimeUnit::Microsecond, None) => {
1136 timestamp_decoder!(TimestampMicrosecondType, data_type, Utc)
1137 },
1138 DataType::Timestamp(TimeUnit::Nanosecond, None) => {
1139 timestamp_decoder!(TimestampNanosecondType, data_type, Utc)
1140 },
1141 DataType::Timestamp(TimeUnit::Second, Some(ref tz)) => {
1142 let tz: Tz = tz.parse()?;
1143 timestamp_decoder!(TimestampSecondType, data_type, tz)
1144 },
1145 DataType::Timestamp(TimeUnit::Millisecond, Some(ref tz)) => {
1146 let tz: Tz = tz.parse()?;
1147 timestamp_decoder!(TimestampMillisecondType, data_type, tz)
1148 },
1149 DataType::Timestamp(TimeUnit::Microsecond, Some(ref tz)) => {
1150 let tz: Tz = tz.parse()?;
1151 timestamp_decoder!(TimestampMicrosecondType, data_type, tz)
1152 },
1153 DataType::Timestamp(TimeUnit::Nanosecond, Some(ref tz)) => {
1154 let tz: Tz = tz.parse()?;
1155 timestamp_decoder!(TimestampNanosecondType, data_type, tz)
1156 },
1157 DataType::Date32 => primitive_decoder!(Date32Type, data_type),
1158 DataType::Date64 => primitive_decoder!(Date64Type, data_type),
1159 DataType::Time32(TimeUnit::Second) => primitive_decoder!(Time32SecondType, data_type),
1160 DataType::Time32(TimeUnit::Millisecond) => primitive_decoder!(Time32MillisecondType, data_type),
1161 DataType::Time64(TimeUnit::Microsecond) => primitive_decoder!(Time64MicrosecondType, data_type),
1162 DataType::Time64(TimeUnit::Nanosecond) => primitive_decoder!(Time64NanosecondType, data_type),
1163 DataType::Duration(TimeUnit::Nanosecond) => primitive_decoder!(DurationNanosecondType, data_type),
1164 DataType::Duration(TimeUnit::Microsecond) => primitive_decoder!(DurationMicrosecondType, data_type),
1165 DataType::Duration(TimeUnit::Millisecond) => primitive_decoder!(DurationMillisecondType, data_type),
1166 DataType::Duration(TimeUnit::Second) => primitive_decoder!(DurationSecondType, data_type),
1167 DataType::Decimal32(p, s) => decimal_decoder!(Decimal32Type, p, s),
1168 DataType::Decimal64(p, s) => decimal_decoder!(Decimal64Type, p, s),
1169 DataType::Decimal128(p, s) => decimal_decoder!(Decimal128Type, p, s),
1170 DataType::Decimal256(p, s) => decimal_decoder!(Decimal256Type, p, s),
1171 DataType::Boolean => Ok(Box::new(BooleanArrayDecoder::new(ctx))),
1172 DataType::Utf8 => Ok(Box::new(StringArrayDecoder::<i32>::new(ctx))),
1173 DataType::Utf8View => Ok(Box::new(StringViewArrayDecoder::new(ctx))),
1174 DataType::LargeUtf8 => Ok(Box::new(StringArrayDecoder::<i64>::new(ctx))),
1175 DataType::List(_) => Ok(Box::new(ListArrayDecoder::<i32>::new(ctx, data_type, is_nullable)?)),
1176 DataType::LargeList(_) => Ok(Box::new(ListArrayDecoder::<i64>::new(ctx, data_type, is_nullable)?)),
1177 DataType::ListView(_) => Ok(Box::new(ListViewArrayDecoder::<i32>::new(ctx, data_type, is_nullable)?)),
1178 DataType::LargeListView(_) => Ok(Box::new(ListViewArrayDecoder::<i64>::new(ctx, data_type, is_nullable)?)),
1179 DataType::FixedSizeList(_, _) => Ok(Box::new(FixedSizeListArrayDecoder::new(ctx, data_type, is_nullable)?)),
1180 DataType::Struct(_) => Ok(Box::new(StructArrayDecoder::new(ctx, data_type, is_nullable)?)),
1181 DataType::Binary => Ok(Box::new(BinaryArrayDecoder::<i32>::default())),
1182 DataType::LargeBinary => Ok(Box::new(BinaryArrayDecoder::<i64>::default())),
1183 DataType::FixedSizeBinary(len) => Ok(Box::new(FixedSizeBinaryArrayDecoder::new(len))),
1184 DataType::BinaryView => Ok(Box::new(BinaryViewDecoder::default())),
1185 DataType::Map(_, _) => Ok(Box::new(MapArrayDecoder::new(ctx, data_type, is_nullable)?)),
1186 DataType::RunEndEncoded(ref r, _) => match r.data_type() {
1187 DataType::Int16 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int16Type>::new(ctx, data_type, is_nullable)?)),
1188 DataType::Int32 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int32Type>::new(ctx, data_type, is_nullable)?)),
1189 DataType::Int64 => Ok(Box::new(RunEndEncodedArrayDecoder::<Int64Type>::new(ctx, data_type, is_nullable)?)),
1190 d => unreachable!("unsupported run end index type: {d}"),
1191 },
1192 _ => Err(ArrowError::NotYetImplemented(format!("Support for {data_type} in JSON reader")))
1193 }
1194}
1195
1196#[cfg(test)]
1197mod tests {
1198 use arrow_array::cast::AsArray;
1199 use arrow_array::{
1200 Array, BooleanArray, Float64Array, GenericListViewArray, Int32Array, ListArray, MapArray,
1201 NullArray, OffsetSizeTrait, StringArray, StringViewArray, StructArray,
1202 };
1203 use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer, ScalarBuffer};
1204 use arrow_cast::display::{ArrayFormatter, FormatOptions};
1205 use arrow_schema::{Field, Fields};
1206 use serde_json::json;
1207 use std::fs::File;
1208 use std::io::{BufReader, Cursor, Seek};
1209 use std::sync::Mutex;
1210
1211 use super::*;
1212
1213 fn do_read(
1214 buf: &str,
1215 batch_size: usize,
1216 coerce_primitive: bool,
1217 strict_mode: bool,
1218 schema: SchemaRef,
1219 ) -> Vec<RecordBatch> {
1220 let config = ReaderBuilder::new(schema)
1221 .with_batch_size(batch_size)
1222 .with_strict_mode(strict_mode)
1223 .with_coerce_primitive(coerce_primitive);
1224 do_read_config(buf, config)
1225 }
1226
1227 fn do_read_config(buf: &str, builder: ReaderBuilder) -> Vec<RecordBatch> {
1228 let mut unbuffered = vec![];
1229
1230 for batch_size in [1, 3, 100, builder.batch_size] {
1232 unbuffered = builder
1233 .clone()
1234 .with_batch_size(batch_size)
1235 .build(Cursor::new(buf.as_bytes()))
1236 .unwrap()
1237 .collect::<Result<Vec<_>, _>>()
1238 .unwrap();
1239
1240 for b in unbuffered.iter().take(unbuffered.len() - 1) {
1241 assert_eq!(b.num_rows(), batch_size)
1242 }
1243
1244 for b in [1, 3, 5] {
1246 let buffered = builder
1247 .clone()
1248 .with_batch_size(batch_size)
1249 .build(BufReader::with_capacity(b, Cursor::new(buf.as_bytes())))
1250 .unwrap()
1251 .collect::<Result<Vec<_>, _>>()
1252 .unwrap();
1253 assert_eq!(unbuffered, buffered);
1254 }
1255 }
1256
1257 unbuffered
1258 }
1259
1260 #[test]
1261 fn test_basic() {
1262 let buf = r#"
1263 {"a": 1, "b": 2, "c": true, "d": 1}
1264 {"a": 2E0, "b": 4, "c": false, "d": 2, "e": 254}
1265
1266 {"b": 6, "a": 2.0, "d": 45}
1267 {"b": "5", "a": 2}
1268 {"b": 4e0}
1269 {"b": 7, "a": null}
1270 "#;
1271
1272 let schema = Arc::new(Schema::new(vec![
1273 Field::new("a", DataType::Int64, true),
1274 Field::new("b", DataType::Int32, true),
1275 Field::new("c", DataType::Boolean, true),
1276 Field::new("d", DataType::Date32, true),
1277 Field::new("e", DataType::Date64, true),
1278 ]));
1279
1280 let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder().unwrap();
1281 assert!(decoder.is_empty());
1282 assert_eq!(decoder.len(), 0);
1283 assert!(!decoder.has_partial_record());
1284 assert_eq!(decoder.decode(buf.as_bytes()).unwrap(), 221);
1285 assert!(!decoder.is_empty());
1286 assert_eq!(decoder.len(), 6);
1287 assert!(!decoder.has_partial_record());
1288 let batch = decoder.flush().unwrap().unwrap();
1289 assert_eq!(batch.num_rows(), 6);
1290 assert!(decoder.is_empty());
1291 assert_eq!(decoder.len(), 0);
1292 assert!(!decoder.has_partial_record());
1293
1294 let batches = do_read(buf, 1024, false, false, schema);
1295 assert_eq!(batches.len(), 1);
1296
1297 let col1 = batches[0].column(0).as_primitive::<Int64Type>();
1298 assert_eq!(col1.null_count(), 2);
1299 assert_eq!(col1.values(), &[1, 2, 2, 2, 0, 0]);
1300 assert!(col1.is_null(4));
1301 assert!(col1.is_null(5));
1302
1303 let col2 = batches[0].column(1).as_primitive::<Int32Type>();
1304 assert_eq!(col2.null_count(), 0);
1305 assert_eq!(col2.values(), &[2, 4, 6, 5, 4, 7]);
1306
1307 let col3 = batches[0].column(2).as_boolean();
1308 assert_eq!(col3.null_count(), 4);
1309 assert!(col3.value(0));
1310 assert!(!col3.is_null(0));
1311 assert!(!col3.value(1));
1312 assert!(!col3.is_null(1));
1313
1314 let col4 = batches[0].column(3).as_primitive::<Date32Type>();
1315 assert_eq!(col4.null_count(), 3);
1316 assert!(col4.is_null(3));
1317 assert_eq!(col4.values(), &[1, 2, 45, 0, 0, 0]);
1318
1319 let col5 = batches[0].column(4).as_primitive::<Date64Type>();
1320 assert_eq!(col5.null_count(), 5);
1321 assert!(col5.is_null(0));
1322 assert!(col5.is_null(2));
1323 assert!(col5.is_null(3));
1324 assert_eq!(col5.values(), &[0, 254, 0, 0, 0, 0]);
1325 }
1326
1327 #[test]
1328 fn test_string() {
1329 let buf = r#"
1330 {"a": "1", "b": "2"}
1331 {"a": "hello", "b": "shoo"}
1332 {"b": "\t😁foo", "a": "\nfoobar\ud83d\ude00\u0061\u0073\u0066\u0067\u00FF"}
1333
1334 {"b": null}
1335 {"b": "", "a": null}
1336
1337 "#;
1338 let schema = Arc::new(Schema::new(vec![
1339 Field::new("a", DataType::Utf8, true),
1340 Field::new("b", DataType::LargeUtf8, true),
1341 ]));
1342
1343 let batches = do_read(buf, 1024, false, false, schema);
1344 assert_eq!(batches.len(), 1);
1345
1346 let col1 = batches[0].column(0).as_string::<i32>();
1347 assert_eq!(col1.null_count(), 2);
1348 assert_eq!(col1.value(0), "1");
1349 assert_eq!(col1.value(1), "hello");
1350 assert_eq!(col1.value(2), "\nfoobar😀asfgÿ");
1351 assert!(col1.is_null(3));
1352 assert!(col1.is_null(4));
1353
1354 let col2 = batches[0].column(1).as_string::<i64>();
1355 assert_eq!(col2.null_count(), 1);
1356 assert_eq!(col2.value(0), "2");
1357 assert_eq!(col2.value(1), "shoo");
1358 assert_eq!(col2.value(2), "\t😁foo");
1359 assert!(col2.is_null(3));
1360 assert_eq!(col2.value(4), "");
1361 }
1362
1363 #[test]
1364 fn test_long_string_view_allocation() {
1365 let expected_capacity: usize = 41;
1375
1376 let buf = r#"
1377 {"a": "short", "b": "dummy"}
1378 {"a": "this is definitely long", "b": "dummy"}
1379 {"a": "hello", "b": "dummy"}
1380 {"a": "\nfoobar😀asfgÿ", "b": "dummy"}
1381 "#;
1382
1383 let schema = Arc::new(Schema::new(vec![
1384 Field::new("a", DataType::Utf8View, true),
1385 Field::new("b", DataType::LargeUtf8, true),
1386 ]));
1387
1388 let batches = do_read(buf, 1024, false, false, schema);
1389 assert_eq!(batches.len(), 1, "Expected one record batch");
1390
1391 let col_a = batches[0].column(0);
1393 let string_view_array = col_a
1394 .as_any()
1395 .downcast_ref::<StringViewArray>()
1396 .expect("Column should be a StringViewArray");
1397
1398 let data_buffer = string_view_array.to_data().buffers()[0].len();
1401
1402 assert!(
1405 data_buffer >= expected_capacity,
1406 "Data buffer length ({data_buffer}) should be at least {expected_capacity}",
1407 );
1408
1409 assert_eq!(string_view_array.value(0), "short");
1411 assert_eq!(string_view_array.value(1), "this is definitely long");
1412 assert_eq!(string_view_array.value(2), "hello");
1413 assert_eq!(string_view_array.value(3), "\nfoobar😀asfgÿ");
1414 }
1415
1416 #[test]
1418 fn test_numeric_view_allocation() {
1419 let expected_capacity: usize = 33;
1427
1428 let buf = r#"
1429 {"n": 123456789}
1430 {"n": 1000000000000}
1431 {"n": 3.1415}
1432 {"n": 2.718281828459045}
1433 "#;
1434
1435 let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Utf8View, true)]));
1436
1437 let batches = do_read(buf, 1024, true, false, schema);
1438 assert_eq!(batches.len(), 1, "Expected one record batch");
1439
1440 let col_n = batches[0].column(0);
1441 let string_view_array = col_n
1442 .as_any()
1443 .downcast_ref::<StringViewArray>()
1444 .expect("Column should be a StringViewArray");
1445
1446 let data_buffer = string_view_array.to_data().buffers()[0].len();
1448 assert!(
1449 data_buffer >= expected_capacity,
1450 "Data buffer length ({data_buffer}) should be at least {expected_capacity}",
1451 );
1452
1453 assert_eq!(string_view_array.value(0), "123456789");
1456 assert_eq!(string_view_array.value(1), "1000000000000");
1457 assert_eq!(string_view_array.value(2), "3.1415");
1458 assert_eq!(string_view_array.value(3), "2.718281828459045");
1459 }
1460
1461 #[test]
1462 fn test_string_with_uft8view() {
1463 let buf = r#"
1464 {"a": "1", "b": "2"}
1465 {"a": "hello", "b": "shoo"}
1466 {"b": "\t😁foo", "a": "\nfoobar\ud83d\ude00\u0061\u0073\u0066\u0067\u00FF"}
1467
1468 {"b": null}
1469 {"b": "", "a": null}
1470
1471 "#;
1472 let schema = Arc::new(Schema::new(vec![
1473 Field::new("a", DataType::Utf8View, true),
1474 Field::new("b", DataType::LargeUtf8, true),
1475 ]));
1476
1477 let batches = do_read(buf, 1024, false, false, schema);
1478 assert_eq!(batches.len(), 1);
1479
1480 let col1 = batches[0].column(0).as_string_view();
1481 assert_eq!(col1.null_count(), 2);
1482 assert_eq!(col1.value(0), "1");
1483 assert_eq!(col1.value(1), "hello");
1484 assert_eq!(col1.value(2), "\nfoobar😀asfgÿ");
1485 assert!(col1.is_null(3));
1486 assert!(col1.is_null(4));
1487 assert_eq!(col1.data_type(), &DataType::Utf8View);
1488
1489 let col2 = batches[0].column(1).as_string::<i64>();
1490 assert_eq!(col2.null_count(), 1);
1491 assert_eq!(col2.value(0), "2");
1492 assert_eq!(col2.value(1), "shoo");
1493 assert_eq!(col2.value(2), "\t😁foo");
1494 assert!(col2.is_null(3));
1495 assert_eq!(col2.value(4), "");
1496 }
1497
1498 #[test]
1499 fn test_complex() {
1500 let buf = r#"
1501 {"list": [], "nested": {"a": 1, "b": 2}, "nested_list": {"list2": [{"c": 3}, {"c": 4}]}}
1502 {"list": [5, 6], "nested": {"a": 7}, "nested_list": {"list2": []}}
1503 {"list": null, "nested": {"a": null}}
1504 "#;
1505
1506 let schema = Arc::new(Schema::new(vec![
1507 Field::new_list("list", Field::new("element", DataType::Int32, false), true),
1508 Field::new_struct(
1509 "nested",
1510 vec![
1511 Field::new("a", DataType::Int32, true),
1512 Field::new("b", DataType::Int32, true),
1513 ],
1514 true,
1515 ),
1516 Field::new_struct(
1517 "nested_list",
1518 vec![Field::new_list(
1519 "list2",
1520 Field::new_struct(
1521 "element",
1522 vec![Field::new("c", DataType::Int32, false)],
1523 false,
1524 ),
1525 true,
1526 )],
1527 true,
1528 ),
1529 ]));
1530
1531 let batches = do_read(buf, 1024, false, false, schema);
1532 assert_eq!(batches.len(), 1);
1533
1534 let list = batches[0].column(0).as_list::<i32>();
1535 assert_eq!(list.len(), 3);
1536 assert_eq!(list.value_offsets(), &[0, 0, 2, 2]);
1537 assert_eq!(list.null_count(), 1);
1538 assert!(list.is_null(2));
1539 let list_values = list.values().as_primitive::<Int32Type>();
1540 assert_eq!(list_values.values(), &[5, 6]);
1541
1542 let nested = batches[0].column(1).as_struct();
1543 let a = nested.column(0).as_primitive::<Int32Type>();
1544 assert_eq!(list.null_count(), 1);
1545 assert_eq!(a.values(), &[1, 7, 0]);
1546 assert!(list.is_null(2));
1547
1548 let b = nested.column(1).as_primitive::<Int32Type>();
1549 assert_eq!(b.null_count(), 2);
1550 assert_eq!(b.len(), 3);
1551 assert_eq!(b.value(0), 2);
1552 assert!(b.is_null(1));
1553 assert!(b.is_null(2));
1554
1555 let nested_list = batches[0].column(2).as_struct();
1556 assert_eq!(nested_list.len(), 3);
1557 assert_eq!(nested_list.null_count(), 1);
1558 assert!(nested_list.is_null(2));
1559
1560 let list2 = nested_list.column(0).as_list::<i32>();
1561 assert_eq!(list2.len(), 3);
1562 assert_eq!(list2.null_count(), 1);
1563 assert_eq!(list2.value_offsets(), &[0, 2, 2, 2]);
1564 assert!(list2.is_null(2));
1565
1566 let list2_values = list2.values().as_struct();
1567
1568 let c = list2_values.column(0).as_primitive::<Int32Type>();
1569 assert_eq!(c.values(), &[3, 4]);
1570 }
1571
1572 #[test]
1573 fn test_projection() {
1574 let buf = r#"
1575 {"list": [], "nested": {"a": 1, "b": 2}, "nested_list": {"list2": [{"c": 3, "d": 5}, {"c": 4}]}}
1576 {"list": [5, 6], "nested": {"a": 7}, "nested_list": {"list2": []}}
1577 "#;
1578
1579 let schema = Arc::new(Schema::new(vec![
1580 Field::new_struct(
1581 "nested",
1582 vec![Field::new("a", DataType::Int32, false)],
1583 true,
1584 ),
1585 Field::new_struct(
1586 "nested_list",
1587 vec![Field::new_list(
1588 "list2",
1589 Field::new_struct(
1590 "element",
1591 vec![Field::new("d", DataType::Int32, true)],
1592 false,
1593 ),
1594 true,
1595 )],
1596 true,
1597 ),
1598 ]));
1599
1600 let batches = do_read(buf, 1024, false, false, schema);
1601 assert_eq!(batches.len(), 1);
1602
1603 let nested = batches[0].column(0).as_struct();
1604 assert_eq!(nested.num_columns(), 1);
1605 let a = nested.column(0).as_primitive::<Int32Type>();
1606 assert_eq!(a.null_count(), 0);
1607 assert_eq!(a.values(), &[1, 7]);
1608
1609 let nested_list = batches[0].column(1).as_struct();
1610 assert_eq!(nested_list.num_columns(), 1);
1611 assert_eq!(nested_list.null_count(), 0);
1612
1613 let list2 = nested_list.column(0).as_list::<i32>();
1614 assert_eq!(list2.value_offsets(), &[0, 2, 2]);
1615 assert_eq!(list2.null_count(), 0);
1616
1617 let child = list2.values().as_struct();
1618 assert_eq!(child.num_columns(), 1);
1619 assert_eq!(child.len(), 2);
1620 assert_eq!(child.null_count(), 0);
1621
1622 let c = child.column(0).as_primitive::<Int32Type>();
1623 assert_eq!(c.values(), &[5, 0]);
1624 assert_eq!(c.null_count(), 1);
1625 assert!(c.is_null(1));
1626 }
1627
1628 #[test]
1629 fn test_map() {
1630 let buf = r#"
1631 {"map": {"a": ["foo", null]}}
1632 {"map": {"a": [null], "b": []}}
1633 {"map": {"c": null, "a": ["baz"]}}
1634 "#;
1635 let map = Field::new_map(
1636 "map",
1637 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1638 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1639 Field::new_list(
1640 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
1641 Field::new("element", DataType::Utf8, true),
1642 true,
1643 ),
1644 false,
1645 true,
1646 );
1647
1648 let schema = Arc::new(Schema::new(vec![map]));
1649
1650 let batches = do_read(buf, 1024, false, false, schema);
1651 assert_eq!(batches.len(), 1);
1652
1653 let map = batches[0].column(0).as_map();
1654 let map_keys = map.keys().as_string::<i32>();
1655 let map_values = map.values().as_list::<i32>();
1656 assert_eq!(map.value_offsets(), &[0, 1, 3, 5]);
1657
1658 let k: Vec<_> = map_keys.iter().flatten().collect();
1659 assert_eq!(&k, &["a", "a", "b", "c", "a"]);
1660
1661 let list_values = map_values.values().as_string::<i32>();
1662 let lv: Vec<_> = list_values.iter().collect();
1663 assert_eq!(&lv, &[Some("foo"), None, None, Some("baz")]);
1664 assert_eq!(map_values.value_offsets(), &[0, 2, 3, 3, 3, 4]);
1665 assert_eq!(map_values.null_count(), 1);
1666 assert!(map_values.is_null(3));
1667
1668 let options = FormatOptions::default().with_null("null");
1669 let formatter = ArrayFormatter::try_new(map, &options).unwrap();
1670 assert_eq!(formatter.value(0).to_string(), "{a: [foo, null]}");
1671 assert_eq!(formatter.value(1).to_string(), "{a: [null], b: []}");
1672 assert_eq!(formatter.value(2).to_string(), "{c: null, a: [baz]}");
1673 }
1674
1675 #[test]
1676 fn test_map_non_nullable_value() {
1677 let map = Field::new_map(
1678 "map",
1679 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1680 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1681 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1682 false,
1683 false,
1684 );
1685 let schema = Arc::new(Schema::new(vec![map]));
1686 let buf = r#"{"map": {"key": null}}"#;
1687
1688 let err = ReaderBuilder::new(schema)
1689 .build(Cursor::new(buf.as_bytes()))
1690 .unwrap()
1691 .read()
1692 .unwrap_err();
1693
1694 assert_eq!(
1695 err.to_string(),
1696 "Invalid argument error: Found unmasked nulls for non-nullable StructArray field \"value\""
1697 );
1698 }
1699
1700 #[test]
1701 fn test_not_coercing_primitive_into_string_without_flag() {
1702 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)]));
1703
1704 let buf = r#"{"a": 1}"#;
1705 let err = ReaderBuilder::new(schema.clone())
1706 .with_batch_size(1024)
1707 .build(Cursor::new(buf.as_bytes()))
1708 .unwrap()
1709 .read()
1710 .unwrap_err();
1711
1712 assert_eq!(
1713 err.to_string(),
1714 "Json error: whilst decoding field 'a': expected string got 1"
1715 );
1716
1717 let buf = r#"{"a": true}"#;
1718 let err = ReaderBuilder::new(schema)
1719 .with_batch_size(1024)
1720 .build(Cursor::new(buf.as_bytes()))
1721 .unwrap()
1722 .read()
1723 .unwrap_err();
1724
1725 assert_eq!(
1726 err.to_string(),
1727 "Json error: whilst decoding field 'a': expected string got true"
1728 );
1729 }
1730
1731 #[test]
1732 fn test_coercing_primitive_into_string() {
1733 let buf = r#"
1734 {"a": 1, "b": 2, "c": true}
1735 {"a": 2E0, "b": 4, "c": false}
1736
1737 {"b": 6, "a": 2.0}
1738 {"b": "5", "a": 2}
1739 {"b": 4e0}
1740 {"b": 7, "a": null}
1741 "#;
1742
1743 let schema = Arc::new(Schema::new(vec![
1744 Field::new("a", DataType::Utf8, true),
1745 Field::new("b", DataType::Utf8, true),
1746 Field::new("c", DataType::Utf8, true),
1747 ]));
1748
1749 let batches = do_read(buf, 1024, true, false, schema);
1750 assert_eq!(batches.len(), 1);
1751
1752 let col1 = batches[0].column(0).as_string::<i32>();
1753 assert_eq!(col1.null_count(), 2);
1754 assert_eq!(col1.value(0), "1");
1755 assert_eq!(col1.value(1), "2E0");
1756 assert_eq!(col1.value(2), "2.0");
1757 assert_eq!(col1.value(3), "2");
1758 assert!(col1.is_null(4));
1759 assert!(col1.is_null(5));
1760
1761 let col2 = batches[0].column(1).as_string::<i32>();
1762 assert_eq!(col2.null_count(), 0);
1763 assert_eq!(col2.value(0), "2");
1764 assert_eq!(col2.value(1), "4");
1765 assert_eq!(col2.value(2), "6");
1766 assert_eq!(col2.value(3), "5");
1767 assert_eq!(col2.value(4), "4e0");
1768 assert_eq!(col2.value(5), "7");
1769
1770 let col3 = batches[0].column(2).as_string::<i32>();
1771 assert_eq!(col3.null_count(), 4);
1772 assert_eq!(col3.value(0), "true");
1773 assert_eq!(col3.value(1), "false");
1774 assert!(col3.is_null(2));
1775 assert!(col3.is_null(3));
1776 assert!(col3.is_null(4));
1777 assert!(col3.is_null(5));
1778 }
1779
1780 fn test_decimal<T: DecimalType>(data_type: DataType) {
1781 let buf = r#"
1782 {"a": 1, "b": 2, "c": 38.30}
1783 {"a": 2, "b": 4, "c": 123.456}
1784
1785 {"b": 1337, "a": "2.0452"}
1786 {"b": "5", "a": "11034.2"}
1787 {"b": 40}
1788 {"b": 1234, "a": null}
1789 "#;
1790
1791 let schema = Arc::new(Schema::new(vec![
1792 Field::new("a", data_type.clone(), true),
1793 Field::new("b", data_type.clone(), true),
1794 Field::new("c", data_type, true),
1795 ]));
1796
1797 let batches = do_read(buf, 1024, true, false, schema);
1798 assert_eq!(batches.len(), 1);
1799
1800 let col1 = batches[0].column(0).as_primitive::<T>();
1801 assert_eq!(col1.null_count(), 2);
1802 assert!(col1.is_null(4));
1803 assert!(col1.is_null(5));
1804 assert_eq!(
1805 col1.values(),
1806 &[100, 200, 205, 1103420, 0, 0].map(T::Native::usize_as)
1807 );
1808
1809 let col2 = batches[0].column(1).as_primitive::<T>();
1810 assert_eq!(col2.null_count(), 0);
1811 assert_eq!(
1812 col2.values(),
1813 &[200, 400, 133700, 500, 4000, 123400].map(T::Native::usize_as)
1814 );
1815
1816 let col3 = batches[0].column(2).as_primitive::<T>();
1817 assert_eq!(col3.null_count(), 4);
1818 assert!(!col3.is_null(0));
1819 assert!(!col3.is_null(1));
1820 assert!(col3.is_null(2));
1821 assert!(col3.is_null(3));
1822 assert!(col3.is_null(4));
1823 assert!(col3.is_null(5));
1824 assert_eq!(
1825 col3.values(),
1826 &[3830, 12346, 0, 0, 0, 0].map(T::Native::usize_as)
1827 );
1828 }
1829
1830 #[test]
1831 fn test_decimal_number_formats() {
1832 let buf = r#"
1836 {"a": 0e0, "b": " 1.5 ", "c": 1234.5}
1837 {"a": 1.5E2, "b": "-0.005", "c": -150}
1838 {"a": 1e-3, "b": "1e2", "c": 1E2}
1839 {"a": -0E+0, "b": "+.5", "c": 5}
1840 "#;
1841 let schema = Arc::new(Schema::new(vec![
1842 Field::new("a", DataType::Decimal128(10, 2), true),
1843 Field::new("b", DataType::Decimal64(18, 2), true),
1844 Field::new("c", DataType::Decimal128(10, -2), true),
1845 ]));
1846 let batches = do_read(buf, 1024, true, false, schema);
1847 assert_eq!(batches.len(), 1);
1848 assert_eq!(
1849 batches[0]
1850 .column(0)
1851 .as_primitive::<Decimal128Type>()
1852 .values(),
1853 &[0, 15000, 0, 0]
1854 );
1855 assert_eq!(
1856 batches[0]
1857 .column(1)
1858 .as_primitive::<Decimal64Type>()
1859 .values(),
1860 &[150, -1, 10000, 50]
1861 );
1862 assert_eq!(
1863 batches[0]
1864 .column(2)
1865 .as_primitive::<Decimal128Type>()
1866 .values(),
1867 &[12, -2, 1, 0]
1868 );
1869
1870 let schema = Arc::new(Schema::new(vec![Field::new(
1873 "a",
1874 DataType::Decimal128(5, 2),
1875 true,
1876 )]));
1877 let long = "1".repeat(300);
1878 for (buf, expected) in [
1879 (
1880 r#"{"a": "abc"}"#.to_string(),
1881 "Invalid decimal format: \"abc\"",
1882 ),
1883 (
1884 r#"{"a": 123456789}"#.to_string(),
1885 "does not fit in Decimal128(5, 2)",
1886 ),
1887 (
1888 r#"{"a": 1e99999}"#.to_string(),
1889 "does not fit in Decimal128(5, 2)",
1890 ),
1891 (
1892 format!(r#"{{"a": {long}}}"#),
1893 "does not fit in Decimal128(5, 2)",
1894 ),
1895 (
1896 r#"{"a": 4825037936439135476.2609835314269495255615E-14}"#.to_string(),
1897 "does not fit in Decimal128(5, 2)",
1898 ),
1899 ] {
1900 let err = ReaderBuilder::new(schema.clone())
1901 .build(Cursor::new(buf.as_bytes()))
1902 .unwrap()
1903 .next()
1904 .unwrap()
1905 .unwrap_err()
1906 .to_string();
1907 assert!(err.contains(expected), "{buf}: {err}");
1908
1909 let batch = ReaderBuilder::new(schema.clone())
1910 .with_ignore_type_conflicts(true)
1911 .build(Cursor::new(buf.as_bytes()))
1912 .unwrap()
1913 .next()
1914 .unwrap()
1915 .unwrap();
1916 assert!(batch.column(0).is_null(0), "{buf}");
1917 }
1918 }
1919
1920 #[test]
1921 fn test_decimals() {
1922 test_decimal::<Decimal32Type>(DataType::Decimal32(8, 2));
1923 test_decimal::<Decimal64Type>(DataType::Decimal64(10, 2));
1924 test_decimal::<Decimal128Type>(DataType::Decimal128(10, 2));
1925 test_decimal::<Decimal256Type>(DataType::Decimal256(10, 2));
1926 }
1927
1928 fn test_timestamp<T: ArrowTimestampType>() {
1929 let buf = r#"
1930 {"a": 1, "b": "2020-09-08T13:42:29.190855+00:00", "c": 38.30, "d": "1997-01-31T09:26:56.123"}
1931 {"a": 2, "b": "2020-09-08T13:42:29.190855Z", "c": 123.456, "d": 123.456}
1932
1933 {"b": 1337, "b": "2020-09-08T13:42:29Z", "c": "1997-01-31T09:26:56.123", "d": "1997-01-31T09:26:56.123Z"}
1934 {"b": 40, "c": "2020-09-08T13:42:29.190855+00:00", "d": "1997-01-31 09:26:56.123-05:00"}
1935 {"b": 1234, "a": null, "c": "1997-01-31 09:26:56.123Z", "d": "1997-01-31 092656"}
1936 {"c": "1997-01-31T14:26:56.123-05:00", "d": "1997-01-31"}
1937 "#;
1938
1939 let with_timezone = DataType::Timestamp(T::UNIT, Some("+08:00".into()));
1940 let schema = Arc::new(Schema::new(vec![
1941 Field::new("a", T::DATA_TYPE, true),
1942 Field::new("b", T::DATA_TYPE, true),
1943 Field::new("c", T::DATA_TYPE, true),
1944 Field::new("d", with_timezone, true),
1945 ]));
1946
1947 let batches = do_read(buf, 1024, true, false, schema);
1948 assert_eq!(batches.len(), 1);
1949
1950 let unit_in_nanos: i64 = match T::UNIT {
1951 TimeUnit::Second => 1_000_000_000,
1952 TimeUnit::Millisecond => 1_000_000,
1953 TimeUnit::Microsecond => 1_000,
1954 TimeUnit::Nanosecond => 1,
1955 };
1956
1957 let col1 = batches[0].column(0).as_primitive::<T>();
1958 assert_eq!(col1.null_count(), 4);
1959 assert!(col1.is_null(2));
1960 assert!(col1.is_null(3));
1961 assert!(col1.is_null(4));
1962 assert!(col1.is_null(5));
1963 assert_eq!(col1.values(), &[1, 2, 0, 0, 0, 0].map(T::Native::usize_as));
1964
1965 let col2 = batches[0].column(1).as_primitive::<T>();
1966 assert_eq!(col2.null_count(), 1);
1967 assert!(col2.is_null(5));
1968 assert_eq!(
1969 col2.values(),
1970 &[
1971 1599572549190855000 / unit_in_nanos,
1972 1599572549190855000 / unit_in_nanos,
1973 1599572549000000000 / unit_in_nanos,
1974 40,
1975 1234,
1976 0
1977 ]
1978 );
1979
1980 let col3 = batches[0].column(2).as_primitive::<T>();
1981 assert_eq!(col3.null_count(), 0);
1982 assert_eq!(
1983 col3.values(),
1984 &[
1985 38,
1986 123,
1987 854702816123000000 / unit_in_nanos,
1988 1599572549190855000 / unit_in_nanos,
1989 854702816123000000 / unit_in_nanos,
1990 854738816123000000 / unit_in_nanos
1991 ]
1992 );
1993
1994 let col4 = batches[0].column(3).as_primitive::<T>();
1995
1996 assert_eq!(col4.null_count(), 0);
1997 assert_eq!(
1998 col4.values(),
1999 &[
2000 854674016123000000 / unit_in_nanos,
2001 123,
2002 854702816123000000 / unit_in_nanos,
2003 854720816123000000 / unit_in_nanos,
2004 854674016000000000 / unit_in_nanos,
2005 854640000000000000 / unit_in_nanos
2006 ]
2007 );
2008 }
2009
2010 #[test]
2011 #[cfg_attr(miri, ignore)] fn test_timestamps() {
2013 test_timestamp::<TimestampSecondType>();
2014 test_timestamp::<TimestampMillisecondType>();
2015 test_timestamp::<TimestampMicrosecondType>();
2016 test_timestamp::<TimestampNanosecondType>();
2017 }
2018
2019 fn test_time<T: ArrowTemporalType>() {
2020 let buf = r#"
2021 {"a": 1, "b": "09:26:56.123 AM", "c": 38.30}
2022 {"a": 2, "b": "23:59:59", "c": 123.456}
2023
2024 {"b": 1337, "b": "6:00 pm", "c": "09:26:56.123"}
2025 {"b": 40, "c": "13:42:29.190855"}
2026 {"b": 1234, "a": null, "c": "09:26:56.123"}
2027 {"c": "14:26:56.123"}
2028 "#;
2029
2030 let (DataType::Time32(unit) | DataType::Time64(unit)) = T::DATA_TYPE else {
2031 unreachable!()
2032 };
2033
2034 let unit_in_nanos = match unit {
2035 TimeUnit::Second => 1_000_000_000,
2036 TimeUnit::Millisecond => 1_000_000,
2037 TimeUnit::Microsecond => 1_000,
2038 TimeUnit::Nanosecond => 1,
2039 };
2040
2041 let schema = Arc::new(Schema::new(vec![
2042 Field::new("a", T::DATA_TYPE, true),
2043 Field::new("b", T::DATA_TYPE, true),
2044 Field::new("c", T::DATA_TYPE, true),
2045 ]));
2046
2047 let batches = do_read(buf, 1024, true, false, schema);
2048 assert_eq!(batches.len(), 1);
2049
2050 let col1 = batches[0].column(0).as_primitive::<T>();
2051 assert_eq!(col1.null_count(), 4);
2052 assert!(col1.is_null(2));
2053 assert!(col1.is_null(3));
2054 assert!(col1.is_null(4));
2055 assert!(col1.is_null(5));
2056 assert_eq!(col1.values(), &[1, 2, 0, 0, 0, 0].map(T::Native::usize_as));
2057
2058 let col2 = batches[0].column(1).as_primitive::<T>();
2059 assert_eq!(col2.null_count(), 1);
2060 assert!(col2.is_null(5));
2061 assert_eq!(
2062 col2.values(),
2063 &[
2064 34016123000000 / unit_in_nanos,
2065 86399000000000 / unit_in_nanos,
2066 64800000000000 / unit_in_nanos,
2067 40,
2068 1234,
2069 0
2070 ]
2071 .map(T::Native::usize_as)
2072 );
2073
2074 let col3 = batches[0].column(2).as_primitive::<T>();
2075 assert_eq!(col3.null_count(), 0);
2076 assert_eq!(
2077 col3.values(),
2078 &[
2079 38,
2080 123,
2081 34016123000000 / unit_in_nanos,
2082 49349190855000 / unit_in_nanos,
2083 34016123000000 / unit_in_nanos,
2084 52016123000000 / unit_in_nanos
2085 ]
2086 .map(T::Native::usize_as)
2087 );
2088 }
2089
2090 #[test]
2091 fn test_times() {
2092 test_time::<Time32MillisecondType>();
2093 test_time::<Time32SecondType>();
2094 test_time::<Time64MicrosecondType>();
2095 test_time::<Time64NanosecondType>();
2096 }
2097
2098 fn test_duration<T: ArrowTemporalType>() {
2099 let buf = r#"
2100 {"a": 1, "b": "2"}
2101 {"a": 3, "b": null}
2102 "#;
2103
2104 let schema = Arc::new(Schema::new(vec![
2105 Field::new("a", T::DATA_TYPE, true),
2106 Field::new("b", T::DATA_TYPE, true),
2107 ]));
2108
2109 let batches = do_read(buf, 1024, true, false, schema);
2110 assert_eq!(batches.len(), 1);
2111
2112 let col_a = batches[0].column_by_name("a").unwrap().as_primitive::<T>();
2113 assert_eq!(col_a.null_count(), 0);
2114 assert_eq!(col_a.values(), &[1, 3].map(T::Native::usize_as));
2115
2116 let col2 = batches[0].column_by_name("b").unwrap().as_primitive::<T>();
2117 assert_eq!(col2.null_count(), 1);
2118 assert_eq!(col2.values(), &[2, 0].map(T::Native::usize_as));
2119 }
2120
2121 #[test]
2122 fn test_durations() {
2123 test_duration::<DurationNanosecondType>();
2124 test_duration::<DurationMicrosecondType>();
2125 test_duration::<DurationMillisecondType>();
2126 test_duration::<DurationSecondType>();
2127 }
2128
2129 #[test]
2130 fn test_delta_checkpoint() {
2131 let json = "{\"protocol\":{\"minReaderVersion\":1,\"minWriterVersion\":2}}";
2132 let schema = Arc::new(Schema::new(vec![
2133 Field::new_struct(
2134 "protocol",
2135 vec![
2136 Field::new("minReaderVersion", DataType::Int32, true),
2137 Field::new("minWriterVersion", DataType::Int32, true),
2138 ],
2139 true,
2140 ),
2141 Field::new_struct(
2142 "add",
2143 vec![Field::new_map(
2144 "partitionValues",
2145 "key_value",
2146 Field::new("key", DataType::Utf8, false),
2147 Field::new("value", DataType::Utf8, true),
2148 false,
2149 false,
2150 )],
2151 true,
2152 ),
2153 ]));
2154
2155 let batches = do_read(json, 1024, true, false, schema);
2156 assert_eq!(batches.len(), 1);
2157
2158 let s: StructArray = batches.into_iter().next().unwrap().into();
2159 let opts = FormatOptions::default().with_null("null");
2160 let formatter = ArrayFormatter::try_new(&s, &opts).unwrap();
2161 assert_eq!(
2162 formatter.value(0).to_string(),
2163 "{protocol: {minReaderVersion: 1, minWriterVersion: 2}, add: null}"
2164 );
2165 }
2166
2167 #[test]
2168 fn struct_nullability() {
2169 let do_test = |child: DataType| {
2170 let non_null = r#"{"foo": {}}"#;
2172 let schema = Arc::new(Schema::new(vec![Field::new_struct(
2173 "foo",
2174 vec![Field::new("bar", child, false)],
2175 true,
2176 )]));
2177 let mut reader = ReaderBuilder::new(schema.clone())
2178 .build(Cursor::new(non_null.as_bytes()))
2179 .unwrap();
2180 assert!(reader.next().unwrap().is_err()); let null = r#"{"foo": {bar: null}}"#;
2183 let mut reader = ReaderBuilder::new(schema.clone())
2184 .build(Cursor::new(null.as_bytes()))
2185 .unwrap();
2186 assert!(reader.next().unwrap().is_err()); let null = r#"{"foo": null}"#;
2190 let mut reader = ReaderBuilder::new(schema)
2191 .build(Cursor::new(null.as_bytes()))
2192 .unwrap();
2193 let batch = reader.next().unwrap().unwrap();
2194 assert_eq!(batch.num_columns(), 1);
2195 let foo = batch.column(0).as_struct();
2196 assert_eq!(foo.len(), 1);
2197 assert!(foo.is_null(0));
2198 assert_eq!(foo.num_columns(), 1);
2199
2200 let bar = foo.column(0);
2201 assert_eq!(bar.len(), 1);
2202 assert!(bar.is_null(0));
2204 };
2205
2206 do_test(DataType::Boolean);
2207 do_test(DataType::Int32);
2208 do_test(DataType::Utf8);
2209 do_test(DataType::Decimal128(2, 1));
2210 do_test(DataType::Timestamp(
2211 TimeUnit::Microsecond,
2212 Some("+00:00".into()),
2213 ));
2214 }
2215
2216 #[test]
2217 fn test_truncation() {
2218 let buf = r#"
2219 {"i64": 9223372036854775807, "u64": 18446744073709551615 }
2220 {"i64": "9223372036854775807", "u64": "18446744073709551615" }
2221 {"i64": -9223372036854775808, "u64": 0 }
2222 {"i64": "-9223372036854775808", "u64": 0 }
2223 "#;
2224
2225 let schema = Arc::new(Schema::new(vec![
2226 Field::new("i64", DataType::Int64, true),
2227 Field::new("u64", DataType::UInt64, true),
2228 ]));
2229
2230 let batches = do_read(buf, 1024, true, false, schema);
2231 assert_eq!(batches.len(), 1);
2232
2233 let i64 = batches[0].column(0).as_primitive::<Int64Type>();
2234 assert_eq!(i64.values(), &[i64::MAX, i64::MAX, i64::MIN, i64::MIN]);
2235
2236 let u64 = batches[0].column(1).as_primitive::<UInt64Type>();
2237 assert_eq!(u64.values(), &[u64::MAX, u64::MAX, u64::MIN, u64::MIN]);
2238 }
2239
2240 #[test]
2241 fn test_timestamp_truncation() {
2242 let buf = r#"
2243 {"time": 9223372036854775807 }
2244 {"time": -9223372036854775808 }
2245 {"time": 9e5 }
2246 "#;
2247
2248 let schema = Arc::new(Schema::new(vec![Field::new(
2249 "time",
2250 DataType::Timestamp(TimeUnit::Nanosecond, None),
2251 true,
2252 )]));
2253
2254 let batches = do_read(buf, 1024, true, false, schema);
2255 assert_eq!(batches.len(), 1);
2256
2257 let i64 = batches[0]
2258 .column(0)
2259 .as_primitive::<TimestampNanosecondType>();
2260 assert_eq!(i64.values(), &[i64::MAX, i64::MIN, 900000]);
2261 }
2262
2263 #[test]
2264 fn test_strict_mode_no_missing_columns_in_schema() {
2265 let buf = r#"
2266 {"a": 1, "b": "2", "c": true}
2267 {"a": 2E0, "b": "4", "c": false}
2268 "#;
2269
2270 let schema = Arc::new(Schema::new(vec![
2271 Field::new("a", DataType::Int16, false),
2272 Field::new("b", DataType::Utf8, false),
2273 Field::new("c", DataType::Boolean, false),
2274 ]));
2275
2276 let batches = do_read(buf, 1024, true, true, schema);
2277 assert_eq!(batches.len(), 1);
2278
2279 let buf = r#"
2280 {"a": 1, "b": "2", "c": {"a": true, "b": 1}}
2281 {"a": 2E0, "b": "4", "c": {"a": false, "b": 2}}
2282 "#;
2283
2284 let schema = Arc::new(Schema::new(vec![
2285 Field::new("a", DataType::Int16, false),
2286 Field::new("b", DataType::Utf8, false),
2287 Field::new_struct(
2288 "c",
2289 vec![
2290 Field::new("a", DataType::Boolean, false),
2291 Field::new("b", DataType::Int16, false),
2292 ],
2293 false,
2294 ),
2295 ]));
2296
2297 let batches = do_read(buf, 1024, true, true, schema);
2298 assert_eq!(batches.len(), 1);
2299 }
2300
2301 #[test]
2302 fn test_strict_mode_missing_columns_in_schema() {
2303 let buf = r#"
2304 {"a": 1, "b": "2", "c": true}
2305 {"a": 2E0, "b": "4", "c": false}
2306 "#;
2307
2308 let schema = Arc::new(Schema::new(vec![
2309 Field::new("a", DataType::Int16, true),
2310 Field::new("c", DataType::Boolean, true),
2311 ]));
2312
2313 let err = ReaderBuilder::new(schema)
2314 .with_batch_size(1024)
2315 .with_strict_mode(true)
2316 .build(Cursor::new(buf.as_bytes()))
2317 .unwrap()
2318 .read()
2319 .unwrap_err();
2320
2321 assert_eq!(
2322 err.to_string(),
2323 "Json error: column 'b' missing from schema"
2324 );
2325
2326 let buf = r#"
2327 {"a": 1, "b": "2", "c": {"a": true, "b": 1}}
2328 {"a": 2E0, "b": "4", "c": {"a": false, "b": 2}}
2329 "#;
2330
2331 let schema = Arc::new(Schema::new(vec![
2332 Field::new("a", DataType::Int16, false),
2333 Field::new("b", DataType::Utf8, false),
2334 Field::new_struct("c", vec![Field::new("a", DataType::Boolean, false)], false),
2335 ]));
2336
2337 let err = ReaderBuilder::new(schema)
2338 .with_batch_size(1024)
2339 .with_strict_mode(true)
2340 .build(Cursor::new(buf.as_bytes()))
2341 .unwrap()
2342 .read()
2343 .unwrap_err();
2344
2345 assert_eq!(
2346 err.to_string(),
2347 "Json error: whilst decoding field 'c': column 'b' missing from schema"
2348 );
2349 }
2350
2351 fn read_file(path: &str, schema: Option<Schema>) -> Reader<BufReader<File>> {
2352 let file = File::open(path).unwrap();
2353 let mut reader = BufReader::new(file);
2354 let schema = schema.unwrap_or_else(|| {
2355 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
2356 reader.rewind().unwrap();
2357 schema
2358 });
2359 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(64);
2360 builder.build(reader).unwrap()
2361 }
2362
2363 #[test]
2364 fn test_json_basic() {
2365 let mut reader = read_file("test/data/basic.json", None);
2366 let batch = reader.next().unwrap().unwrap();
2367
2368 assert_eq!(8, batch.num_columns());
2369 assert_eq!(12, batch.num_rows());
2370
2371 let schema = reader.schema();
2372 let batch_schema = batch.schema();
2373 assert_eq!(schema, batch_schema);
2374
2375 let a = schema.column_with_name("a").unwrap();
2376 assert_eq!(0, a.0);
2377 assert_eq!(&DataType::Int64, a.1.data_type());
2378 let b = schema.column_with_name("b").unwrap();
2379 assert_eq!(1, b.0);
2380 assert_eq!(&DataType::Float64, b.1.data_type());
2381 let c = schema.column_with_name("c").unwrap();
2382 assert_eq!(2, c.0);
2383 assert_eq!(&DataType::Boolean, c.1.data_type());
2384 let d = schema.column_with_name("d").unwrap();
2385 assert_eq!(3, d.0);
2386 assert_eq!(&DataType::Utf8, d.1.data_type());
2387
2388 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2389 assert_eq!(1, aa.value(0));
2390 assert_eq!(-10, aa.value(1));
2391 let bb = batch.column(b.0).as_primitive::<Float64Type>();
2392 assert_eq!(2.0, bb.value(0));
2393 assert_eq!(-3.5, bb.value(1));
2394 let cc = batch.column(c.0).as_boolean();
2395 assert!(!cc.value(0));
2396 assert!(cc.value(10));
2397 let dd = batch.column(d.0).as_string::<i32>();
2398 assert_eq!("4", dd.value(0));
2399 assert_eq!("text", dd.value(8));
2400 }
2401
2402 #[test]
2403 fn test_json_empty_projection() {
2404 let mut reader = read_file("test/data/basic.json", Some(Schema::empty()));
2405 let batch = reader.next().unwrap().unwrap();
2406
2407 assert_eq!(0, batch.num_columns());
2408 assert_eq!(12, batch.num_rows());
2409 }
2410
2411 #[test]
2412 fn test_json_basic_with_nulls() {
2413 let mut reader = read_file("test/data/basic_nulls.json", None);
2414 let batch = reader.next().unwrap().unwrap();
2415
2416 assert_eq!(4, batch.num_columns());
2417 assert_eq!(12, batch.num_rows());
2418
2419 let schema = reader.schema();
2420 let batch_schema = batch.schema();
2421 assert_eq!(schema, batch_schema);
2422
2423 let a = schema.column_with_name("a").unwrap();
2424 assert_eq!(&DataType::Int64, a.1.data_type());
2425 let b = schema.column_with_name("b").unwrap();
2426 assert_eq!(&DataType::Float64, b.1.data_type());
2427 let c = schema.column_with_name("c").unwrap();
2428 assert_eq!(&DataType::Boolean, c.1.data_type());
2429 let d = schema.column_with_name("d").unwrap();
2430 assert_eq!(&DataType::Utf8, d.1.data_type());
2431
2432 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2433 assert!(aa.is_valid(0));
2434 assert!(!aa.is_valid(1));
2435 assert!(!aa.is_valid(11));
2436 let bb = batch.column(b.0).as_primitive::<Float64Type>();
2437 assert!(bb.is_valid(0));
2438 assert!(!bb.is_valid(2));
2439 assert!(!bb.is_valid(11));
2440 let cc = batch.column(c.0).as_boolean();
2441 assert!(cc.is_valid(0));
2442 assert!(!cc.is_valid(4));
2443 assert!(!cc.is_valid(11));
2444 let dd = batch.column(d.0).as_string::<i32>();
2445 assert!(!dd.is_valid(0));
2446 assert!(dd.is_valid(1));
2447 assert!(!dd.is_valid(4));
2448 assert!(!dd.is_valid(11));
2449 }
2450
2451 #[test]
2452 fn test_json_basic_schema() {
2453 let schema = Schema::new(vec![
2454 Field::new("a", DataType::Int64, true),
2455 Field::new("b", DataType::Float32, false),
2456 Field::new("c", DataType::Boolean, false),
2457 Field::new("d", DataType::Utf8, false),
2458 ]);
2459
2460 let mut reader = read_file("test/data/basic.json", Some(schema.clone()));
2461 let reader_schema = reader.schema();
2462 assert_eq!(reader_schema.as_ref(), &schema);
2463 let batch = reader.next().unwrap().unwrap();
2464
2465 assert_eq!(4, batch.num_columns());
2466 assert_eq!(12, batch.num_rows());
2467
2468 let schema = batch.schema();
2469
2470 let a = schema.column_with_name("a").unwrap();
2471 assert_eq!(&DataType::Int64, a.1.data_type());
2472 let b = schema.column_with_name("b").unwrap();
2473 assert_eq!(&DataType::Float32, b.1.data_type());
2474 let c = schema.column_with_name("c").unwrap();
2475 assert_eq!(&DataType::Boolean, c.1.data_type());
2476 let d = schema.column_with_name("d").unwrap();
2477 assert_eq!(&DataType::Utf8, d.1.data_type());
2478
2479 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2480 assert_eq!(1, aa.value(0));
2481 assert_eq!(100000000000000, aa.value(11));
2482 let bb = batch.column(b.0).as_primitive::<Float32Type>();
2483 assert_eq!(2.0, bb.value(0));
2484 assert_eq!(-3.5, bb.value(1));
2485 }
2486
2487 #[test]
2488 fn test_json_basic_schema_projection() {
2489 let schema = Schema::new(vec![
2490 Field::new("a", DataType::Int64, true),
2491 Field::new("c", DataType::Boolean, false),
2492 ]);
2493
2494 let mut reader = read_file("test/data/basic.json", Some(schema.clone()));
2495 let batch = reader.next().unwrap().unwrap();
2496
2497 assert_eq!(2, batch.num_columns());
2498 assert_eq!(2, batch.schema().fields().len());
2499 assert_eq!(12, batch.num_rows());
2500
2501 assert_eq!(batch.schema().as_ref(), &schema);
2502
2503 let a = schema.column_with_name("a").unwrap();
2504 assert_eq!(0, a.0);
2505 assert_eq!(&DataType::Int64, a.1.data_type());
2506 let c = schema.column_with_name("c").unwrap();
2507 assert_eq!(1, c.0);
2508 assert_eq!(&DataType::Boolean, c.1.data_type());
2509 }
2510
2511 #[test]
2512 fn test_json_arrays() {
2513 let mut reader = read_file("test/data/arrays.json", None);
2514 let batch = reader.next().unwrap().unwrap();
2515
2516 assert_eq!(4, batch.num_columns());
2517 assert_eq!(3, batch.num_rows());
2518
2519 let schema = batch.schema();
2520
2521 let a = schema.column_with_name("a").unwrap();
2522 assert_eq!(&DataType::Int64, a.1.data_type());
2523 let b = schema.column_with_name("b").unwrap();
2524 assert_eq!(
2525 &DataType::List(Arc::new(Field::new_list_field(DataType::Float64, true))),
2526 b.1.data_type()
2527 );
2528 let c = schema.column_with_name("c").unwrap();
2529 assert_eq!(
2530 &DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))),
2531 c.1.data_type()
2532 );
2533 let d = schema.column_with_name("d").unwrap();
2534 assert_eq!(&DataType::Utf8, d.1.data_type());
2535
2536 let aa = batch.column(a.0).as_primitive::<Int64Type>();
2537 assert_eq!(1, aa.value(0));
2538 assert_eq!(-10, aa.value(1));
2539 assert_eq!(1627668684594000000, aa.value(2));
2540 let bb = batch.column(b.0).as_list::<i32>();
2541 let bb = bb.values().as_primitive::<Float64Type>();
2542 assert_eq!(9, bb.len());
2543 assert_eq!(2.0, bb.value(0));
2544 assert_eq!(-6.1, bb.value(5));
2545 assert!(!bb.is_valid(7));
2546
2547 let cc = batch
2548 .column(c.0)
2549 .as_any()
2550 .downcast_ref::<ListArray>()
2551 .unwrap();
2552 let cc = cc.values().as_boolean();
2553 assert_eq!(6, cc.len());
2554 assert!(!cc.value(0));
2555 assert!(!cc.value(4));
2556 assert!(!cc.is_valid(5));
2557 }
2558
2559 #[test]
2560 fn test_empty_json_arrays() {
2561 let json_content = r#"
2562 {"items": []}
2563 {"items": null}
2564 {}
2565 "#;
2566
2567 let schema = Arc::new(Schema::new(vec![Field::new(
2568 "items",
2569 DataType::List(FieldRef::new(Field::new_list_field(DataType::Null, true))),
2570 true,
2571 )]));
2572
2573 let batches = do_read(json_content, 1024, false, false, schema);
2574 assert_eq!(batches.len(), 1);
2575
2576 let col1 = batches[0].column(0).as_list::<i32>();
2577 assert_eq!(col1.null_count(), 2);
2578 assert!(col1.value(0).is_empty());
2579 assert_eq!(col1.value(0).data_type(), &DataType::Null);
2580 assert!(col1.is_null(1));
2581 assert!(col1.is_null(2));
2582 }
2583
2584 #[test]
2585 fn test_nested_empty_json_arrays() {
2586 let json_content = r#"
2587 {"items": [[],[]]}
2588 {"items": [[null, null],[null]]}
2589 "#;
2590
2591 let schema = Arc::new(Schema::new(vec![Field::new(
2592 "items",
2593 DataType::List(FieldRef::new(Field::new_list_field(
2594 DataType::List(FieldRef::new(Field::new_list_field(DataType::Null, true))),
2595 true,
2596 ))),
2597 true,
2598 )]));
2599
2600 let batches = do_read(json_content, 1024, false, false, schema);
2601 assert_eq!(batches.len(), 1);
2602
2603 let col1 = batches[0].column(0).as_list::<i32>();
2604 assert_eq!(col1.null_count(), 0);
2605 assert_eq!(col1.value(0).len(), 2);
2606 assert!(col1.value(0).as_list::<i32>().value(0).is_empty());
2607 assert!(col1.value(0).as_list::<i32>().value(1).is_empty());
2608
2609 assert_eq!(col1.value(1).len(), 2);
2610 assert_eq!(col1.value(1).as_list::<i32>().value(0).len(), 2);
2611 assert_eq!(col1.value(1).as_list::<i32>().value(1).len(), 1);
2612 }
2613
2614 #[test]
2615 fn test_nested_list_json_arrays() {
2616 let c_field = Field::new_struct("c", vec![Field::new("d", DataType::Utf8, true)], true);
2617 let a_struct_field = Field::new_struct(
2618 "a",
2619 vec![Field::new("b", DataType::Boolean, true), c_field.clone()],
2620 true,
2621 );
2622 let a_field = Field::new("a", DataType::List(Arc::new(a_struct_field.clone())), true);
2623 let schema = Arc::new(Schema::new(vec![a_field.clone()]));
2624 let builder = ReaderBuilder::new(schema).with_batch_size(64);
2625 let json_content = r#"
2626 {"a": [{"b": true, "c": {"d": "a_text"}}, {"b": false, "c": {"d": "b_text"}}]}
2627 {"a": [{"b": false, "c": null}]}
2628 {"a": [{"b": true, "c": {"d": "c_text"}}, {"b": null, "c": {"d": "d_text"}}, {"b": true, "c": {"d": null}}]}
2629 {"a": null}
2630 {"a": []}
2631 {"a": [null]}
2632 "#;
2633 let mut reader = builder.build(Cursor::new(json_content)).unwrap();
2634
2635 let d = StringArray::from(vec![
2637 Some("a_text"),
2638 Some("b_text"),
2639 None,
2640 Some("c_text"),
2641 Some("d_text"),
2642 None,
2643 None,
2644 ]);
2645 let c = StructArray::new(
2646 vec![Field::new("d", DataType::Utf8, true)].into(),
2647 vec![Arc::new(d.clone()) as ArrayRef],
2648 Some(NullBuffer::from(vec![
2649 true, true, false, true, true, true, false,
2650 ])),
2651 );
2652 let b = BooleanArray::from(vec![
2653 Some(true),
2654 Some(false),
2655 Some(false),
2656 Some(true),
2657 None,
2658 Some(true),
2659 None,
2660 ]);
2661 let a = StructArray::new(
2662 vec![Field::new("b", DataType::Boolean, true), c_field.clone()].into(),
2663 vec![
2664 Arc::new(b.clone()) as ArrayRef,
2665 Arc::new(c.clone()) as ArrayRef,
2666 ],
2667 Some(NullBuffer::from(vec![
2668 true, true, true, true, true, true, false,
2669 ])),
2670 );
2671 let a_list = ListArray::new(
2672 Arc::new(a_struct_field.clone()),
2673 OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 2, 3, 6, 6, 6, 7])),
2674 Arc::new(a),
2675 Some(NullBuffer::from(vec![true, true, true, false, true, true])),
2676 );
2677
2678 let batch = reader.next().unwrap().unwrap();
2680 let read = batch.column(0);
2681 assert_eq!(read.len(), 6);
2682 let read: &ListArray = read.as_list::<i32>();
2684 let expected = &a_list;
2685 assert_eq!(read.value_offsets(), &[0, 2, 3, 6, 6, 6, 7]);
2686 assert_eq!(read.nulls(), expected.nulls());
2688 let struct_array = read.values().as_struct();
2690 let expected_struct_array = expected.values().as_struct();
2691
2692 assert_eq!(7, struct_array.len());
2693 assert_eq!(1, struct_array.null_count());
2694 assert_eq!(7, expected_struct_array.len());
2695 assert_eq!(1, expected_struct_array.null_count());
2696 assert_eq!(struct_array.nulls(), expected_struct_array.nulls());
2698 let read_b = struct_array.column(0);
2700 assert_eq!(read_b.as_ref(), &b);
2701 let read_c = struct_array.column(1);
2702 assert_eq!(read_c.as_struct(), &c);
2703 let read_c = read_c.as_struct();
2704 let read_d = read_c.column(0);
2705 assert_eq!(read_d.as_ref(), &d);
2706
2707 assert_eq!(read, expected);
2708 }
2709
2710 fn assert_read_list_view<O: OffsetSizeTrait>() {
2711 let field = Arc::new(Field::new("item", DataType::Int32, true));
2712 let data_type = GenericListViewArray::<O>::DATA_TYPE_CONSTRUCTOR(field.clone());
2713 let schema = Arc::new(Schema::new(vec![Field::new("lv", data_type, true)]));
2714
2715 let buf = r#"
2716 {"lv": [1, 2, 3]}
2717 {"lv": [4, null]}
2718 {"lv": null}
2719 {"lv": [6]}
2720 {"lv": []}
2721 "#;
2722
2723 let batches = do_read(buf, 1024, false, false, schema);
2724 assert_eq!(batches.len(), 1);
2725 let batch = &batches[0];
2726 let col = batch.column(0);
2727 let list_view = col
2728 .as_any()
2729 .downcast_ref::<GenericListViewArray<O>>()
2730 .unwrap();
2731
2732 assert_eq!(list_view.len(), 5);
2733
2734 let expected_offsets: Vec<O> = vec![0, 3, 5, 5, 6]
2736 .into_iter()
2737 .map(|v| O::usize_as(v))
2738 .collect();
2739 let expected_sizes: Vec<O> = vec![3, 2, 0, 1, 0]
2740 .into_iter()
2741 .map(|v| O::usize_as(v))
2742 .collect();
2743 assert_eq!(list_view.value_offsets(), &expected_offsets);
2744 assert_eq!(list_view.value_sizes(), &expected_sizes);
2745
2746 assert!(list_view.is_valid(0));
2748 let vals = list_view.value(0);
2749 let ints = vals.as_primitive::<Int32Type>();
2750 assert_eq!(ints.values(), &[1, 2, 3]);
2751
2752 assert!(list_view.is_valid(1));
2754 let vals = list_view.value(1);
2755 let ints = vals.as_primitive::<Int32Type>();
2756 assert_eq!(ints.len(), 2);
2757 assert_eq!(ints.value(0), 4);
2758 assert!(ints.is_null(1));
2759
2760 assert!(list_view.is_null(2));
2762
2763 assert!(list_view.is_valid(3));
2765 let vals = list_view.value(3);
2766 let ints = vals.as_primitive::<Int32Type>();
2767 assert_eq!(ints.values(), &[6]);
2768
2769 assert!(list_view.is_valid(4));
2771 let vals = list_view.value(4);
2772 assert_eq!(vals.len(), 0);
2773 }
2774
2775 #[test]
2776 fn test_read_list_view() {
2777 assert_read_list_view::<i32>();
2778 assert_read_list_view::<i64>();
2779 }
2780
2781 #[test]
2782 fn test_read_list_view_rejects_null_non_nullable_child() {
2783 let field = Arc::new(Field::new("item", DataType::Int32, false));
2784 for (data_type, array_type) in [
2785 (DataType::ListView(field.clone()), "ListViewArray"),
2786 (DataType::LargeListView(field.clone()), "LargeListViewArray"),
2787 ] {
2788 let schema = Arc::new(Schema::new(vec![Field::new("lv", data_type, true)]));
2789 let buf = r#"
2790 {"lv": [1, 2, 3]}
2791 {"lv": [4, null]}
2792 "#;
2793
2794 let error = ReaderBuilder::new(schema)
2795 .build(Cursor::new(buf.as_bytes()))
2796 .unwrap()
2797 .collect::<Result<Vec<_>, _>>()
2798 .unwrap_err();
2799
2800 assert_eq!(
2801 error.to_string(),
2802 format!(
2803 "Invalid argument error: Non-nullable field of {array_type} \"item\" cannot contain nulls"
2804 )
2805 );
2806 }
2807 }
2808
2809 #[test]
2810 fn test_fixed_size_list() {
2811 let buf = r#"
2812 {"a": [1, 2, 3]}
2813 {"a": [4, 5, 6]}
2814 {"a": [7, 8, 9]}
2815 "#;
2816
2817 let field = Field::new_list_field(DataType::Int32, true);
2818 let schema = Arc::new(Schema::new(vec![Field::new(
2819 "a",
2820 DataType::FixedSizeList(Arc::new(field), 3),
2821 false,
2822 )]));
2823
2824 let batches = do_read(buf, 1024, false, false, schema);
2825 assert_eq!(batches.len(), 1);
2826
2827 let col = batches[0].column(0).as_fixed_size_list();
2828 assert_eq!(col.len(), 3);
2829 assert_eq!(col.value_length(), 3);
2830
2831 let values = col.values().as_primitive::<Int32Type>();
2832 assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
2833 }
2834
2835 #[test]
2836 fn test_fixed_size_list_nullable() {
2837 let buf = r#"
2838 {"a": [1, 2]}
2839 {"a": null}
2840 {"a": [3, null]}
2841 "#;
2842
2843 let field = Field::new_list_field(DataType::Int32, true);
2844 let schema = Arc::new(Schema::new(vec![Field::new(
2845 "a",
2846 DataType::FixedSizeList(Arc::new(field), 2),
2847 true,
2848 )]));
2849
2850 let batches = do_read(buf, 1024, false, false, schema);
2851 assert_eq!(batches.len(), 1);
2852
2853 let col = batches[0].column(0).as_fixed_size_list();
2854 assert_eq!(col.len(), 3);
2855 assert!(col.is_valid(0));
2856 assert!(col.is_null(1));
2857 assert!(col.is_valid(2));
2858
2859 let values = col.values().as_primitive::<Int32Type>();
2860 assert_eq!(values.value(0), 1);
2861 assert_eq!(values.value(1), 2);
2862 assert_eq!(values.value(4), 3);
2863 assert!(values.is_null(5));
2864 }
2865
2866 #[test]
2867 fn test_fixed_size_list_zero_size_non_nullable() {
2868 let buf = r#"
2869 {"a": []}
2870 {"a": []}
2871 {"a": []}
2872 "#;
2873
2874 let field = Field::new_list_field(DataType::Int32, true);
2875 let schema = Arc::new(Schema::new(vec![Field::new(
2876 "a",
2877 DataType::FixedSizeList(Arc::new(field), 0),
2878 false,
2879 )]));
2880
2881 let batches = do_read(buf, 1024, false, false, schema);
2882 assert_eq!(batches.len(), 1);
2883
2884 let col = batches[0].column(0).as_fixed_size_list();
2885 assert_eq!(col.len(), 3);
2886 assert_eq!(col.value_length(), 0);
2887
2888 let values = col.values().as_primitive::<Int32Type>();
2889 assert!(values.values().is_empty());
2890 }
2891
2892 #[test]
2893 fn test_fixed_size_list_wrong_size() {
2894 let buf = r#"{"a": [1, 2, 3]}"#;
2895
2896 let field = Field::new_list_field(DataType::Int32, true);
2897 let schema = Arc::new(Schema::new(vec![Field::new(
2898 "a",
2899 DataType::FixedSizeList(Arc::new(field), 2),
2900 false,
2901 )]));
2902
2903 let err = ReaderBuilder::new(schema)
2904 .build(Cursor::new(buf.as_bytes()))
2905 .unwrap()
2906 .next()
2907 .unwrap()
2908 .unwrap_err();
2909
2910 assert!(err.to_string().contains("expected 2 but got 3"), "{}", err);
2911 }
2912
2913 #[test]
2914 fn test_fixed_size_list_nested() {
2915 let buf = r#"
2916 {"a": [[1, 2], [3, 4]]}
2917 {"a": [[5, 6], [7, 8]]}
2918 "#;
2919
2920 let inner_field = Field::new_list_field(DataType::Int32, true);
2921 let inner_type = DataType::FixedSizeList(Arc::new(inner_field), 2);
2922 let outer_field = Arc::new(Field::new_list_field(inner_type.clone(), true));
2923 let schema = Arc::new(Schema::new(vec![Field::new(
2924 "a",
2925 DataType::FixedSizeList(outer_field, 2),
2926 false,
2927 )]));
2928
2929 let batches = do_read(buf, 1024, false, false, schema);
2930 assert_eq!(batches.len(), 1);
2931
2932 let col = batches[0].column(0).as_fixed_size_list();
2933 assert_eq!(col.len(), 2);
2934 assert_eq!(col.value_length(), 2);
2935
2936 let inner = col.values().as_fixed_size_list();
2937 assert_eq!(inner.len(), 4);
2938 assert_eq!(inner.value_length(), 2);
2939
2940 let values = inner.values().as_primitive::<Int32Type>();
2941 assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6, 7, 8]);
2942 }
2943
2944 #[test]
2945 fn test_fixed_size_list_ignore_type_conflicts() {
2946 let field = Field::new("item", DataType::Int32, true);
2947 let schema = Arc::new(Schema::new(vec![Field::new(
2948 "a",
2949 DataType::FixedSizeList(Arc::new(field), 2),
2950 true,
2951 )]));
2952
2953 let json = vec![
2954 json!({"a": [1, 2]}),
2955 json!({"a": "not a list"}),
2956 json!({"a": 42}),
2957 json!({"a": [6, 7]}),
2958 ];
2959
2960 let mut decoder = ReaderBuilder::new(schema)
2961 .with_ignore_type_conflicts(true)
2962 .build_decoder()
2963 .unwrap();
2964 decoder.serialize(&json).unwrap();
2965 let batch = decoder.flush().unwrap().unwrap();
2966
2967 let col = batch.column(0).as_fixed_size_list();
2968 assert_eq!(col.len(), 4);
2969 assert!(col.is_valid(0));
2970 assert!(col.is_null(1)); assert!(col.is_null(2)); assert!(col.is_valid(3));
2973
2974 let values = col.values().as_primitive::<Int32Type>();
2975 assert_eq!(values.value(0), 1);
2976 assert_eq!(values.value(1), 2);
2977 assert_eq!(values.value(6), 6);
2978 assert_eq!(values.value(7), 7);
2979 }
2980
2981 #[test]
2982 fn test_skip_empty_lines() {
2983 let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
2984 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(64);
2985 let json_content = "
2986 {\"a\": 1}
2987 {\"a\": 2}
2988 {\"a\": 3}";
2989 let mut reader = builder.build(Cursor::new(json_content)).unwrap();
2990 let batch = reader.next().unwrap().unwrap();
2991
2992 assert_eq!(1, batch.num_columns());
2993 assert_eq!(3, batch.num_rows());
2994
2995 let schema = reader.schema();
2996 let c = schema.column_with_name("a").unwrap();
2997 assert_eq!(&DataType::Int64, c.1.data_type());
2998 }
2999
3000 #[test]
3001 fn test_with_multiple_batches() {
3002 let file = File::open("test/data/basic_nulls.json").unwrap();
3003 let mut reader = BufReader::new(file);
3004 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
3005 reader.rewind().unwrap();
3006
3007 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(5);
3008 let mut reader = builder.build(reader).unwrap();
3009
3010 let mut num_records = Vec::new();
3011 while let Some(rb) = reader.next().transpose().unwrap() {
3012 num_records.push(rb.num_rows());
3013 }
3014
3015 assert_eq!(vec![5, 5, 2], num_records);
3016 }
3017
3018 #[test]
3019 fn test_timestamp_from_json_seconds() {
3020 let schema = Schema::new(vec![Field::new(
3021 "a",
3022 DataType::Timestamp(TimeUnit::Second, None),
3023 true,
3024 )]);
3025
3026 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
3027 let batch = reader.next().unwrap().unwrap();
3028
3029 assert_eq!(1, batch.num_columns());
3030 assert_eq!(12, batch.num_rows());
3031
3032 let schema = reader.schema();
3033 let batch_schema = batch.schema();
3034 assert_eq!(schema, batch_schema);
3035
3036 let a = schema.column_with_name("a").unwrap();
3037 assert_eq!(
3038 &DataType::Timestamp(TimeUnit::Second, None),
3039 a.1.data_type()
3040 );
3041
3042 let aa = batch.column(a.0).as_primitive::<TimestampSecondType>();
3043 assert!(aa.is_valid(0));
3044 assert!(!aa.is_valid(1));
3045 assert!(!aa.is_valid(2));
3046 assert_eq!(1, aa.value(0));
3047 assert_eq!(1, aa.value(3));
3048 assert_eq!(5, aa.value(7));
3049 }
3050
3051 #[test]
3052 fn test_timestamp_from_json_milliseconds() {
3053 let schema = Schema::new(vec![Field::new(
3054 "a",
3055 DataType::Timestamp(TimeUnit::Millisecond, None),
3056 true,
3057 )]);
3058
3059 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
3060 let batch = reader.next().unwrap().unwrap();
3061
3062 assert_eq!(1, batch.num_columns());
3063 assert_eq!(12, batch.num_rows());
3064
3065 let schema = reader.schema();
3066 let batch_schema = batch.schema();
3067 assert_eq!(schema, batch_schema);
3068
3069 let a = schema.column_with_name("a").unwrap();
3070 assert_eq!(
3071 &DataType::Timestamp(TimeUnit::Millisecond, None),
3072 a.1.data_type()
3073 );
3074
3075 let aa = batch.column(a.0).as_primitive::<TimestampMillisecondType>();
3076 assert!(aa.is_valid(0));
3077 assert!(!aa.is_valid(1));
3078 assert!(!aa.is_valid(2));
3079 assert_eq!(1, aa.value(0));
3080 assert_eq!(1, aa.value(3));
3081 assert_eq!(5, aa.value(7));
3082 }
3083
3084 #[test]
3085 fn test_date_from_json_milliseconds() {
3086 let schema = Schema::new(vec![Field::new("a", DataType::Date64, true)]);
3087
3088 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
3089 let batch = reader.next().unwrap().unwrap();
3090
3091 assert_eq!(1, batch.num_columns());
3092 assert_eq!(12, batch.num_rows());
3093
3094 let schema = reader.schema();
3095 let batch_schema = batch.schema();
3096 assert_eq!(schema, batch_schema);
3097
3098 let a = schema.column_with_name("a").unwrap();
3099 assert_eq!(&DataType::Date64, a.1.data_type());
3100
3101 let aa = batch.column(a.0).as_primitive::<Date64Type>();
3102 assert!(aa.is_valid(0));
3103 assert!(!aa.is_valid(1));
3104 assert!(!aa.is_valid(2));
3105 assert_eq!(1, aa.value(0));
3106 assert_eq!(1, aa.value(3));
3107 assert_eq!(5, aa.value(7));
3108 }
3109
3110 #[test]
3111 fn test_time_from_json_nanoseconds() {
3112 let schema = Schema::new(vec![Field::new(
3113 "a",
3114 DataType::Time64(TimeUnit::Nanosecond),
3115 true,
3116 )]);
3117
3118 let mut reader = read_file("test/data/basic_nulls.json", Some(schema));
3119 let batch = reader.next().unwrap().unwrap();
3120
3121 assert_eq!(1, batch.num_columns());
3122 assert_eq!(12, batch.num_rows());
3123
3124 let schema = reader.schema();
3125 let batch_schema = batch.schema();
3126 assert_eq!(schema, batch_schema);
3127
3128 let a = schema.column_with_name("a").unwrap();
3129 assert_eq!(&DataType::Time64(TimeUnit::Nanosecond), a.1.data_type());
3130
3131 let aa = batch.column(a.0).as_primitive::<Time64NanosecondType>();
3132 assert!(aa.is_valid(0));
3133 assert!(!aa.is_valid(1));
3134 assert!(!aa.is_valid(2));
3135 assert_eq!(1, aa.value(0));
3136 assert_eq!(1, aa.value(3));
3137 assert_eq!(5, aa.value(7));
3138 }
3139
3140 #[test]
3141 fn test_json_iterator() {
3142 let file = File::open("test/data/basic.json").unwrap();
3143 let mut reader = BufReader::new(file);
3144 let (schema, _) = infer_json_schema(&mut reader, None).unwrap();
3145 reader.rewind().unwrap();
3146
3147 let builder = ReaderBuilder::new(Arc::new(schema)).with_batch_size(5);
3148 let reader = builder.build(reader).unwrap();
3149 let schema = reader.schema();
3150 let (col_a_index, _) = schema.column_with_name("a").unwrap();
3151
3152 let mut sum_num_rows = 0;
3153 let mut num_batches = 0;
3154 let mut sum_a = 0;
3155 for batch in reader {
3156 let batch = batch.unwrap();
3157 assert_eq!(8, batch.num_columns());
3158 sum_num_rows += batch.num_rows();
3159 num_batches += 1;
3160 let batch_schema = batch.schema();
3161 assert_eq!(schema, batch_schema);
3162 let a_array = batch.column(col_a_index).as_primitive::<Int64Type>();
3163 sum_a += (0..a_array.len()).map(|i| a_array.value(i)).sum::<i64>();
3164 }
3165 assert_eq!(12, sum_num_rows);
3166 assert_eq!(3, num_batches);
3167 assert_eq!(100000000000011, sum_a);
3168 }
3169
3170 #[test]
3171 fn test_decoder_error() {
3172 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3173 "a",
3174 vec![Field::new("child", DataType::Int32, false)],
3175 true,
3176 )]));
3177
3178 let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder().unwrap();
3179 let _ = decoder.decode(br#"{"a": { "child":"#).unwrap();
3180 assert!(decoder.tape_decoder.has_partial_row());
3181 assert_eq!(decoder.tape_decoder.num_buffered_rows(), 1);
3182 let _ = decoder.flush().unwrap_err();
3183 assert!(decoder.tape_decoder.has_partial_row());
3184 assert_eq!(decoder.tape_decoder.num_buffered_rows(), 1);
3185
3186 let parse_err = |s: &str| {
3187 ReaderBuilder::new(schema.clone())
3188 .build(Cursor::new(s.as_bytes()))
3189 .unwrap()
3190 .next()
3191 .unwrap()
3192 .unwrap_err()
3193 .to_string()
3194 };
3195
3196 let err = parse_err(r#"{"a": 123}"#);
3197 assert_eq!(
3198 err,
3199 "Json error: whilst decoding field 'a': expected { got 123"
3200 );
3201
3202 let err = parse_err(r#"{"a": ["bar"]}"#);
3203 assert_eq!(
3204 err,
3205 r#"Json error: whilst decoding field 'a': expected { got ["bar"]"#
3206 );
3207
3208 let err = parse_err(r#"{"a": []}"#);
3209 assert_eq!(
3210 err,
3211 "Json error: whilst decoding field 'a': expected { got []"
3212 );
3213
3214 let err = parse_err(r#"{"a": [{"child": 234}]}"#);
3215 assert_eq!(
3216 err,
3217 r#"Json error: whilst decoding field 'a': expected { got [{"child": 234}]"#
3218 );
3219
3220 let err = parse_err(r#"{"a": [{"child": {"foo": [{"foo": ["bar"]}]}}]}"#);
3221 assert_eq!(
3222 err,
3223 r#"Json error: whilst decoding field 'a': expected { got [{"child": {"foo": [{"foo": ["bar"]}]}}]"#
3224 );
3225
3226 let err = parse_err(r#"{"a": true}"#);
3227 assert_eq!(
3228 err,
3229 "Json error: whilst decoding field 'a': expected { got true"
3230 );
3231
3232 let err = parse_err(r#"{"a": false}"#);
3233 assert_eq!(
3234 err,
3235 "Json error: whilst decoding field 'a': expected { got false"
3236 );
3237
3238 let err = parse_err(r#"{"a": "foo"}"#);
3239 assert_eq!(
3240 err,
3241 "Json error: whilst decoding field 'a': expected { got \"foo\""
3242 );
3243
3244 let err = parse_err(r#"{"a": {"child": false}}"#);
3245 assert_eq!(
3246 err,
3247 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got false"
3248 );
3249
3250 let err = parse_err(r#"{"a": {"child": []}}"#);
3251 assert_eq!(
3252 err,
3253 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got []"
3254 );
3255
3256 let err = parse_err(r#"{"a": {"child": [123]}}"#);
3257 assert_eq!(
3258 err,
3259 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got [123]"
3260 );
3261
3262 let err = parse_err(r#"{"a": {"child": [123, 3465346]}}"#);
3263 assert_eq!(
3264 err,
3265 "Json error: whilst decoding field 'a': whilst decoding field 'child': expected primitive got [123, 3465346]"
3266 );
3267 }
3268
3269 #[test]
3270 fn test_serialize_timestamp() {
3271 let json = vec![
3272 json!({"timestamp": 1681319393}),
3273 json!({"timestamp": "1970-01-01T00:00:00+02:00"}),
3274 ];
3275 let schema = Schema::new(vec![Field::new(
3276 "timestamp",
3277 DataType::Timestamp(TimeUnit::Second, None),
3278 true,
3279 )]);
3280 let mut decoder = ReaderBuilder::new(Arc::new(schema))
3281 .build_decoder()
3282 .unwrap();
3283 decoder.serialize(&json).unwrap();
3284 let batch = decoder.flush().unwrap().unwrap();
3285 assert_eq!(batch.num_rows(), 2);
3286 assert_eq!(batch.num_columns(), 1);
3287 let values = batch.column(0).as_primitive::<TimestampSecondType>();
3288 assert_eq!(values.values(), &[1681319393, -7200]);
3289 }
3290
3291 #[test]
3292 fn test_serialize_decimal() {
3293 let json = vec![
3294 json!({"decimal": 1.234}),
3295 json!({"decimal": "1.234"}),
3296 json!({"decimal": 1234}),
3297 json!({"decimal": "1234"}),
3298 ];
3299 let schema = Schema::new(vec![Field::new(
3300 "decimal",
3301 DataType::Decimal128(10, 3),
3302 true,
3303 )]);
3304 let mut decoder = ReaderBuilder::new(Arc::new(schema))
3305 .build_decoder()
3306 .unwrap();
3307 decoder.serialize(&json).unwrap();
3308 let batch = decoder.flush().unwrap().unwrap();
3309 assert_eq!(batch.num_rows(), 4);
3310 assert_eq!(batch.num_columns(), 1);
3311 let values = batch.column(0).as_primitive::<Decimal128Type>();
3312 assert_eq!(values.values(), &[1234, 1234, 1234000, 1234000]);
3313 }
3314
3315 #[test]
3316 fn test_serde_field() {
3317 let field = Field::new("int", DataType::Int32, true);
3318 let mut decoder = ReaderBuilder::new_with_field(field)
3319 .build_decoder()
3320 .unwrap();
3321 decoder.serialize(&[1_i32, 2, 3, 4]).unwrap();
3322 let b = decoder.flush().unwrap().unwrap();
3323 let values = b.column(0).as_primitive::<Int32Type>().values();
3324 assert_eq!(values, &[1, 2, 3, 4]);
3325 }
3326
3327 #[test]
3328 fn test_serde_large_numbers() {
3329 let field = Field::new("int", DataType::Int64, true);
3330 let mut decoder = ReaderBuilder::new_with_field(field)
3331 .build_decoder()
3332 .unwrap();
3333
3334 decoder.serialize(&[1699148028689_u64, 2, 3, 4]).unwrap();
3335 let b = decoder.flush().unwrap().unwrap();
3336 let values = b.column(0).as_primitive::<Int64Type>().values();
3337 assert_eq!(values, &[1699148028689, 2, 3, 4]);
3338
3339 let field = Field::new(
3340 "int",
3341 DataType::Timestamp(TimeUnit::Microsecond, None),
3342 true,
3343 );
3344 let mut decoder = ReaderBuilder::new_with_field(field)
3345 .build_decoder()
3346 .unwrap();
3347
3348 decoder.serialize(&[1699148028689_u64, 2, 3, 4]).unwrap();
3349 let b = decoder.flush().unwrap().unwrap();
3350 let values = b
3351 .column(0)
3352 .as_primitive::<TimestampMicrosecondType>()
3353 .values();
3354 assert_eq!(values, &[1699148028689, 2, 3, 4]);
3355 }
3356
3357 #[test]
3358 fn test_coercing_primitive_into_string_decoder() {
3359 let buf = &format!(
3360 r#"[{{"a": 1, "b": "A", "c": "T"}}, {{"a": 2, "b": "BB", "c": "F"}}, {{"a": {}, "b": 123, "c": false}}, {{"a": {}, "b": 789, "c": true}}]"#,
3361 (i32::MAX as i64 + 10),
3362 i64::MAX - 10
3363 );
3364 let schema = Schema::new(vec![
3365 Field::new("a", DataType::Float64, true),
3366 Field::new("b", DataType::Utf8, true),
3367 Field::new("c", DataType::Utf8, true),
3368 ]);
3369 let json_array: Vec<serde_json::Value> = serde_json::from_str(buf).unwrap();
3370 let schema_ref = Arc::new(schema);
3371
3372 let reader = ReaderBuilder::new(schema_ref.clone()).with_coerce_primitive(true);
3374 let mut decoder = reader.build_decoder().unwrap();
3375 decoder.serialize(json_array.as_slice()).unwrap();
3376 let batch = decoder.flush().unwrap().unwrap();
3377 assert_eq!(
3378 batch,
3379 RecordBatch::try_new(
3380 schema_ref,
3381 vec![
3382 Arc::new(Float64Array::from(vec![
3383 1.0,
3384 2.0,
3385 (i32::MAX as i64 + 10) as f64,
3386 (i64::MAX - 10) as f64
3387 ])),
3388 Arc::new(StringArray::from(vec!["A", "BB", "123", "789"])),
3389 Arc::new(StringArray::from(vec!["T", "F", "false", "true"])),
3390 ]
3391 )
3392 .unwrap()
3393 );
3394 }
3395
3396 #[test]
3397 fn test_serialize_f32_into_string() {
3398 let field = Field::new("f", DataType::Utf8, true);
3400 let mut decoder = ReaderBuilder::new_with_field(field)
3401 .with_coerce_primitive(true)
3402 .build_decoder()
3403 .unwrap();
3404 decoder.serialize(&[1.5_f32, -2.25_f32]).unwrap();
3405 let batch = decoder.flush().unwrap().unwrap();
3406 let values = batch.column(0).as_string::<i32>();
3407 assert_eq!(values.value(0), "1.5");
3408 assert_eq!(values.value(1), "-2.25");
3409 }
3410
3411 fn _parse_structs(
3416 row: &str,
3417 struct_mode: StructMode,
3418 fields: Fields,
3419 as_struct: bool,
3420 ) -> Result<RecordBatch, ArrowError> {
3421 let builder = if as_struct {
3422 ReaderBuilder::new_with_field(Field::new("r", DataType::Struct(fields), true))
3423 } else {
3424 ReaderBuilder::new(Arc::new(Schema::new(fields)))
3425 };
3426 builder
3427 .with_struct_mode(struct_mode)
3428 .build(Cursor::new(row.as_bytes()))
3429 .unwrap()
3430 .next()
3431 .unwrap()
3432 }
3433
3434 #[test]
3435 fn test_struct_decoding_list_length() {
3436 use arrow_array::array;
3437
3438 let row = "[1, 2]";
3439
3440 let mut fields = vec![Field::new("a", DataType::Int32, true)];
3441 let too_few_fields = Fields::from(fields.clone());
3442 fields.push(Field::new("b", DataType::Int32, true));
3443 let correct_fields = Fields::from(fields.clone());
3444 fields.push(Field::new("c", DataType::Int32, true));
3445 let too_many_fields = Fields::from(fields.clone());
3446
3447 let parse = |fields: Fields, as_struct: bool| {
3448 _parse_structs(row, StructMode::ListOnly, fields, as_struct)
3449 };
3450
3451 let expected_row = StructArray::new(
3452 correct_fields.clone(),
3453 vec![
3454 Arc::new(array::Int32Array::from(vec![1])),
3455 Arc::new(array::Int32Array::from(vec![2])),
3456 ],
3457 None,
3458 );
3459 let row_field = Field::new("r", DataType::Struct(correct_fields.clone()), true);
3460
3461 assert_eq!(
3462 parse(too_few_fields.clone(), true).unwrap_err().to_string(),
3463 "Json error: found extra columns for 1 fields".to_string()
3464 );
3465 assert_eq!(
3466 parse(too_few_fields, false).unwrap_err().to_string(),
3467 "Json error: found extra columns for 1 fields".to_string()
3468 );
3469 assert_eq!(
3470 parse(correct_fields.clone(), true).unwrap(),
3471 RecordBatch::try_new(
3472 Arc::new(Schema::new(vec![row_field])),
3473 vec![Arc::new(expected_row.clone())]
3474 )
3475 .unwrap()
3476 );
3477 assert_eq!(
3478 parse(correct_fields, false).unwrap(),
3479 RecordBatch::from(expected_row)
3480 );
3481 assert_eq!(
3482 parse(too_many_fields.clone(), true)
3483 .unwrap_err()
3484 .to_string(),
3485 "Json error: found 2 columns for 3 fields".to_string()
3486 );
3487 assert_eq!(
3488 parse(too_many_fields, false).unwrap_err().to_string(),
3489 "Json error: found 2 columns for 3 fields".to_string()
3490 );
3491 }
3492
3493 #[test]
3494 fn test_struct_decoding() {
3495 use arrow_array::builder;
3496
3497 let nested_object_json = r#"{"a": {"b": [1, 2], "c": {"d": 3}}}"#;
3498 let nested_list_json = r#"[[[1, 2], {"d": 3}]]"#;
3499 let nested_mixed_json = r#"{"a": [[1, 2], {"d": 3}]}"#;
3500
3501 let struct_fields = Fields::from(vec![
3502 Field::new("b", DataType::new_list(DataType::Int32, true), true),
3503 Field::new_map(
3504 "c",
3505 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3506 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3507 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
3508 false,
3509 false,
3510 ),
3511 ]);
3512
3513 let list_array =
3514 ListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![Some(1), Some(2)])]);
3515
3516 let map_array = {
3517 let mut map_builder = builder::MapBuilder::new(
3518 None,
3519 builder::StringBuilder::new(),
3520 builder::Int32Builder::new(),
3521 );
3522 map_builder.keys().append_value("d");
3523 map_builder.values().append_value(3);
3524 map_builder.append(true).unwrap();
3525 map_builder.finish()
3526 };
3527
3528 let struct_array = StructArray::new(
3529 struct_fields.clone(),
3530 vec![Arc::new(list_array), Arc::new(map_array)],
3531 None,
3532 );
3533
3534 let fields = Fields::from(vec![Field::new("a", DataType::Struct(struct_fields), true)]);
3535 let schema = Arc::new(Schema::new(fields.clone()));
3536 let expected = RecordBatch::try_new(schema.clone(), vec![Arc::new(struct_array)]).unwrap();
3537
3538 let parse = |row: &str, struct_mode: StructMode| {
3539 _parse_structs(row, struct_mode, fields.clone(), false)
3540 };
3541
3542 assert_eq!(
3543 parse(nested_object_json, StructMode::ObjectOnly).unwrap(),
3544 expected
3545 );
3546 assert_eq!(
3547 parse(nested_list_json, StructMode::ObjectOnly)
3548 .unwrap_err()
3549 .to_string(),
3550 "Json error: expected { got [[[1, 2], {\"d\": 3}]]".to_owned()
3551 );
3552 assert_eq!(
3553 parse(nested_mixed_json, StructMode::ObjectOnly)
3554 .unwrap_err()
3555 .to_string(),
3556 "Json error: whilst decoding field 'a': expected { got [[1, 2], {\"d\": 3}]".to_owned()
3557 );
3558
3559 assert_eq!(
3560 parse(nested_list_json, StructMode::ListOnly).unwrap(),
3561 expected
3562 );
3563 assert_eq!(
3564 parse(nested_object_json, StructMode::ListOnly)
3565 .unwrap_err()
3566 .to_string(),
3567 "Json error: expected [ got {\"a\": {\"b\": [1, 2]\"c\": {\"d\": 3}}}".to_owned()
3568 );
3569 assert_eq!(
3570 parse(nested_mixed_json, StructMode::ListOnly)
3571 .unwrap_err()
3572 .to_string(),
3573 "Json error: expected [ got {\"a\": [[1, 2], {\"d\": 3}]}".to_owned()
3574 );
3575 }
3576
3577 #[test]
3583 fn test_struct_decoding_empty_list() {
3584 let int_field = Field::new("a", DataType::Int32, true);
3585 let struct_field = Field::new(
3586 "r",
3587 DataType::Struct(Fields::from(vec![int_field.clone()])),
3588 true,
3589 );
3590
3591 let parse = |row: &str, as_struct: bool, field: Field| {
3592 _parse_structs(
3593 row,
3594 StructMode::ListOnly,
3595 Fields::from(vec![field]),
3596 as_struct,
3597 )
3598 };
3599
3600 assert_eq!(
3602 parse("[]", true, struct_field.clone())
3603 .unwrap_err()
3604 .to_string(),
3605 "Json error: found 0 columns for 1 fields".to_owned()
3606 );
3607 assert_eq!(
3608 parse("[]", false, int_field.clone())
3609 .unwrap_err()
3610 .to_string(),
3611 "Json error: found 0 columns for 1 fields".to_owned()
3612 );
3613 assert_eq!(
3614 parse("[]", false, struct_field.clone())
3615 .unwrap_err()
3616 .to_string(),
3617 "Json error: found 0 columns for 1 fields".to_owned()
3618 );
3619 assert_eq!(
3620 parse("[[]]", false, struct_field.clone())
3621 .unwrap_err()
3622 .to_string(),
3623 "Json error: whilst decoding field 'r': found 0 columns for 1 fields".to_owned()
3624 );
3625 }
3626
3627 #[test]
3628 fn test_decode_list_struct_with_wrong_types() {
3629 let int_field = Field::new("a", DataType::Int32, true);
3630 let struct_field = Field::new(
3631 "r",
3632 DataType::Struct(Fields::from(vec![int_field.clone()])),
3633 true,
3634 );
3635
3636 let parse = |row: &str, as_struct: bool, field: Field| {
3637 _parse_structs(
3638 row,
3639 StructMode::ListOnly,
3640 Fields::from(vec![field]),
3641 as_struct,
3642 )
3643 };
3644
3645 assert_eq!(
3647 parse(r#"[["a"]]"#, false, struct_field.clone())
3648 .unwrap_err()
3649 .to_string(),
3650 "Json error: whilst decoding field 'r': whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3651 );
3652 assert_eq!(
3653 parse(r#"[["a"]]"#, true, struct_field.clone())
3654 .unwrap_err()
3655 .to_string(),
3656 "Json error: whilst decoding field 'r': whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3657 );
3658 assert_eq!(
3659 parse(r#"["a"]"#, true, int_field.clone())
3660 .unwrap_err()
3661 .to_string(),
3662 "Json error: whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3663 );
3664 assert_eq!(
3665 parse(r#"["a"]"#, false, int_field.clone())
3666 .unwrap_err()
3667 .to_string(),
3668 "Json error: whilst decoding field 'a': failed to parse \"a\" as Int32".to_owned()
3669 );
3670 }
3671
3672 #[test]
3673 fn test_type_conflict_nulls() {
3674 let schema = Schema::new(vec![
3675 Field::new("null", DataType::Null, true),
3676 Field::new("bool", DataType::Boolean, true),
3677 Field::new("primitive", DataType::Int32, true),
3678 Field::new("numeric", DataType::Decimal128(10, 3), true),
3679 Field::new("string", DataType::Utf8, true),
3680 Field::new("string_view", DataType::Utf8View, true),
3681 Field::new(
3682 "timestamp",
3683 DataType::Timestamp(TimeUnit::Second, None),
3684 true,
3685 ),
3686 Field::new(
3687 "array",
3688 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
3689 true,
3690 ),
3691 Field::new(
3692 "map",
3693 DataType::Map(
3694 Arc::new(Field::new(
3695 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3696 DataType::Struct(Fields::from(vec![
3697 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3698 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3699 ])),
3700 false, )),
3702 false, ),
3704 true, ),
3706 Field::new(
3707 "struct",
3708 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
3709 true,
3710 ),
3711 ]);
3712
3713 let json_values = vec![
3715 json!(null),
3716 json!(true),
3717 json!(42),
3718 json!(1.234),
3719 json!("hi"),
3720 json!("ho"),
3721 json!("1970-01-01T00:00:00+02:00"),
3722 json!([1, "ho", 3]),
3723 json!({"k": "value"}),
3724 json!({"a": 1}),
3725 ];
3726
3727 let json: Vec<_> = (0..json_values.len())
3729 .map(|i| {
3730 let pairs = json_values[i..]
3731 .iter()
3732 .chain(json_values[..i].iter())
3733 .zip(&schema.fields)
3734 .map(|(v, f)| (f.name().clone(), v.clone()))
3735 .collect();
3736 serde_json::Value::Object(pairs)
3737 })
3738 .collect();
3739 let mut decoder = ReaderBuilder::new(Arc::new(schema))
3740 .with_ignore_type_conflicts(true)
3741 .with_coerce_primitive(true)
3742 .build_decoder()
3743 .unwrap();
3744 decoder.serialize(&json).unwrap();
3745 let batch = decoder.flush().unwrap().unwrap();
3746 assert_eq!(batch.num_rows(), 10);
3747 assert_eq!(batch.num_columns(), 10);
3748
3749 let _ = batch
3751 .column(0)
3752 .as_any()
3753 .downcast_ref::<NullArray>()
3754 .unwrap();
3755
3756 assert!(
3757 batch
3758 .column(1)
3759 .as_any()
3760 .downcast_ref::<BooleanArray>()
3761 .unwrap()
3762 .iter()
3763 .eq([
3764 Some(true),
3765 None,
3766 None,
3767 None,
3768 None,
3769 None,
3770 None,
3771 None,
3772 None,
3773 None
3774 ])
3775 );
3776
3777 assert!(batch.column(2).as_primitive::<Int32Type>().iter().eq([
3778 Some(42),
3779 Some(1),
3780 None,
3781 None,
3782 None,
3783 None,
3784 None,
3785 None,
3786 None,
3787 None
3788 ]));
3789
3790 assert!(batch.column(3).as_primitive::<Decimal128Type>().iter().eq([
3791 Some(1234),
3792 None,
3793 None,
3794 None,
3795 None,
3796 None,
3797 None,
3798 None,
3799 None,
3800 Some(42000)
3801 ]));
3802
3803 assert!(
3804 batch
3805 .column(4)
3806 .as_any()
3807 .downcast_ref::<StringArray>()
3808 .unwrap()
3809 .iter()
3810 .eq([
3811 Some("hi"),
3812 Some("ho"),
3813 Some("1970-01-01T00:00:00+02:00"),
3814 None,
3815 None,
3816 None,
3817 None,
3818 Some("true"),
3819 Some("42"),
3820 Some("1.234"),
3821 ])
3822 );
3823
3824 assert!(
3825 batch
3826 .column(5)
3827 .as_any()
3828 .downcast_ref::<StringViewArray>()
3829 .unwrap()
3830 .iter()
3831 .eq([
3832 Some("ho"),
3833 Some("1970-01-01T00:00:00+02:00"),
3834 None,
3835 None,
3836 None,
3837 None,
3838 Some("true"),
3839 Some("42"),
3840 Some("1.234"),
3841 Some("hi"),
3842 ])
3843 );
3844
3845 assert!(
3846 batch
3847 .column(6)
3848 .as_primitive::<TimestampSecondType>()
3849 .iter()
3850 .eq([
3851 Some(-7200),
3852 None,
3853 None,
3854 None,
3855 None,
3856 None,
3857 Some(42),
3858 None,
3859 None,
3860 None,
3861 ])
3862 );
3863
3864 let arrays = batch
3865 .column(7)
3866 .as_any()
3867 .downcast_ref::<ListArray>()
3868 .unwrap();
3869 assert_eq!(
3870 arrays.nulls(),
3871 Some(&NullBuffer::from(
3872 &[
3873 true, false, false, false, false, false, false, false, false, false
3874 ][..]
3875 ))
3876 );
3877 assert_eq!(arrays.offsets()[1], 3);
3878 let array_values = arrays
3879 .values()
3880 .as_any()
3881 .downcast_ref::<Int32Array>()
3882 .unwrap();
3883 assert!(array_values.iter().eq([Some(1), None, Some(3)]));
3884
3885 let maps = batch.column(8).as_any().downcast_ref::<MapArray>().unwrap();
3886 assert_eq!(
3887 maps.nulls(),
3888 Some(&NullBuffer::from(
3889 &[
3891 true, true, false, false, false, false, false, false, false, false
3892 ][..]
3893 ))
3894 );
3895 let map_keys = maps.keys().as_any().downcast_ref::<StringArray>().unwrap();
3896 assert!(map_keys.iter().eq([Some("k"), Some("a")]));
3897 let map_values = maps
3898 .values()
3899 .as_any()
3900 .downcast_ref::<StringArray>()
3901 .unwrap();
3902 assert!(map_values.iter().eq([Some("value"), Some("1")]));
3903
3904 let structs = batch
3905 .column(9)
3906 .as_any()
3907 .downcast_ref::<StructArray>()
3908 .unwrap();
3909 assert_eq!(
3910 structs.nulls(),
3911 Some(&NullBuffer::from(
3912 &[
3914 true, false, false, false, false, false, false, false, false, true
3915 ][..]
3916 ))
3917 );
3918 let struct_fields = structs
3919 .column(0)
3920 .as_any()
3921 .downcast_ref::<Int32Array>()
3922 .unwrap();
3923 assert!(struct_fields.slice(0, 2).iter().eq([Some(1), None]));
3924 }
3925
3926 #[test]
3927 fn test_type_conflict_non_nullable() {
3928 let fields = [
3929 Field::new("bool", DataType::Boolean, false),
3930 Field::new("primitive", DataType::Int32, false),
3931 Field::new("numeric", DataType::Decimal128(10, 3), false),
3932 Field::new("string", DataType::Utf8, false),
3933 Field::new("string_view", DataType::Utf8View, false),
3934 Field::new(
3935 "timestamp",
3936 DataType::Timestamp(TimeUnit::Second, None),
3937 false,
3938 ),
3939 Field::new(
3940 "array",
3941 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
3942 false,
3943 ),
3944 Field::new(
3945 "fixed_size_list",
3946 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 2),
3947 false,
3948 ),
3949 Field::new(
3950 "map",
3951 DataType::Map(
3952 Arc::new(Field::new(
3953 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3954 DataType::Struct(Fields::from(vec![
3955 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
3956 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3957 ])),
3958 false, )),
3960 false, ),
3962 false, ),
3964 Field::new(
3965 "struct",
3966 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
3967 false,
3968 ),
3969 ];
3970
3971 let json_values = vec![json!(true), json!({"a": 1})];
3973
3974 for field in fields {
3975 let mut decoder = ReaderBuilder::new_with_field(field)
3976 .with_ignore_type_conflicts(true)
3977 .build_decoder()
3978 .unwrap();
3979 decoder.serialize(&json_values).unwrap();
3980 decoder
3981 .flush()
3982 .expect_err("type conflict on non-nullable type");
3983 }
3984 }
3985
3986 #[test]
3987 fn test_ignore_type_conflicts_disabled() {
3988 let fields = [
3989 Field::new("null", DataType::Null, true),
3990 Field::new("bool", DataType::Boolean, true),
3991 Field::new("primitive", DataType::Int32, true),
3992 Field::new("numeric", DataType::Decimal128(10, 3), true),
3993 Field::new("string", DataType::Utf8, true),
3994 Field::new("string_view", DataType::Utf8View, true),
3995 Field::new(
3996 "timestamp",
3997 DataType::Timestamp(TimeUnit::Second, None),
3998 true,
3999 ),
4000 Field::new(
4001 "array",
4002 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
4003 true,
4004 ),
4005 Field::new(
4006 "fixed_size_list",
4007 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 2),
4008 true,
4009 ),
4010 Field::new(
4011 "map",
4012 DataType::Map(
4013 Arc::new(Field::new(
4014 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4015 DataType::Struct(Fields::from(vec![
4016 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4017 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
4018 ])),
4019 false, )),
4021 false, ),
4023 true, ),
4025 Field::new(
4026 "struct",
4027 DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
4028 true,
4029 ),
4030 ];
4031
4032 let json_values = vec![json!(true), json!({"a": 1})];
4034
4035 for field in fields {
4036 let mut decoder = ReaderBuilder::new_with_field(field)
4037 .build_decoder()
4038 .unwrap();
4039 decoder.serialize(&json_values).unwrap();
4040 decoder
4041 .flush()
4042 .expect_err("type conflict on non-nullable type");
4043 }
4044 }
4045
4046 #[test]
4047 fn test_read_run_end_encoded() {
4048 let buf = r#"
4049 {"a": "x"}
4050 {"a": "x"}
4051 {"a": "y"}
4052 {"a": "y"}
4053 {"a": "y"}
4054 "#;
4055
4056 let ree_type = DataType::RunEndEncoded(
4057 Arc::new(Field::new(
4058 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4059 DataType::Int32,
4060 false,
4061 )),
4062 Arc::new(Field::new(
4063 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4064 DataType::Utf8,
4065 true,
4066 )),
4067 );
4068 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
4069 let batches = do_read(buf, 1024, false, false, schema);
4070 assert_eq!(batches.len(), 1);
4071
4072 let col = batches[0].column(0);
4073 let run_array = col.as_run::<arrow_array::types::Int32Type>();
4074
4075 assert_eq!(run_array.len(), 5);
4077 assert_eq!(run_array.run_ends().values(), &[2, 5]);
4078
4079 let values = run_array.values().as_string::<i32>();
4080 assert_eq!(values.len(), 2);
4081 assert_eq!(values.value(0), "x");
4082 assert_eq!(values.value(1), "y");
4083 }
4084
4085 #[test]
4086 fn test_read_run_end_encoded_consecutive_nulls() {
4087 let buf = r#"
4088 {"a": "x"}
4089 {}
4090 {}
4091 {}
4092 {"a": "y"}
4093 "#;
4094
4095 let ree_type = DataType::RunEndEncoded(
4096 Arc::new(Field::new(
4097 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4098 DataType::Int32,
4099 false,
4100 )),
4101 Arc::new(Field::new(
4102 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4103 DataType::Utf8,
4104 true,
4105 )),
4106 );
4107 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
4108 let batches = do_read(buf, 1024, false, false, schema);
4109 assert_eq!(batches.len(), 1);
4110
4111 let col = batches[0].column(0);
4112 let run_array = col.as_run::<arrow_array::types::Int32Type>();
4113
4114 assert_eq!(run_array.len(), 5);
4116 assert_eq!(run_array.run_ends().values(), &[1, 4, 5]);
4117
4118 let values = run_array.values().as_string::<i32>();
4119 assert_eq!(values.len(), 3);
4120 assert_eq!(values.value(0), "x");
4121 assert!(values.is_null(1));
4122 assert_eq!(values.value(2), "y");
4123 }
4124
4125 #[test]
4126 fn test_read_run_end_encoded_nullability() {
4127 for field_nullable in [false, true] {
4128 for values_nullable in [false, true] {
4129 let ree_type = DataType::RunEndEncoded(
4130 Arc::new(Field::new(
4131 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4132 DataType::Int32,
4133 false,
4134 )),
4135 Arc::new(Field::new(
4136 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4137 DataType::Utf8,
4138 values_nullable,
4139 )),
4140 );
4141 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, field_nullable)]));
4142
4143 for buf in [
4144 r#"{"a": "x"}
4145 {"a": null}
4146 {"a": "y"}"#,
4147 r#"{"a": "x"}
4148 {}
4149 {"a": "y"}"#,
4150 ] {
4151 let mut decoder = ReaderBuilder::new(schema.clone()).build_decoder().unwrap();
4152 let result = decoder.decode(buf.as_bytes()).and_then(|_| decoder.flush());
4153
4154 if field_nullable && values_nullable {
4155 result.expect("REE field and values are both nullable");
4156 } else {
4157 let err = result.expect_err(
4158 "REE nulls require both the field and values to be nullable",
4159 );
4160 assert!(
4161 err.to_string().contains(
4162 "Encountered nulls in non-nullable values of RunEndEncoded"
4163 ),
4164 "unexpected error: {err}"
4165 );
4166 }
4167 }
4168 }
4169 }
4170 }
4171
4172 #[test]
4173 fn test_read_run_end_encoded_all_unique() {
4174 let buf = r#"
4175 {"a": 1}
4176 {"a": 2}
4177 {"a": 3}
4178 "#;
4179
4180 let ree_type = DataType::RunEndEncoded(
4181 Arc::new(Field::new(
4182 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4183 DataType::Int32,
4184 false,
4185 )),
4186 Arc::new(Field::new(
4187 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4188 DataType::Int32,
4189 true,
4190 )),
4191 );
4192 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
4193 let batches = do_read(buf, 1024, false, false, schema);
4194 assert_eq!(batches.len(), 1);
4195
4196 let col = batches[0].column(0);
4197 let run_array = col.as_run::<arrow_array::types::Int32Type>();
4198
4199 assert_eq!(run_array.len(), 3);
4201 assert_eq!(run_array.run_ends().values(), &[1, 2, 3]);
4202 }
4203
4204 #[test]
4205 fn test_read_run_end_encoded_int16_run_ends() {
4206 let buf = r#"
4207 {"a": "x"}
4208 {"a": "x"}
4209 {"a": "y"}
4210 "#;
4211
4212 let ree_type = DataType::RunEndEncoded(
4213 Arc::new(Field::new(
4214 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4215 DataType::Int16,
4216 false,
4217 )),
4218 Arc::new(Field::new(
4219 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4220 DataType::Utf8,
4221 true,
4222 )),
4223 );
4224 let schema = Arc::new(Schema::new(vec![Field::new("a", ree_type, true)]));
4225 let batches = do_read(buf, 1024, false, false, schema);
4226 assert_eq!(batches.len(), 1);
4227
4228 let col = batches[0].column(0);
4229 let run_array = col.as_run::<arrow_array::types::Int16Type>();
4230
4231 assert_eq!(run_array.len(), 3);
4232 assert_eq!(run_array.run_ends().values(), &[2i16, 3]);
4233 }
4234
4235 #[test]
4236 fn test_read_nested_run_end_encoded() {
4237 let buf = r#"
4238 {"a": "x"}
4239 {"a": "x"}
4240 {"a": "y"}
4241 "#;
4242
4243 let inner_type = DataType::RunEndEncoded(
4246 Arc::new(Field::new(
4247 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4248 DataType::Int64,
4249 false,
4250 )),
4251 Arc::new(Field::new(
4252 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4253 DataType::Utf8,
4254 true,
4255 )),
4256 );
4257 let outer_type = DataType::RunEndEncoded(
4258 Arc::new(Field::new(
4259 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
4260 DataType::Int64,
4261 false,
4262 )),
4263 Arc::new(Field::new(
4264 Field::REE_VALUES_FIELD_DEFAULT_NAME,
4265 inner_type,
4266 true,
4267 )),
4268 );
4269 let schema = Arc::new(Schema::new(vec![Field::new("a", outer_type, true)]));
4270 let batches = do_read(buf, 1024, false, false, schema);
4271 assert_eq!(batches.len(), 1);
4272
4273 let col = batches[0].column(0);
4274 let outer = col.as_run::<arrow_array::types::Int64Type>();
4275 assert_eq!(outer.len(), 3);
4277 assert_eq!(outer.run_ends().values(), &[2, 3]);
4278
4279 let nested = outer.values().as_run::<arrow_array::types::Int64Type>();
4280 assert_eq!(nested.len(), 2);
4282 assert_eq!(nested.run_ends().values(), &[1, 2]);
4283
4284 let nested_values = nested.values().as_string::<i32>();
4285 assert_eq!(nested_values.len(), 2);
4286 assert_eq!(nested_values.value(0), "x");
4287 assert_eq!(nested_values.value(1), "y");
4288 }
4289
4290 #[test]
4291 fn test_flatten_top_level_arrays() {
4292 let buf = r#"
4293 [
4294 {"a": 1},
4295 {"a": 2}
4296 ]
4297 {"a": 3}
4298 [{"a": 4}, {"a": 5}, {"a": 6}]
4299 "#;
4300
4301 let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
4302 let batches = do_read_config(buf, ReaderBuilder::new(schema).with_flatten(true));
4303 assert_eq!(batches.len(), 1);
4304
4305 let col = batches[0].column(0);
4306 let col = col.as_any().downcast_ref::<Int32Array>().unwrap();
4307 assert_eq!(col.len(), 6);
4308 for i in 0..6 {
4309 assert_eq!(col.value(i), (i as i32) + 1);
4310 }
4311 }
4312
4313 #[test]
4315 fn test_decoder_factory_declining_everything_is_transparent() {
4316 #[derive(Debug, Default)]
4317 struct DeclineAll {
4318 seen: Mutex<Vec<FieldRef>>,
4319 }
4320
4321 impl DecoderFactory for DeclineAll {
4322 fn make_default_decoder(
4323 &self,
4324 _ctx: &DecoderContext,
4325 field: &FieldRef,
4326 _is_nullable: bool,
4327 ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
4328 self.seen.lock().unwrap().push(field.clone());
4329 Ok(None)
4330 }
4331 }
4332
4333 let schema = Arc::new(Schema::new(vec![
4334 Field::new("a", DataType::Int32, true),
4335 Field::new("b", DataType::Utf8, true),
4336 Field::new_list("c", Field::new("item", DataType::Float64, true), true),
4337 ]));
4338 let buf = r#"{"a": 1, "b": "x", "c": [1.5, 2.5]}
4339 {"a": null, "c": []}
4340 "#;
4341
4342 let read = |factory: Option<Arc<dyn DecoderFactory>>| {
4343 let mut builder = ReaderBuilder::new(schema.clone());
4344 if let Some(factory) = factory {
4345 builder = builder.with_decoder_factory(factory);
4346 }
4347 builder
4348 .build(Cursor::new(buf.as_bytes()))
4349 .unwrap()
4350 .next()
4351 .unwrap()
4352 .unwrap()
4353 };
4354
4355 let factory = Arc::new(DeclineAll::default());
4356 assert_eq!(read(None), read(Some(factory.clone())));
4357
4358 let seen = factory.seen.lock().unwrap();
4361 assert!(
4362 matches!(seen[0].data_type(), DataType::Struct(_)),
4363 "{seen:?}"
4364 );
4365 assert_eq!(seen[0].name(), "", "root field should be nameless");
4366
4367 let by_name: Vec<(&str, &DataType)> = seen
4368 .iter()
4369 .map(|f| (f.name().as_str(), f.data_type()))
4370 .collect();
4371 for expected in [
4372 ("a", &DataType::Int32),
4373 ("b", &DataType::Utf8),
4374 ("item", &DataType::Float64),
4375 ] {
4376 assert!(
4377 by_name.contains(&expected),
4378 "{expected:?} missing from {by_name:?}"
4379 );
4380 }
4381 }
4382
4383 #[test]
4386 fn test_decoder_factory_output_is_validated() {
4387 struct Bad {
4389 wrong_type: bool,
4390 }
4391
4392 impl ArrayDecoder for Bad {
4393 fn decode(&mut self, _tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
4394 Ok(match self.wrong_type {
4395 true => Arc::new(StringViewArray::from(vec!["x"; pos.len()])) as ArrayRef,
4396 false => Arc::new(StringArray::from(Vec::<&str>::new())),
4397 })
4398 }
4399 }
4400
4401 #[derive(Debug)]
4402 struct BadFactory {
4403 wrong_type: bool,
4404 }
4405
4406 impl DecoderFactory for BadFactory {
4407 fn make_default_decoder(
4408 &self,
4409 _ctx: &DecoderContext,
4410 field: &FieldRef,
4411 _is_nullable: bool,
4412 ) -> Result<Option<Box<dyn ArrayDecoder>>, ArrowError> {
4413 match field.data_type() {
4414 DataType::Utf8 => Ok(Some(Box::new(Bad {
4415 wrong_type: self.wrong_type,
4416 }))),
4417 _ => Ok(None),
4418 }
4419 }
4420 }
4421
4422 let read = |wrong_type: bool| {
4423 let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
4424 ReaderBuilder::new(schema)
4425 .with_decoder_factory(Arc::new(BadFactory { wrong_type }))
4426 .build(Cursor::new(br#"{"s": "hello"}"#.as_slice()))
4427 .unwrap()
4428 .next()
4429 .unwrap()
4430 .unwrap_err()
4431 .to_string()
4432 };
4433
4434 let err = read(true);
4435 assert!(
4436 err.contains("returned Utf8View for a field of type Utf8"),
4437 "{err}"
4438 );
4439
4440 let err = read(false);
4441 assert!(err.contains("returned 0 values for 1 rows"), "{err}");
4442 }
4443}