Skip to main content

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}