arrow_integration_test/schema.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::{field_from_json, field_to_json};
19use arrow::datatypes::{Fields, Schema};
20use arrow::error::{ArrowError, Result};
21use std::collections::HashMap;
22
23/// Generate a JSON representation of the `Schema`.
24pub fn schema_to_json(schema: &Schema) -> serde_json::Value {
25 let metadata: serde_json::Map<String, serde_json::Value> = schema
26 .metadata()
27 .iter()
28 .map(|(k, v)| (k.clone(), serde_json::Value::String(v.clone())))
29 .collect();
30 serde_json::json!({
31 "fields": schema.fields().iter().map(|f| field_to_json(f.as_ref())).collect::<Vec<_>>(),
32 "metadata": metadata
33 })
34}
35
36/// Parse a `Schema` definition from a JSON representation.
37pub fn schema_from_json(json: &serde_json::Value) -> Result<Schema> {
38 use serde_json::Value;
39 match *json {
40 Value::Object(ref schema) => {
41 let fields: Fields = match schema.get("fields") {
42 Some(Value::Array(fields)) => {
43 fields.iter().map(field_from_json).collect::<Result<_>>()?
44 }
45 _ => {
46 return Err(ArrowError::ParseError(
47 "Schema fields should be an array".to_string(),
48 ));
49 }
50 };
51
52 let metadata = if let Some(value) = schema.get("metadata") {
53 from_metadata(value)?
54 } else {
55 HashMap::default()
56 };
57
58 Ok(Schema::new_with_metadata(fields, metadata))
59 }
60 _ => Err(ArrowError::ParseError(
61 "Invalid json value type for schema".to_string(),
62 )),
63 }
64}
65
66/// Parse a `metadata` definition from a JSON representation.
67/// The JSON can either be an Object or an Array of Objects.
68fn from_metadata(json: &serde_json::Value) -> Result<HashMap<String, String>> {
69 use serde_json::Value;
70 match json {
71 Value::Array(_) => {
72 let mut hashmap = HashMap::new();
73 let values: Vec<MetadataKeyValue> =
74 serde_json::from_value(json.clone()).map_err(|_| {
75 ArrowError::JsonError("Unable to parse object into key-value pair".to_string())
76 })?;
77 for meta in values {
78 hashmap.insert(meta.key.clone(), meta.value);
79 }
80 Ok(hashmap)
81 }
82 Value::Object(md) => md
83 .iter()
84 .map(|(k, v)| {
85 if let Value::String(v) = v {
86 Ok((k.clone(), v.clone()))
87 } else {
88 Err(ArrowError::ParseError(
89 "metadata `value` field must be a string".to_string(),
90 ))
91 }
92 })
93 .collect::<Result<_>>(),
94 _ => Err(ArrowError::ParseError(
95 "`metadata` field must be an object".to_string(),
96 )),
97 }
98}
99
100#[derive(serde::Deserialize)]
101struct MetadataKeyValue {
102 key: String,
103 value: String,
104}
105
106#[cfg(test)]
107mod tests {
108 use super::*;
109 use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit};
110 use serde_json::Value;
111 use std::sync::Arc;
112
113 #[test]
114 fn schema_json() {
115 // Add some custom metadata
116 let metadata = HashMap::from([("Key".to_string(), "Value".to_string())]);
117
118 let schema = Schema::new_with_metadata(
119 vec![
120 Field::new("c1", DataType::Utf8, false),
121 Field::new("c2", DataType::Binary, false),
122 Field::new("c3", DataType::FixedSizeBinary(3), false),
123 Field::new("c4", DataType::Boolean, false),
124 Field::new("c5", DataType::Date32, false),
125 Field::new("c6", DataType::Date64, false),
126 Field::new("c7", DataType::Time32(TimeUnit::Second), false),
127 Field::new("c8", DataType::Time32(TimeUnit::Millisecond), false),
128 Field::new("c9", DataType::Time32(TimeUnit::Microsecond), false),
129 Field::new("c10", DataType::Time32(TimeUnit::Nanosecond), false),
130 Field::new("c11", DataType::Time64(TimeUnit::Second), false),
131 Field::new("c12", DataType::Time64(TimeUnit::Millisecond), false),
132 Field::new("c13", DataType::Time64(TimeUnit::Microsecond), false),
133 Field::new("c14", DataType::Time64(TimeUnit::Nanosecond), false),
134 Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
135 Field::new(
136 "c16",
137 DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
138 false,
139 ),
140 Field::new(
141 "c17",
142 DataType::Timestamp(TimeUnit::Microsecond, Some("Africa/Johannesburg".into())),
143 false,
144 ),
145 Field::new(
146 "c18",
147 DataType::Timestamp(TimeUnit::Nanosecond, None),
148 false,
149 ),
150 Field::new("c19", DataType::Interval(IntervalUnit::DayTime), false),
151 Field::new("c20", DataType::Interval(IntervalUnit::YearMonth), false),
152 Field::new("c21", DataType::Interval(IntervalUnit::MonthDayNano), false),
153 Field::new(
154 "c22",
155 DataType::List(Arc::new(Field::new_list_field(DataType::Boolean, true))),
156 false,
157 ),
158 Field::new(
159 "c23",
160 DataType::FixedSizeList(
161 Arc::new(Field::new("bools", DataType::Boolean, false)),
162 5,
163 ),
164 false,
165 ),
166 Field::new(
167 "c24",
168 DataType::List(Arc::new(Field::new(
169 "inner_list",
170 DataType::List(Arc::new(Field::new(
171 "struct",
172 DataType::Struct(Fields::empty()),
173 true,
174 ))),
175 false,
176 ))),
177 true,
178 ),
179 Field::new(
180 "c25",
181 DataType::Struct(Fields::from(vec![
182 Field::new("a", DataType::Utf8, false),
183 Field::new("b", DataType::UInt16, false),
184 ])),
185 false,
186 ),
187 Field::new("c26", DataType::Interval(IntervalUnit::YearMonth), true),
188 Field::new("c27", DataType::Interval(IntervalUnit::DayTime), true),
189 Field::new("c28", DataType::Interval(IntervalUnit::MonthDayNano), true),
190 Field::new("c29", DataType::Duration(TimeUnit::Second), false),
191 Field::new("c30", DataType::Duration(TimeUnit::Millisecond), false),
192 Field::new("c31", DataType::Duration(TimeUnit::Microsecond), false),
193 Field::new("c32", DataType::Duration(TimeUnit::Nanosecond), false),
194 #[expect(deprecated)]
195 Field::new_dict(
196 "c33",
197 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
198 true,
199 123,
200 true,
201 ),
202 Field::new("c34", DataType::LargeBinary, true),
203 Field::new("c35", DataType::LargeUtf8, true),
204 Field::new(
205 "c36",
206 DataType::LargeList(Arc::new(Field::new(
207 "inner_large_list",
208 DataType::LargeList(Arc::new(Field::new(
209 "struct",
210 DataType::Struct(Fields::empty()),
211 false,
212 ))),
213 true,
214 ))),
215 true,
216 ),
217 Field::new(
218 "c37",
219 DataType::Map(
220 Arc::new(Field::new(
221 "my_entries",
222 DataType::Struct(Fields::from(vec![
223 Field::new("my_keys", DataType::Utf8, false),
224 Field::new("my_values", DataType::UInt16, true),
225 ])),
226 false,
227 )),
228 true,
229 ),
230 false,
231 ),
232 ],
233 metadata,
234 );
235
236 let expected = schema_to_json(&schema);
237 let json = r#"{
238 "fields": [
239 {
240 "name": "c1",
241 "nullable": false,
242 "type": {
243 "name": "utf8"
244 },
245 "children": []
246 },
247 {
248 "name": "c2",
249 "nullable": false,
250 "type": {
251 "name": "binary"
252 },
253 "children": []
254 },
255 {
256 "name": "c3",
257 "nullable": false,
258 "type": {
259 "name": "fixedsizebinary",
260 "byteWidth": 3
261 },
262 "children": []
263 },
264 {
265 "name": "c4",
266 "nullable": false,
267 "type": {
268 "name": "bool"
269 },
270 "children": []
271 },
272 {
273 "name": "c5",
274 "nullable": false,
275 "type": {
276 "name": "date",
277 "unit": "DAY"
278 },
279 "children": []
280 },
281 {
282 "name": "c6",
283 "nullable": false,
284 "type": {
285 "name": "date",
286 "unit": "MILLISECOND"
287 },
288 "children": []
289 },
290 {
291 "name": "c7",
292 "nullable": false,
293 "type": {
294 "name": "time",
295 "bitWidth": 32,
296 "unit": "SECOND"
297 },
298 "children": []
299 },
300 {
301 "name": "c8",
302 "nullable": false,
303 "type": {
304 "name": "time",
305 "bitWidth": 32,
306 "unit": "MILLISECOND"
307 },
308 "children": []
309 },
310 {
311 "name": "c9",
312 "nullable": false,
313 "type": {
314 "name": "time",
315 "bitWidth": 32,
316 "unit": "MICROSECOND"
317 },
318 "children": []
319 },
320 {
321 "name": "c10",
322 "nullable": false,
323 "type": {
324 "name": "time",
325 "bitWidth": 32,
326 "unit": "NANOSECOND"
327 },
328 "children": []
329 },
330 {
331 "name": "c11",
332 "nullable": false,
333 "type": {
334 "name": "time",
335 "bitWidth": 64,
336 "unit": "SECOND"
337 },
338 "children": []
339 },
340 {
341 "name": "c12",
342 "nullable": false,
343 "type": {
344 "name": "time",
345 "bitWidth": 64,
346 "unit": "MILLISECOND"
347 },
348 "children": []
349 },
350 {
351 "name": "c13",
352 "nullable": false,
353 "type": {
354 "name": "time",
355 "bitWidth": 64,
356 "unit": "MICROSECOND"
357 },
358 "children": []
359 },
360 {
361 "name": "c14",
362 "nullable": false,
363 "type": {
364 "name": "time",
365 "bitWidth": 64,
366 "unit": "NANOSECOND"
367 },
368 "children": []
369 },
370 {
371 "name": "c15",
372 "nullable": false,
373 "type": {
374 "name": "timestamp",
375 "unit": "SECOND"
376 },
377 "children": []
378 },
379 {
380 "name": "c16",
381 "nullable": false,
382 "type": {
383 "name": "timestamp",
384 "unit": "MILLISECOND",
385 "timezone": "UTC"
386 },
387 "children": []
388 },
389 {
390 "name": "c17",
391 "nullable": false,
392 "type": {
393 "name": "timestamp",
394 "unit": "MICROSECOND",
395 "timezone": "Africa/Johannesburg"
396 },
397 "children": []
398 },
399 {
400 "name": "c18",
401 "nullable": false,
402 "type": {
403 "name": "timestamp",
404 "unit": "NANOSECOND"
405 },
406 "children": []
407 },
408 {
409 "name": "c19",
410 "nullable": false,
411 "type": {
412 "name": "interval",
413 "unit": "DAY_TIME"
414 },
415 "children": []
416 },
417 {
418 "name": "c20",
419 "nullable": false,
420 "type": {
421 "name": "interval",
422 "unit": "YEAR_MONTH"
423 },
424 "children": []
425 },
426 {
427 "name": "c21",
428 "nullable": false,
429 "type": {
430 "name": "interval",
431 "unit": "MONTH_DAY_NANO"
432 },
433 "children": []
434 },
435 {
436 "name": "c22",
437 "nullable": false,
438 "type": {
439 "name": "list"
440 },
441 "children": [
442 {
443 "name": "item",
444 "nullable": true,
445 "type": {
446 "name": "bool"
447 },
448 "children": []
449 }
450 ]
451 },
452 {
453 "name": "c23",
454 "nullable": false,
455 "type": {
456 "name": "fixedsizelist",
457 "listSize": 5
458 },
459 "children": [
460 {
461 "name": "bools",
462 "nullable": false,
463 "type": {
464 "name": "bool"
465 },
466 "children": []
467 }
468 ]
469 },
470 {
471 "name": "c24",
472 "nullable": true,
473 "type": {
474 "name": "list"
475 },
476 "children": [
477 {
478 "name": "inner_list",
479 "nullable": false,
480 "type": {
481 "name": "list"
482 },
483 "children": [
484 {
485 "name": "struct",
486 "nullable": true,
487 "type": {
488 "name": "struct"
489 },
490 "children": []
491 }
492 ]
493 }
494 ]
495 },
496 {
497 "name": "c25",
498 "nullable": false,
499 "type": {
500 "name": "struct"
501 },
502 "children": [
503 {
504 "name": "a",
505 "nullable": false,
506 "type": {
507 "name": "utf8"
508 },
509 "children": []
510 },
511 {
512 "name": "b",
513 "nullable": false,
514 "type": {
515 "name": "int",
516 "bitWidth": 16,
517 "isSigned": false
518 },
519 "children": []
520 }
521 ]
522 },
523 {
524 "name": "c26",
525 "nullable": true,
526 "type": {
527 "name": "interval",
528 "unit": "YEAR_MONTH"
529 },
530 "children": []
531 },
532 {
533 "name": "c27",
534 "nullable": true,
535 "type": {
536 "name": "interval",
537 "unit": "DAY_TIME"
538 },
539 "children": []
540 },
541 {
542 "name": "c28",
543 "nullable": true,
544 "type": {
545 "name": "interval",
546 "unit": "MONTH_DAY_NANO"
547 },
548 "children": []
549 },
550 {
551 "name": "c29",
552 "nullable": false,
553 "type": {
554 "name": "duration",
555 "unit": "SECOND"
556 },
557 "children": []
558 },
559 {
560 "name": "c30",
561 "nullable": false,
562 "type": {
563 "name": "duration",
564 "unit": "MILLISECOND"
565 },
566 "children": []
567 },
568 {
569 "name": "c31",
570 "nullable": false,
571 "type": {
572 "name": "duration",
573 "unit": "MICROSECOND"
574 },
575 "children": []
576 },
577 {
578 "name": "c32",
579 "nullable": false,
580 "type": {
581 "name": "duration",
582 "unit": "NANOSECOND"
583 },
584 "children": []
585 },
586 {
587 "name": "c33",
588 "nullable": true,
589 "children": [],
590 "type": {
591 "name": "utf8"
592 },
593 "dictionary": {
594 "id": 123,
595 "indexType": {
596 "name": "int",
597 "bitWidth": 32,
598 "isSigned": true
599 },
600 "isOrdered": true
601 }
602 },
603 {
604 "name": "c34",
605 "nullable": true,
606 "type": {
607 "name": "largebinary"
608 },
609 "children": []
610 },
611 {
612 "name": "c35",
613 "nullable": true,
614 "type": {
615 "name": "largeutf8"
616 },
617 "children": []
618 },
619 {
620 "name": "c36",
621 "nullable": true,
622 "type": {
623 "name": "largelist"
624 },
625 "children": [
626 {
627 "name": "inner_large_list",
628 "nullable": true,
629 "type": {
630 "name": "largelist"
631 },
632 "children": [
633 {
634 "name": "struct",
635 "nullable": false,
636 "type": {
637 "name": "struct"
638 },
639 "children": []
640 }
641 ]
642 }
643 ]
644 },
645 {
646 "name": "c37",
647 "nullable": false,
648 "type": {
649 "name": "map",
650 "keysSorted": true
651 },
652 "children": [
653 {
654 "name": "my_entries",
655 "nullable": false,
656 "type": {
657 "name": "struct"
658 },
659 "children": [
660 {
661 "name": "my_keys",
662 "nullable": false,
663 "type": {
664 "name": "utf8"
665 },
666 "children": []
667 },
668 {
669 "name": "my_values",
670 "nullable": true,
671 "type": {
672 "name": "int",
673 "bitWidth": 16,
674 "isSigned": false
675 },
676 "children": []
677 }
678 ]
679 }
680 ]
681 }
682 ],
683 "metadata" : {
684 "Key": "Value"
685 }
686 }"#;
687 let value: Value = serde_json::from_str(json).unwrap();
688 assert_eq!(expected, value);
689
690 // convert back to a schema
691 let value: Value = serde_json::from_str(json).unwrap();
692 let schema2 = schema_from_json(&value).unwrap();
693
694 assert_eq!(schema, schema2);
695
696 // Check that empty metadata produces empty value in JSON and can be parsed
697 let json = r#"{
698 "fields": [
699 {
700 "name": "c1",
701 "nullable": false,
702 "type": {
703 "name": "utf8"
704 },
705 "children": []
706 }
707 ],
708 "metadata": {}
709 }"#;
710 let value: Value = serde_json::from_str(json).unwrap();
711 let schema = schema_from_json(&value).unwrap();
712 assert!(schema.metadata.is_empty());
713
714 // Check that metadata field is not required in the JSON.
715 let json = r#"{
716 "fields": [
717 {
718 "name": "c1",
719 "nullable": false,
720 "type": {
721 "name": "utf8"
722 },
723 "children": []
724 }
725 ]
726 }"#;
727 let value: Value = serde_json::from_str(json).unwrap();
728 let schema = schema_from_json(&value).unwrap();
729 assert!(schema.metadata.is_empty());
730 }
731}