arrow_schema/
ffi.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
18//! Contains declarations to bind to the [C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html).
19//!
20//! ```
21//! # use arrow_schema::{DataType, Field, Schema};
22//! # use arrow_schema::ffi::FFI_ArrowSchema;
23//!
24//! // Create from data type
25//! let ffi_data_type = FFI_ArrowSchema::try_from(&DataType::LargeUtf8).unwrap();
26//! let back = DataType::try_from(&ffi_data_type).unwrap();
27//! assert_eq!(back, DataType::LargeUtf8);
28//!
29//! // Create from schema
30//! let schema = Schema::new(vec![Field::new("foo", DataType::Int64, false)]);
31//! let ffi_schema = FFI_ArrowSchema::try_from(&schema).unwrap();
32//! let back = Schema::try_from(&ffi_schema).unwrap();
33//!
34//! assert_eq!(schema, back);
35//! ```
36
37use crate::{
38    ArrowError, DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, UnionFields, UnionMode,
39};
40use bitflags::bitflags;
41use std::borrow::Cow;
42use std::sync::Arc;
43use std::{
44    collections::HashMap,
45    ffi::{CStr, CString, c_char, c_void},
46};
47
48bitflags! {
49    /// Flags for [`FFI_ArrowSchema`]
50    ///
51    /// Old Workaround at <https://github.com/bitflags/bitflags/issues/356>
52    /// is no longer required as `bitflags` [fixed the issue](https://github.com/bitflags/bitflags/pull/355).
53    pub struct Flags: i64 {
54        /// Indicates that the dictionary is ordered
55        const DICTIONARY_ORDERED = 0b00000001;
56        /// Indicates that the field is nullable
57        const NULLABLE = 0b00000010;
58        /// Indicates that the map keys are sorted
59        const MAP_KEYS_SORTED = 0b00000100;
60    }
61}
62
63/// ABI-compatible struct for `ArrowSchema` from C Data Interface
64/// See <https://arrow.apache.org/docs/format/CDataInterface.html#the-arrowschema-structure>
65///
66/// ```
67/// # use arrow_schema::DataType;
68/// # use arrow_schema::ffi::FFI_ArrowSchema;
69/// fn array_schema(data_type: &DataType) -> FFI_ArrowSchema {
70///     FFI_ArrowSchema::try_from(data_type).unwrap()
71/// }
72/// ```
73///
74#[repr(C)]
75#[derive(Debug)]
76#[allow(non_camel_case_types)]
77pub struct FFI_ArrowSchema {
78    /// Null-terminated, UTF8-encoded string describing the data type
79    pub format: *const c_char,
80    /// Null-terminated, UTF8-encoded string of the field or array name
81    pub name: *const c_char,
82    /// Binary string describing the type’s metadata
83    pub metadata: *const c_char,
84    /// A bitfield of flags enriching the type description
85    /// Refer to [Arrow Flags](https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.flags)
86    pub flags: i64,
87    /// The number of children this type has
88    pub n_children: i64,
89    /// C array of pointers to each child type of this type
90    pub children: *mut *mut FFI_ArrowSchema,
91    /// Pointer to the type of dictionary values
92    pub dictionary: *mut FFI_ArrowSchema,
93    /// Pointer to a producer-provided release callback
94    pub release: Option<unsafe extern "C" fn(arg1: *mut FFI_ArrowSchema)>,
95    /// Opaque pointer to producer-provided private data
96    pub private_data: *mut c_void,
97}
98
99struct SchemaPrivateData {
100    children: Box<[*mut FFI_ArrowSchema]>,
101    dictionary: *mut FFI_ArrowSchema,
102    metadata: Option<Vec<u8>>,
103}
104
105// callback used to drop [FFI_ArrowSchema] when it is exported.
106unsafe extern "C" fn release_schema(schema: *mut FFI_ArrowSchema) {
107    if schema.is_null() {
108        return;
109    }
110    let schema = unsafe { &mut *schema };
111
112    // take ownership back to release it.
113    drop(unsafe { CString::from_raw(schema.format as *mut c_char) });
114    if !schema.name.is_null() {
115        drop(unsafe { CString::from_raw(schema.name as *mut c_char) });
116    }
117    if !schema.private_data.is_null() {
118        let private_data = unsafe { Box::from_raw(schema.private_data as *mut SchemaPrivateData) };
119        for child in private_data.children.iter() {
120            drop(unsafe { Box::from_raw(*child) })
121        }
122        if !private_data.dictionary.is_null() {
123            drop(unsafe { Box::from_raw(private_data.dictionary) });
124        }
125
126        drop(private_data);
127    }
128
129    schema.release = None;
130}
131
132impl FFI_ArrowSchema {
133    /// create a new [`FFI_ArrowSchema`]. This fails if the fields'
134    /// [`DataType`] is not supported.
135    pub fn try_new(
136        format: &str,
137        children: Vec<FFI_ArrowSchema>,
138        dictionary: Option<FFI_ArrowSchema>,
139    ) -> Result<Self, ArrowError> {
140        let mut this = Self::empty();
141
142        let children_ptr = children
143            .into_iter()
144            .map(Box::new)
145            .map(Box::into_raw)
146            .collect::<Box<_>>();
147
148        this.format = CString::new(format).unwrap().into_raw();
149        this.release = Some(release_schema);
150        this.n_children = children_ptr.len() as i64;
151
152        let dictionary_ptr = dictionary
153            .map(|d| Box::into_raw(Box::new(d)))
154            .unwrap_or(std::ptr::null_mut());
155
156        let mut private_data = Box::new(SchemaPrivateData {
157            children: children_ptr,
158            dictionary: dictionary_ptr,
159            metadata: None,
160        });
161
162        // intentionally set from private_data (see https://github.com/apache/arrow-rs/issues/580)
163        this.children = private_data.children.as_mut_ptr();
164
165        this.dictionary = dictionary_ptr;
166
167        this.private_data = Box::into_raw(private_data) as *mut c_void;
168
169        Ok(this)
170    }
171
172    /// Set the name of the schema
173    pub fn with_name(mut self, name: &str) -> Result<Self, ArrowError> {
174        self.name = CString::new(name)
175            .map_err(|e| {
176                ArrowError::CDataInterface(format!(
177                    "Null byte at position {} not allowed in name",
178                    e.nul_position()
179                ))
180            })?
181            .into_raw();
182        Ok(self)
183    }
184
185    /// Set the flags of the schema
186    pub fn with_flags(mut self, flags: Flags) -> Result<Self, ArrowError> {
187        self.flags = flags.bits();
188        Ok(self)
189    }
190
191    /// Add metadata to the schema
192    pub fn with_metadata<I, S>(mut self, metadata: I) -> Result<Self, ArrowError>
193    where
194        I: IntoIterator<Item = (S, S)>,
195        S: AsRef<str>,
196    {
197        let metadata: Vec<(S, S)> = metadata.into_iter().collect();
198        // https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.metadata
199        let new_metadata = if !metadata.is_empty() {
200            let mut metadata_serialized: Vec<u8> = Vec::new();
201            let num_entries: i32 = metadata.len().try_into().map_err(|_| {
202                ArrowError::CDataInterface(format!(
203                    "metadata can only have {} entries, but {} were provided",
204                    i32::MAX,
205                    metadata.len()
206                ))
207            })?;
208            metadata_serialized.extend(num_entries.to_ne_bytes());
209
210            for (key, value) in metadata.into_iter() {
211                let key_len: i32 = key.as_ref().len().try_into().map_err(|_| {
212                    ArrowError::CDataInterface(format!(
213                        "metadata key can only have {} bytes, but {} were provided",
214                        i32::MAX,
215                        key.as_ref().len()
216                    ))
217                })?;
218                let value_len: i32 = value.as_ref().len().try_into().map_err(|_| {
219                    ArrowError::CDataInterface(format!(
220                        "metadata value can only have {} bytes, but {} were provided",
221                        i32::MAX,
222                        value.as_ref().len()
223                    ))
224                })?;
225
226                metadata_serialized.extend(key_len.to_ne_bytes());
227                metadata_serialized.extend_from_slice(key.as_ref().as_bytes());
228                metadata_serialized.extend(value_len.to_ne_bytes());
229                metadata_serialized.extend_from_slice(value.as_ref().as_bytes());
230            }
231
232            self.metadata = metadata_serialized.as_ptr() as *const c_char;
233            Some(metadata_serialized)
234        } else {
235            self.metadata = std::ptr::null_mut();
236            None
237        };
238
239        unsafe {
240            let mut private_data = Box::from_raw(self.private_data as *mut SchemaPrivateData);
241            private_data.metadata = new_metadata;
242            self.private_data = Box::into_raw(private_data) as *mut c_void;
243        }
244
245        Ok(self)
246    }
247
248    /// Takes ownership of the pointed to [`FFI_ArrowSchema`]
249    ///
250    /// This acts to [move] the data out of `schema`, setting the release callback to NULL
251    ///
252    /// # Safety
253    ///
254    /// * `schema` must be [valid] for reads and writes
255    /// * `schema` must be properly aligned
256    /// * `schema` must point to a properly initialized value of [`FFI_ArrowSchema`]
257    ///
258    /// [move]: https://arrow.apache.org/docs/format/CDataInterface.html#moving-an-array
259    /// [valid]: https://doc.rust-lang.org/std/ptr/index.html#safety
260    pub unsafe fn from_raw(schema: *mut FFI_ArrowSchema) -> Self {
261        unsafe { std::ptr::replace(schema, Self::empty()) }
262    }
263
264    /// Create an empty [`FFI_ArrowSchema`]
265    pub fn empty() -> Self {
266        Self {
267            format: std::ptr::null_mut(),
268            name: std::ptr::null_mut(),
269            metadata: std::ptr::null_mut(),
270            flags: 0,
271            n_children: 0,
272            children: std::ptr::null_mut(),
273            dictionary: std::ptr::null_mut(),
274            release: None,
275            private_data: std::ptr::null_mut(),
276        }
277    }
278
279    /// Returns the format of this schema.
280    pub fn format(&self) -> &str {
281        assert!(!self.format.is_null());
282        // safe because the lifetime of `self.format` equals `self`
283        unsafe { CStr::from_ptr(self.format) }
284            .to_str()
285            .expect("The external API has a non-utf8 as format")
286    }
287
288    /// Returns the name of this schema.
289    pub fn name(&self) -> Option<&str> {
290        if self.name.is_null() {
291            None
292        } else {
293            // safe because the lifetime of `self.name` equals `self`
294            Some(
295                unsafe { CStr::from_ptr(self.name) }
296                    .to_str()
297                    .expect("The external API has a non-utf8 as name"),
298            )
299        }
300    }
301
302    /// Returns the flags of this schema.
303    pub fn flags(&self) -> Option<Flags> {
304        Flags::from_bits(self.flags)
305    }
306
307    /// Returns the child of this schema at `index`.
308    ///
309    /// # Panics
310    ///
311    /// Panics if `index` is greater than or equal to the number of children.
312    ///
313    /// This is to make sure that the unsafe acces to raw pointer is sound.
314    pub fn child(&self, index: usize) -> &Self {
315        assert!(index < self.n_children as usize);
316        unsafe { self.children.add(index).as_ref().unwrap().as_ref().unwrap() }
317    }
318
319    /// Returns an iterator to the schema's children.
320    pub fn children(&self) -> impl Iterator<Item = &Self> {
321        (0..self.n_children as usize).map(move |i| self.child(i))
322    }
323
324    /// Returns if the field is semantically nullable,
325    /// regardless of whether it actually has null values.
326    pub fn nullable(&self) -> bool {
327        (self.flags / 2) & 1 == 1
328    }
329
330    /// Returns the reference to the underlying dictionary of the schema.
331    /// Check [ArrowSchema.dictionary](https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.dictionary).
332    ///
333    /// This must be `Some` if the schema represents a dictionary-encoded type, `None` otherwise.
334    pub fn dictionary(&self) -> Option<&Self> {
335        unsafe { self.dictionary.as_ref() }
336    }
337
338    /// For map types, returns whether the keys within each map value are sorted.
339    ///
340    /// Refer to [Arrow Flags](https://arrow.apache.org/docs/format/CDataInterface.html#c.ArrowSchema.flags)
341    pub fn map_keys_sorted(&self) -> bool {
342        self.flags & 0b00000100 != 0
343    }
344
345    /// For dictionary-encoded types, returns whether the ordering of dictionary indices is semantically meaningful.
346    pub fn dictionary_ordered(&self) -> bool {
347        self.flags & 0b00000001 != 0
348    }
349
350    /// Returns the metadata in the schema as `Key-Value` pairs
351    pub fn metadata(&self) -> Result<HashMap<String, String>, ArrowError> {
352        if self.metadata.is_null() {
353            Ok(HashMap::new())
354        } else {
355            let mut pos = 0;
356
357            // On some platforms, c_char = u8, and on some, c_char = i8. Where c_char = u8, clippy
358            // wants to complain that we're casting to the same type, but if we remove the cast,
359            // this will fail to compile on the other platforms. So we must allow it.
360            #[allow(clippy::unnecessary_cast)]
361            let buffer: *const u8 = self.metadata as *const u8;
362
363            fn next_four_bytes(buffer: *const u8, pos: &mut isize) -> [u8; 4] {
364                let out = unsafe {
365                    [
366                        *buffer.offset(*pos),
367                        *buffer.offset(*pos + 1),
368                        *buffer.offset(*pos + 2),
369                        *buffer.offset(*pos + 3),
370                    ]
371                };
372                *pos += 4;
373                out
374            }
375
376            fn next_n_bytes(buffer: *const u8, pos: &mut isize, n: i32) -> &[u8] {
377                let out = unsafe {
378                    std::slice::from_raw_parts(buffer.offset(*pos), n.try_into().unwrap())
379                };
380                *pos += isize::try_from(n).unwrap();
381                out
382            }
383
384            let num_entries = i32::from_ne_bytes(next_four_bytes(buffer, &mut pos));
385            if num_entries < 0 {
386                return Err(ArrowError::CDataInterface(
387                    "Negative number of metadata entries".to_string(),
388                ));
389            }
390
391            let mut metadata =
392                HashMap::with_capacity(num_entries.try_into().expect("Too many metadata entries"));
393
394            for _ in 0..num_entries {
395                let key_length = i32::from_ne_bytes(next_four_bytes(buffer, &mut pos));
396                if key_length < 0 {
397                    return Err(ArrowError::CDataInterface(
398                        "Negative key length in metadata".to_string(),
399                    ));
400                }
401                let key = String::from_utf8(next_n_bytes(buffer, &mut pos, key_length).to_vec())?;
402                let value_length = i32::from_ne_bytes(next_four_bytes(buffer, &mut pos));
403                if value_length < 0 {
404                    return Err(ArrowError::CDataInterface(
405                        "Negative value length in metadata".to_string(),
406                    ));
407                }
408                let value =
409                    String::from_utf8(next_n_bytes(buffer, &mut pos, value_length).to_vec())?;
410                metadata.insert(key, value);
411            }
412
413            Ok(metadata)
414        }
415    }
416}
417
418impl Drop for FFI_ArrowSchema {
419    fn drop(&mut self) {
420        match self.release {
421            None => (),
422            Some(release) => unsafe { release(self) },
423        };
424    }
425}
426
427unsafe impl Send for FFI_ArrowSchema {}
428
429impl TryFrom<&FFI_ArrowSchema> for DataType {
430    type Error = ArrowError;
431
432    /// See [CDataInterface docs](https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings)
433    fn try_from(c_schema: &FFI_ArrowSchema) -> Result<Self, ArrowError> {
434        let mut dtype = match c_schema.format() {
435            "n" => DataType::Null,
436            "b" => DataType::Boolean,
437            "c" => DataType::Int8,
438            "C" => DataType::UInt8,
439            "s" => DataType::Int16,
440            "S" => DataType::UInt16,
441            "i" => DataType::Int32,
442            "I" => DataType::UInt32,
443            "l" => DataType::Int64,
444            "L" => DataType::UInt64,
445            "e" => DataType::Float16,
446            "f" => DataType::Float32,
447            "g" => DataType::Float64,
448            "vz" => DataType::BinaryView,
449            "z" => DataType::Binary,
450            "Z" => DataType::LargeBinary,
451            "vu" => DataType::Utf8View,
452            "u" => DataType::Utf8,
453            "U" => DataType::LargeUtf8,
454            "tdD" => DataType::Date32,
455            "tdm" => DataType::Date64,
456            "tts" => DataType::Time32(TimeUnit::Second),
457            "ttm" => DataType::Time32(TimeUnit::Millisecond),
458            "ttu" => DataType::Time64(TimeUnit::Microsecond),
459            "ttn" => DataType::Time64(TimeUnit::Nanosecond),
460            "tDs" => DataType::Duration(TimeUnit::Second),
461            "tDm" => DataType::Duration(TimeUnit::Millisecond),
462            "tDu" => DataType::Duration(TimeUnit::Microsecond),
463            "tDn" => DataType::Duration(TimeUnit::Nanosecond),
464            "tiM" => DataType::Interval(IntervalUnit::YearMonth),
465            "tiD" => DataType::Interval(IntervalUnit::DayTime),
466            "tin" => DataType::Interval(IntervalUnit::MonthDayNano),
467            "+l" => {
468                let c_child = c_schema.child(0);
469                DataType::List(Arc::new(Field::try_from(c_child)?))
470            }
471            "+L" => {
472                let c_child = c_schema.child(0);
473                DataType::LargeList(Arc::new(Field::try_from(c_child)?))
474            }
475            "+vl" => {
476                let c_child = c_schema.child(0);
477                DataType::ListView(Arc::new(Field::try_from(c_child)?))
478            }
479            "+vL" => {
480                let c_child = c_schema.child(0);
481                DataType::LargeListView(Arc::new(Field::try_from(c_child)?))
482            }
483            "+s" => {
484                let fields = c_schema.children().map(Field::try_from);
485                DataType::Struct(fields.collect::<Result<_, ArrowError>>()?)
486            }
487            "+m" => {
488                let c_child = c_schema.child(0);
489                let map_keys_sorted = c_schema.map_keys_sorted();
490                DataType::Map(Arc::new(Field::try_from(c_child)?), map_keys_sorted)
491            }
492            "+r" => {
493                let c_run_ends = c_schema.child(0);
494                let c_values = c_schema.child(1);
495                DataType::RunEndEncoded(
496                    Arc::new(Field::try_from(c_run_ends)?),
497                    Arc::new(Field::try_from(c_values)?),
498                )
499            }
500            // Parametrized types, requiring string parse
501            other => {
502                match other.splitn(2, ':').collect::<Vec<&str>>().as_slice() {
503                    // FixedSizeBinary type in format "w:num_bytes"
504                    ["w", num_bytes] => {
505                        let parsed_num_bytes = num_bytes.parse::<i32>().map_err(|_| {
506                            ArrowError::CDataInterface(
507                                "FixedSizeBinary requires an integer parameter representing number of bytes per element".to_string())
508                        })?;
509                        DataType::FixedSizeBinary(parsed_num_bytes)
510                    }
511                    // FixedSizeList type in format "+w:num_elems"
512                    ["+w", num_elems] => {
513                        let c_child = c_schema.child(0);
514                        let parsed_num_elems = num_elems.parse::<i32>().map_err(|_| {
515                            ArrowError::CDataInterface(
516                                "The FixedSizeList type requires an integer parameter representing number of elements per list".to_string())
517                        })?;
518                        DataType::FixedSizeList(
519                            Arc::new(Field::try_from(c_child)?),
520                            parsed_num_elems,
521                        )
522                    }
523                    // Decimal types in format "d:precision,scale" or "d:precision,scale,bitWidth"
524                    ["d", extra] => match extra.splitn(3, ',').collect::<Vec<&str>>().as_slice() {
525                        [precision, scale] => {
526                            let parsed_precision = precision.parse::<u8>().map_err(|_| {
527                                ArrowError::CDataInterface(
528                                    "The decimal type requires an integer precision".to_string(),
529                                )
530                            })?;
531                            let parsed_scale = scale.parse::<i8>().map_err(|_| {
532                                ArrowError::CDataInterface(
533                                    "The decimal type requires an integer scale".to_string(),
534                                )
535                            })?;
536                            DataType::Decimal128(parsed_precision, parsed_scale)
537                        }
538                        [precision, scale, bits] => {
539                            let parsed_precision = precision.parse::<u8>().map_err(|_| {
540                                ArrowError::CDataInterface(
541                                    "The decimal type requires an integer precision".to_string(),
542                                )
543                            })?;
544                            let parsed_scale = scale.parse::<i8>().map_err(|_| {
545                                ArrowError::CDataInterface(
546                                    "The decimal type requires an integer scale".to_string(),
547                                )
548                            })?;
549                            match *bits {
550                                    "32" => DataType::Decimal32(parsed_precision, parsed_scale),
551                                    "64" => DataType::Decimal64(parsed_precision, parsed_scale),
552                                    "128" => DataType::Decimal128(parsed_precision, parsed_scale),
553                                    "256" => DataType::Decimal256(parsed_precision, parsed_scale),
554                                    _ => return Err(ArrowError::CDataInterface("Only 32/64/128/256 bit wide decimals are supported in the Rust implementation".to_string())),
555                                }
556                        }
557                        _ => {
558                            return Err(ArrowError::CDataInterface(format!(
559                                "The decimal pattern \"d:{extra:?}\" is not supported in the Rust implementation"
560                            )));
561                        }
562                    },
563                    // DenseUnion
564                    ["+ud", extra] => {
565                        let type_ids = extra
566                            .split(',')
567                            .map(|t| {
568                                t.parse::<i8>().map_err(|_| {
569                                    ArrowError::CDataInterface(
570                                        "The Union type requires an integer type id".to_string(),
571                                    )
572                                })
573                            })
574                            .collect::<Result<Vec<_>, ArrowError>>()?;
575                        let mut fields = Vec::with_capacity(type_ids.len());
576                        for idx in 0..c_schema.n_children {
577                            let c_child = c_schema.child(idx as usize);
578                            let field = Field::try_from(c_child)?;
579                            fields.push(field);
580                        }
581
582                        if fields.len() != type_ids.len() {
583                            return Err(ArrowError::CDataInterface(
584                                "The Union type requires same number of fields and type ids"
585                                    .to_string(),
586                            ));
587                        }
588
589                        DataType::Union(UnionFields::try_new(type_ids, fields)?, UnionMode::Dense)
590                    }
591                    // SparseUnion
592                    ["+us", extra] => {
593                        let type_ids = extra
594                            .split(',')
595                            .map(|t| {
596                                t.parse::<i8>().map_err(|_| {
597                                    ArrowError::CDataInterface(
598                                        "The Union type requires an integer type id".to_string(),
599                                    )
600                                })
601                            })
602                            .collect::<Result<Vec<_>, ArrowError>>()?;
603                        let mut fields = Vec::with_capacity(type_ids.len());
604                        for idx in 0..c_schema.n_children {
605                            let c_child = c_schema.child(idx as usize);
606                            let field = Field::try_from(c_child)?;
607                            fields.push(field);
608                        }
609
610                        if fields.len() != type_ids.len() {
611                            return Err(ArrowError::CDataInterface(
612                                "The Union type requires same number of fields and type ids"
613                                    .to_string(),
614                            ));
615                        }
616
617                        DataType::Union(UnionFields::try_new(type_ids, fields)?, UnionMode::Sparse)
618                    }
619
620                    // Timestamps in format "tts:" and "tts:America/New_York" for no timezones and timezones resp.
621                    ["tss", ""] => DataType::Timestamp(TimeUnit::Second, None),
622                    ["tsm", ""] => DataType::Timestamp(TimeUnit::Millisecond, None),
623                    ["tsu", ""] => DataType::Timestamp(TimeUnit::Microsecond, None),
624                    ["tsn", ""] => DataType::Timestamp(TimeUnit::Nanosecond, None),
625                    ["tss", tz] => DataType::Timestamp(TimeUnit::Second, Some(Arc::from(*tz))),
626                    ["tsm", tz] => DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from(*tz))),
627                    ["tsu", tz] => DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from(*tz))),
628                    ["tsn", tz] => DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from(*tz))),
629                    _ => {
630                        return Err(ArrowError::CDataInterface(format!(
631                            "The datatype \"{other:?}\" is still not supported in Rust implementation"
632                        )));
633                    }
634                }
635            }
636        };
637
638        if let Some(dict_schema) = c_schema.dictionary() {
639            let value_type = Self::try_from(dict_schema)?;
640            dtype = DataType::Dictionary(Box::new(dtype), Box::new(value_type));
641        }
642
643        Ok(dtype)
644    }
645}
646
647impl TryFrom<&FFI_ArrowSchema> for Field {
648    type Error = ArrowError;
649
650    fn try_from(c_schema: &FFI_ArrowSchema) -> Result<Self, ArrowError> {
651        let dtype = DataType::try_from(c_schema)?;
652        let mut field = Field::new(c_schema.name().unwrap_or(""), dtype, c_schema.nullable());
653        field.set_metadata(c_schema.metadata()?);
654        Ok(field)
655    }
656}
657
658impl TryFrom<&FFI_ArrowSchema> for Schema {
659    type Error = ArrowError;
660
661    fn try_from(c_schema: &FFI_ArrowSchema) -> Result<Self, ArrowError> {
662        // interpret it as a struct type then extract its fields
663        let dtype = DataType::try_from(c_schema)?;
664        if let DataType::Struct(fields) = dtype {
665            Ok(Schema::new(fields).with_metadata(c_schema.metadata()?))
666        } else {
667            Err(ArrowError::CDataInterface(
668                "Unable to interpret C data struct as a Schema".to_string(),
669            ))
670        }
671    }
672}
673
674impl TryFrom<&DataType> for FFI_ArrowSchema {
675    type Error = ArrowError;
676
677    /// See [CDataInterface docs](https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings)
678    fn try_from(dtype: &DataType) -> Result<Self, ArrowError> {
679        let format = get_format_string(dtype)?;
680        // allocate and hold the children
681        let children = match dtype {
682            DataType::List(child)
683            | DataType::LargeList(child)
684            | DataType::ListView(child)
685            | DataType::LargeListView(child)
686            | DataType::FixedSizeList(child, _)
687            | DataType::Map(child, _) => {
688                vec![FFI_ArrowSchema::try_from(child.as_ref())?]
689            }
690            DataType::Union(fields, _) => fields
691                .iter()
692                .map(|(_, f)| f.as_ref().try_into())
693                .collect::<Result<Vec<_>, ArrowError>>()?,
694            DataType::Struct(fields) => fields
695                .iter()
696                .map(FFI_ArrowSchema::try_from)
697                .collect::<Result<Vec<_>, ArrowError>>()?,
698            DataType::RunEndEncoded(run_ends, values) => vec![
699                FFI_ArrowSchema::try_from(run_ends.as_ref())?,
700                FFI_ArrowSchema::try_from(values.as_ref())?,
701            ],
702            _ => vec![],
703        };
704        let dictionary = if let DataType::Dictionary(_, value_data_type) = dtype {
705            Some(Self::try_from(value_data_type.as_ref())?)
706        } else {
707            None
708        };
709
710        let flags = match dtype {
711            DataType::Map(_, true) => Flags::MAP_KEYS_SORTED,
712            _ => Flags::empty(),
713        };
714
715        FFI_ArrowSchema::try_new(&format, children, dictionary)?.with_flags(flags)
716    }
717}
718
719fn get_format_string(dtype: &DataType) -> Result<Cow<'static, str>, ArrowError> {
720    match dtype {
721        DataType::Null => Ok("n".into()),
722        DataType::Boolean => Ok("b".into()),
723        DataType::Int8 => Ok("c".into()),
724        DataType::UInt8 => Ok("C".into()),
725        DataType::Int16 => Ok("s".into()),
726        DataType::UInt16 => Ok("S".into()),
727        DataType::Int32 => Ok("i".into()),
728        DataType::UInt32 => Ok("I".into()),
729        DataType::Int64 => Ok("l".into()),
730        DataType::UInt64 => Ok("L".into()),
731        DataType::Float16 => Ok("e".into()),
732        DataType::Float32 => Ok("f".into()),
733        DataType::Float64 => Ok("g".into()),
734        DataType::BinaryView => Ok("vz".into()),
735        DataType::Binary => Ok("z".into()),
736        DataType::LargeBinary => Ok("Z".into()),
737        DataType::Utf8View => Ok("vu".into()),
738        DataType::Utf8 => Ok("u".into()),
739        DataType::LargeUtf8 => Ok("U".into()),
740        DataType::FixedSizeBinary(num_bytes) => Ok(Cow::Owned(format!("w:{num_bytes}"))),
741        DataType::FixedSizeList(_, num_elems) => Ok(Cow::Owned(format!("+w:{num_elems}"))),
742        DataType::Decimal32(precision, scale) => {
743            Ok(Cow::Owned(format!("d:{precision},{scale},32")))
744        }
745        DataType::Decimal64(precision, scale) => {
746            Ok(Cow::Owned(format!("d:{precision},{scale},64")))
747        }
748        DataType::Decimal128(precision, scale) => Ok(Cow::Owned(format!("d:{precision},{scale}"))),
749        DataType::Decimal256(precision, scale) => {
750            Ok(Cow::Owned(format!("d:{precision},{scale},256")))
751        }
752        DataType::Date32 => Ok("tdD".into()),
753        DataType::Date64 => Ok("tdm".into()),
754        DataType::Time32(TimeUnit::Second) => Ok("tts".into()),
755        DataType::Time32(TimeUnit::Millisecond) => Ok("ttm".into()),
756        DataType::Time64(TimeUnit::Microsecond) => Ok("ttu".into()),
757        DataType::Time64(TimeUnit::Nanosecond) => Ok("ttn".into()),
758        DataType::Timestamp(TimeUnit::Second, None) => Ok("tss:".into()),
759        DataType::Timestamp(TimeUnit::Millisecond, None) => Ok("tsm:".into()),
760        DataType::Timestamp(TimeUnit::Microsecond, None) => Ok("tsu:".into()),
761        DataType::Timestamp(TimeUnit::Nanosecond, None) => Ok("tsn:".into()),
762        DataType::Timestamp(TimeUnit::Second, Some(tz)) => Ok(Cow::Owned(format!("tss:{tz}"))),
763        DataType::Timestamp(TimeUnit::Millisecond, Some(tz)) => Ok(Cow::Owned(format!("tsm:{tz}"))),
764        DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => Ok(Cow::Owned(format!("tsu:{tz}"))),
765        DataType::Timestamp(TimeUnit::Nanosecond, Some(tz)) => Ok(Cow::Owned(format!("tsn:{tz}"))),
766        DataType::Duration(TimeUnit::Second) => Ok("tDs".into()),
767        DataType::Duration(TimeUnit::Millisecond) => Ok("tDm".into()),
768        DataType::Duration(TimeUnit::Microsecond) => Ok("tDu".into()),
769        DataType::Duration(TimeUnit::Nanosecond) => Ok("tDn".into()),
770        DataType::Interval(IntervalUnit::YearMonth) => Ok("tiM".into()),
771        DataType::Interval(IntervalUnit::DayTime) => Ok("tiD".into()),
772        DataType::Interval(IntervalUnit::MonthDayNano) => Ok("tin".into()),
773        DataType::List(_) => Ok("+l".into()),
774        DataType::LargeList(_) => Ok("+L".into()),
775        DataType::ListView(_) => Ok("+vl".into()),
776        DataType::LargeListView(_) => Ok("+vL".into()),
777        DataType::Struct(_) => Ok("+s".into()),
778        DataType::Map(_, _) => Ok("+m".into()),
779        DataType::RunEndEncoded(_, _) => Ok("+r".into()),
780        DataType::Dictionary(key_data_type, _) => get_format_string(key_data_type),
781        DataType::Union(fields, mode) => {
782            let formats = fields
783                .iter()
784                .map(|(t, _)| t.to_string())
785                .collect::<Vec<_>>();
786            match mode {
787                UnionMode::Dense => Ok(Cow::Owned(format!("{}:{}", "+ud", formats.join(",")))),
788                UnionMode::Sparse => Ok(Cow::Owned(format!("{}:{}", "+us", formats.join(",")))),
789            }
790        }
791        other => Err(ArrowError::CDataInterface(format!(
792            "The datatype \"{other:?}\" is still not supported in Rust implementation"
793        ))),
794    }
795}
796
797impl TryFrom<&FieldRef> for FFI_ArrowSchema {
798    type Error = ArrowError;
799
800    fn try_from(value: &FieldRef) -> Result<Self, Self::Error> {
801        value.as_ref().try_into()
802    }
803}
804
805impl TryFrom<&Field> for FFI_ArrowSchema {
806    type Error = ArrowError;
807
808    fn try_from(field: &Field) -> Result<Self, ArrowError> {
809        let mut flags = if field.is_nullable() {
810            Flags::NULLABLE
811        } else {
812            Flags::empty()
813        };
814
815        if let Some(true) = field.dict_is_ordered() {
816            flags |= Flags::DICTIONARY_ORDERED;
817        }
818
819        FFI_ArrowSchema::try_from(field.data_type())?
820            .with_name(field.name())?
821            .with_flags(flags)?
822            .with_metadata(field.metadata())
823    }
824}
825
826impl TryFrom<&Schema> for FFI_ArrowSchema {
827    type Error = ArrowError;
828
829    fn try_from(schema: &Schema) -> Result<Self, ArrowError> {
830        let dtype = DataType::Struct(schema.fields().clone());
831        let c_schema = FFI_ArrowSchema::try_from(&dtype)?.with_metadata(&schema.metadata)?;
832        Ok(c_schema)
833    }
834}
835
836impl TryFrom<DataType> for FFI_ArrowSchema {
837    type Error = ArrowError;
838
839    fn try_from(dtype: DataType) -> Result<Self, ArrowError> {
840        FFI_ArrowSchema::try_from(&dtype)
841    }
842}
843
844impl TryFrom<Field> for FFI_ArrowSchema {
845    type Error = ArrowError;
846
847    fn try_from(field: Field) -> Result<Self, ArrowError> {
848        FFI_ArrowSchema::try_from(&field)
849    }
850}
851
852impl TryFrom<Schema> for FFI_ArrowSchema {
853    type Error = ArrowError;
854
855    fn try_from(schema: Schema) -> Result<Self, ArrowError> {
856        FFI_ArrowSchema::try_from(&schema)
857    }
858}
859
860#[cfg(test)]
861mod tests {
862    use super::*;
863    use crate::Fields;
864
865    fn round_trip_type(dtype: DataType) {
866        let c_schema = FFI_ArrowSchema::try_from(&dtype).unwrap();
867        let restored = DataType::try_from(&c_schema).unwrap();
868        assert_eq!(restored, dtype);
869    }
870
871    fn round_trip_field(field: Field) {
872        let c_schema = FFI_ArrowSchema::try_from(&field).unwrap();
873        let restored = Field::try_from(&c_schema).unwrap();
874        assert_eq!(restored, field);
875    }
876
877    fn round_trip_schema(schema: Schema) {
878        let c_schema = FFI_ArrowSchema::try_from(&schema).unwrap();
879        let restored = Schema::try_from(&c_schema).unwrap();
880        assert_eq!(restored, schema);
881    }
882
883    #[test]
884    fn test_type() {
885        round_trip_type(DataType::Int64);
886        round_trip_type(DataType::UInt64);
887        round_trip_type(DataType::Float64);
888        round_trip_type(DataType::Date64);
889        round_trip_type(DataType::Time64(TimeUnit::Nanosecond));
890        round_trip_type(DataType::FixedSizeBinary(12));
891        round_trip_type(DataType::FixedSizeList(
892            Arc::new(Field::new("a", DataType::Int64, false)),
893            5,
894        ));
895        round_trip_type(DataType::Utf8);
896        round_trip_type(DataType::Utf8View);
897        round_trip_type(DataType::BinaryView);
898        round_trip_type(DataType::Binary);
899        round_trip_type(DataType::LargeBinary);
900        round_trip_type(DataType::List(Arc::new(Field::new(
901            "a",
902            DataType::Int16,
903            false,
904        ))));
905        round_trip_type(DataType::ListView(Arc::new(Field::new(
906            "a",
907            DataType::Int16,
908            false,
909        ))));
910        round_trip_type(DataType::LargeListView(Arc::new(Field::new(
911            "a",
912            DataType::Int16,
913            false,
914        ))));
915        round_trip_type(DataType::Struct(Fields::from(vec![Field::new(
916            "a",
917            DataType::Utf8,
918            true,
919        )])));
920        round_trip_type(DataType::RunEndEncoded(
921            Arc::new(Field::new("run_ends", DataType::Int32, false)),
922            Arc::new(Field::new("values", DataType::Binary, true)),
923        ));
924    }
925
926    #[test]
927    fn test_field() {
928        let dtype = DataType::Struct(vec![Field::new("a", DataType::Utf8, true)].into());
929        round_trip_field(Field::new("test", dtype, true));
930    }
931
932    #[test]
933    fn test_schema() {
934        let schema = Schema::new(vec![
935            Field::new("name", DataType::Utf8, false),
936            Field::new("address", DataType::Utf8, false),
937            Field::new("priority", DataType::UInt8, false),
938        ])
939        .with_metadata([("hello".to_string(), "world".to_string())].into());
940
941        round_trip_schema(schema);
942
943        // test that we can interpret struct types as schema
944        let dtype = DataType::Struct(Fields::from(vec![
945            Field::new("a", DataType::Utf8, true),
946            Field::new("b", DataType::Int16, false),
947        ]));
948        let c_schema = FFI_ArrowSchema::try_from(&dtype).unwrap();
949        let schema = Schema::try_from(&c_schema).unwrap();
950        assert_eq!(schema.fields().len(), 2);
951
952        // test that we assert the input type
953        let c_schema = FFI_ArrowSchema::try_from(&DataType::Float64).unwrap();
954        let result = Schema::try_from(&c_schema);
955        assert!(result.is_err());
956    }
957
958    #[test]
959    fn test_map_keys_sorted() {
960        let keys = Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false);
961        let values = Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::UInt32, false);
962        let entry_struct = DataType::Struct(vec![keys, values].into());
963
964        // Construct a map array from the above two
965        let map_data_type = DataType::Map(
966            Arc::new(Field::new(
967                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
968                entry_struct,
969                false,
970            )),
971            true,
972        );
973
974        let arrow_schema = FFI_ArrowSchema::try_from(map_data_type).unwrap();
975        assert!(arrow_schema.map_keys_sorted());
976    }
977
978    #[test]
979    fn test_dictionary_ordered() {
980        #[allow(deprecated)]
981        let schema = Schema::new(vec![Field::new_dict(
982            "dict",
983            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
984            false,
985            0,
986            true,
987        )]);
988
989        let arrow_schema = FFI_ArrowSchema::try_from(schema).unwrap();
990        assert!(arrow_schema.child(0).dictionary_ordered());
991    }
992
993    #[test]
994    fn test_set_field_metadata() {
995        let metadata_cases: Vec<HashMap<String, String>> = vec![
996            [].into(),
997            [("key".to_string(), "value".to_string())].into(),
998            [
999                ("key".to_string(), "".to_string()),
1000                ("ascii123".to_string(), "你好".to_string()),
1001                ("".to_string(), "value".to_string()),
1002            ]
1003            .into(),
1004        ];
1005
1006        let mut schema = FFI_ArrowSchema::try_new("b", vec![], None)
1007            .unwrap()
1008            .with_name("test")
1009            .unwrap();
1010
1011        for metadata in metadata_cases {
1012            schema = schema.with_metadata(&metadata).unwrap();
1013            let field = Field::try_from(&schema).unwrap();
1014            assert_eq!(field.metadata(), &metadata);
1015        }
1016    }
1017
1018    #[test]
1019    fn test_name_with_null_byte() {
1020        let schema = FFI_ArrowSchema::try_new("i", vec![], None).unwrap();
1021        assert!(schema.with_name("ab\0cd").is_err());
1022    }
1023
1024    #[test]
1025    fn test_import_field_with_null_name() {
1026        let dtype = DataType::Int16;
1027        let c_schema = FFI_ArrowSchema::try_from(&dtype).unwrap();
1028        assert!(c_schema.name().is_none());
1029        let field = Field::try_from(&c_schema).unwrap();
1030        assert_eq!(field.name(), "");
1031    }
1032}
arrow_schema/ffi.rs

arrow_schema/
ffi.rs
