arrow_ipc/
convert.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Utilities for converting between IPC types and native Arrow types
6
7use arrow_buffer::Buffer;
8use arrow_schema::*;
9use core::panic;
10use flatbuffers::{
11    FlatBufferBuilder, ForwardsUOffset, UnionWIPOffset, Vector, Verifiable, Verifier,
12    VerifierOptions, WIPOffset,
13};
14use std::collections::HashMap;
15use std::fmt::{Debug, Formatter};
16use std::sync::Arc;
17
18use crate::writer::DictionaryTracker;
19use crate::{CONTINUATION_MARKER, KeyValue, Message};
20use DataType::*;
21
22/// Low level Arrow [Schema] to IPC bytes converter
23///
24/// See also [`try_fb_to_schema`] for the reverse operation
25///
26/// # Example
27/// ```
28/// # use arrow_ipc::convert::{try_fb_to_schema, IpcSchemaEncoder};
29/// # use arrow_ipc::root_as_schema;
30/// # use arrow_ipc::writer::DictionaryTracker;
31/// # use arrow_schema::{DataType, Field, Schema};
32/// // given an arrow schema to serialize
33/// let schema = Schema::new(vec![
34///    Field::new("a", DataType::Int32, false),
35/// ]);
36///
37/// // Use a dictionary tracker to track dictionary id if needed
38///  let mut dictionary_tracker = DictionaryTracker::new(true);
39/// // create a FlatBuffersBuilder that contains the encoded bytes
40///  let fb = IpcSchemaEncoder::new()
41///    .with_dictionary_tracker(&mut dictionary_tracker)
42///    .schema_to_fb(&schema);
43///
44/// // the bytes are in `fb.finished_data()`
45/// let ipc_bytes = fb.finished_data();
46///
47///  // convert the IPC bytes back to an Arrow schema
48///  let ipc_schema = root_as_schema(ipc_bytes).unwrap();
49///  let schema2 = try_fb_to_schema(ipc_schema).unwrap();
50/// assert_eq!(schema, schema2);
51/// ```
52#[derive(Debug)]
53pub struct IpcSchemaEncoder<'a> {
54    dictionary_tracker: Option<&'a mut DictionaryTracker>,
55}
56
57impl Default for IpcSchemaEncoder<'_> {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl<'a> IpcSchemaEncoder<'a> {
64    /// Create a new schema encoder
65    pub fn new() -> IpcSchemaEncoder<'a> {
66        IpcSchemaEncoder {
67            dictionary_tracker: None,
68        }
69    }
70
71    /// Specify a dictionary tracker to use
72    pub fn with_dictionary_tracker(
73        mut self,
74        dictionary_tracker: &'a mut DictionaryTracker,
75    ) -> Self {
76        self.dictionary_tracker = Some(dictionary_tracker);
77        self
78    }
79
80    /// Serialize a schema in IPC format, returning a completed [`FlatBufferBuilder`]
81    ///
82    /// Note: Call [`FlatBufferBuilder::finished_data`] to get the serialized bytes
83    pub fn schema_to_fb<'b>(&mut self, schema: &Schema) -> FlatBufferBuilder<'b> {
84        let mut fbb = FlatBufferBuilder::new();
85
86        let root = self.schema_to_fb_offset(&mut fbb, schema);
87
88        fbb.finish(root, None);
89
90        fbb
91    }
92
93    /// Serialize a schema to an in progress [`FlatBufferBuilder`], returning the in progress offset.
94    pub fn schema_to_fb_offset<'b>(
95        &mut self,
96        fbb: &mut FlatBufferBuilder<'b>,
97        schema: &Schema,
98    ) -> WIPOffset<crate::Schema<'b>> {
99        let fields = schema
100            .fields()
101            .iter()
102            .map(|field| build_field(fbb, &mut self.dictionary_tracker, field))
103            .collect::<Vec<_>>();
104        let fb_field_list = fbb.create_vector(&fields);
105
106        let fb_metadata_list =
107            (!schema.metadata().is_empty()).then(|| metadata_to_fb(fbb, schema.metadata()));
108
109        let mut builder = crate::SchemaBuilder::new(fbb);
110        builder.add_fields(fb_field_list);
111        if let Some(fb_metadata_list) = fb_metadata_list {
112            builder.add_custom_metadata(fb_metadata_list);
113        }
114        builder.finish()
115    }
116}
117
118/// Push a key-value metadata into a FlatBufferBuilder and return [WIPOffset]
119pub fn metadata_to_fb<'a>(
120    fbb: &mut FlatBufferBuilder<'a>,
121    metadata: &Metadata,
122) -> WIPOffset<Vector<'a, ForwardsUOffset<KeyValue<'a>>>> {
123    // `Metadata` iterates in deterministic (sorted) key order
124    let custom_metadata = metadata
125        .iter()
126        .map(|(k, v)| {
127            let fb_key_name = fbb.create_string(k);
128            let fb_val_name = fbb.create_string(v);
129
130            let mut kv_builder = crate::KeyValueBuilder::new(fbb);
131            kv_builder.add_key(fb_key_name);
132            kv_builder.add_value(fb_val_name);
133            kv_builder.finish()
134        })
135        .collect::<Vec<_>>();
136    fbb.create_vector(&custom_metadata)
137}
138
139/// Adds a [Schema] to a flatbuffer and returns the offset
140pub fn schema_to_fb_offset<'a>(
141    fbb: &mut FlatBufferBuilder<'a>,
142    schema: &Schema,
143) -> WIPOffset<crate::Schema<'a>> {
144    IpcSchemaEncoder::new().schema_to_fb_offset(fbb, schema)
145}
146
147/// Convert an IPC Field to Arrow Field
148///
149/// This panics on malformed input; every reader path uses the fallible
150/// conversion instead. kept for backwards compatibility only.
151impl From<crate::Field<'_>> for Field {
152    fn from(field: crate::Field) -> Field {
153        try_field_from(field).expect("invalid IPC field")
154    }
155}
156
157/// Convert an IPC Field to Arrow Field
158fn try_field_from(field: crate::Field) -> Result<Field, ArrowError> {
159    let arrow_field = if let Some(dictionary) = field.dictionary() {
160        #[expect(deprecated)]
161        Field::new_dict(
162            field.name().unwrap_or_default(),
163            get_data_type(field, true)?,
164            field.nullable(),
165            dictionary.id(),
166            dictionary.isOrdered(),
167        )
168    } else {
169        Field::new(
170            field.name().unwrap_or_default(),
171            get_data_type(field, true)?,
172            field.nullable(),
173        )
174    };
175
176    let mut metadata_map = HashMap::default();
177    if let Some(list) = field.custom_metadata() {
178        for kv in list {
179            if let (Some(k), Some(v)) = (kv.key(), kv.value()) {
180                metadata_map.insert(k.to_string(), v.to_string());
181            }
182        }
183    }
184
185    Ok(arrow_field.with_metadata(metadata_map))
186}
187
188/// Deserialize an ipc [`crate::Schema`] from flat buffers to an arrow [Schema].
189#[deprecated(since = "60.0.0", note = "Use `try_fb_to_schema` instead")]
190pub fn fb_to_schema(fb: crate::Schema) -> Schema {
191    try_fb_to_schema(fb).expect("invalid IPC schema")
192}
193
194/// Deserialize an ipc [`crate::Schema`] from flat buffers to an arrow [Schema].
195///
196/// Returns an error on schema messages that the flatbuffer
197/// verifier accepts but that are not valid Arrow.
198pub fn try_fb_to_schema(fb: crate::Schema) -> Result<Schema, ArrowError> {
199    let mut fields: Vec<Field> = vec![];
200    let c_fields = fb
201        .fields()
202        .ok_or_else(|| ArrowError::ParseError("IPC schema has no fields".to_string()))?;
203    let len = c_fields.len();
204    for i in 0..len {
205        let c_field: crate::Field = c_fields.get(i);
206        match c_field.type_type() {
207            crate::Type::Decimal if fb.endianness() == crate::Endianness::Big => {
208                return Err(ArrowError::ParseError(
209                    "Big Endian is not supported for Decimal!".to_string(),
210                ));
211            }
212            _ => (),
213        }
214        fields.push(try_field_from(c_field)?);
215    }
216
217    let mut metadata: HashMap<String, String> = HashMap::default();
218    if let Some(md_fields) = fb.custom_metadata() {
219        let len = md_fields.len();
220        for i in 0..len {
221            let kv = md_fields.get(i);
222            let k_str = kv.key();
223            let v_str = kv.value();
224            if let Some(k) = k_str
225                && let Some(v) = v_str
226            {
227                metadata.insert(k.to_string(), v.to_string());
228            }
229        }
230    }
231    Ok(Schema::new_with_metadata(fields, metadata))
232}
233
234/// Try deserialize flat buffer format bytes into a schema
235pub fn try_schema_from_flatbuffer_bytes(bytes: &[u8]) -> Result<Schema, ArrowError> {
236    if let Ok(ipc) = crate::root_as_message(bytes) {
237        if let Some(schema) = ipc.header_as_schema() {
238            try_fb_to_schema(schema)
239        } else {
240            Err(ArrowError::ParseError(
241                "Unable to get head as schema".to_string(),
242            ))
243        }
244    } else {
245        Err(ArrowError::ParseError(
246            "Unable to get root as message".to_string(),
247        ))
248    }
249}
250
251/// Try deserialize the IPC format bytes into a schema
252pub fn try_schema_from_ipc_buffer(buffer: &[u8]) -> Result<Schema, ArrowError> {
253    // There are two protocol types: https://issues.apache.org/jira/browse/ARROW-6313
254    // The original protocol is:
255    //   4 bytes - the byte length of the payload
256    //   a flatbuffer Message whose header is the Schema
257    // The latest version of protocol is:
258    // The schema of the dataset in its IPC form:
259    //   4 bytes - an optional IPC_CONTINUATION_TOKEN prefix
260    //   4 bytes - the byte length of the payload
261    //   a flatbuffer Message whose header is the Schema
262    if buffer.len() < 4 {
263        return Err(ArrowError::ParseError(
264            "The buffer length is less than 4 and missing the continuation marker or length of buffer".to_string()
265        ));
266    }
267
268    let rest = if buffer[..4] == CONTINUATION_MARKER {
269        &buffer[4..]
270    } else {
271        buffer
272    };
273    let Some((len, buffer)) = rest.split_first_chunk::<4>() else {
274        return Err(ArrowError::ParseError(
275            "The buffer length is less than 8 and missing the length of buffer".to_string(),
276        ));
277    };
278
279    let len = i32::from_le_bytes(*len);
280    if len < 0 {
281        return Err(ArrowError::ParseError(format!(
282            "The encapsulated message's reported length is negative ({len})"
283        )));
284    }
285
286    if buffer.len() < len as usize {
287        let actual_len = buffer.len();
288        return Err(ArrowError::ParseError(format!(
289            "The buffer length ({actual_len}) is less than the encapsulated message's reported length ({len})"
290        )));
291    }
292
293    let msg = crate::root_as_message(buffer)
294        .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))?;
295    let ipc_schema = msg.header_as_schema().ok_or_else(|| {
296        ArrowError::ParseError("Unable to convert flight info to a schema".to_string())
297    })?;
298    try_fb_to_schema(ipc_schema)
299}
300
301/// Get the Arrow data type from the flatbuffer Field table
302pub(crate) fn get_data_type(
303    field: crate::Field,
304    may_be_dictionary: bool,
305) -> Result<DataType, ArrowError> {
306    // helper: the flatbuffer verifier accepts a Field whose `type_type` tag does
307    // not match the union value present, so every `type_as_*` accessor can return
308    // None. surface that as an error instead of unwrapping.
309    fn type_err(what: &str) -> ArrowError {
310        ArrowError::ParseError(format!("IPC schema field is missing its {what} type"))
311    }
312
313    if let Some(dictionary) = field.dictionary()
314        && may_be_dictionary
315    {
316        let int = dictionary
317            .indexType()
318            .ok_or_else(|| type_err("dictionary index"))?;
319        let index_type = match (int.bitWidth(), int.is_signed()) {
320            (8, true) => DataType::Int8,
321            (8, false) => DataType::UInt8,
322            (16, true) => DataType::Int16,
323            (16, false) => DataType::UInt16,
324            (32, true) => DataType::Int32,
325            (32, false) => DataType::UInt32,
326            (64, true) => DataType::Int64,
327            (64, false) => DataType::UInt64,
328            _ => {
329                return Err(ArrowError::ParseError(format!(
330                    "Index type with bit width of {} and signed of {} not supported",
331                    int.bitWidth(),
332                    int.is_signed()
333                )));
334            }
335        };
336        return Ok(DataType::Dictionary(
337            Box::new(index_type),
338            Box::new(get_data_type(field, false)?),
339        ));
340    }
341
342    let data_type = match field.type_type() {
343        crate::Type::Null => DataType::Null,
344        crate::Type::Bool => DataType::Boolean,
345        crate::Type::Int => {
346            let int = field.type_as_int().ok_or_else(|| type_err("int"))?;
347            match (int.bitWidth(), int.is_signed()) {
348                (8, true) => DataType::Int8,
349                (8, false) => DataType::UInt8,
350                (16, true) => DataType::Int16,
351                (16, false) => DataType::UInt16,
352                (32, true) => DataType::Int32,
353                (32, false) => DataType::UInt32,
354                (64, true) => DataType::Int64,
355                (64, false) => DataType::UInt64,
356                z => {
357                    return Err(ArrowError::ParseError(format!(
358                        "Int type with bit width of {} and signed of {} not supported",
359                        z.0, z.1
360                    )));
361                }
362            }
363        }
364        crate::Type::Binary => DataType::Binary,
365        crate::Type::BinaryView => DataType::BinaryView,
366        crate::Type::LargeBinary => DataType::LargeBinary,
367        crate::Type::Utf8 => DataType::Utf8,
368        crate::Type::Utf8View => DataType::Utf8View,
369        crate::Type::LargeUtf8 => DataType::LargeUtf8,
370        crate::Type::FixedSizeBinary => {
371            let fsb = field
372                .type_as_fixed_size_binary()
373                .ok_or_else(|| type_err("fixed size binary"))?;
374            let byte_width = fsb.byteWidth();
375            if byte_width < 0 {
376                return Err(ArrowError::ParseError(format!(
377                    "FixedSizeBinary byte width must be non-negative, got {byte_width}"
378                )));
379            }
380            DataType::FixedSizeBinary(byte_width)
381        }
382        crate::Type::FloatingPoint => {
383            let float = field
384                .type_as_floating_point()
385                .ok_or_else(|| type_err("floating point"))?;
386            match float.precision() {
387                crate::Precision::HALF => DataType::Float16,
388                crate::Precision::SINGLE => DataType::Float32,
389                crate::Precision::DOUBLE => DataType::Float64,
390                z => {
391                    return Err(ArrowError::ParseError(format!(
392                        "FloatingPoint type with precision of {z:?} not supported"
393                    )));
394                }
395            }
396        }
397        crate::Type::Date => {
398            let date = field.type_as_date().ok_or_else(|| type_err("date"))?;
399            match date.unit() {
400                crate::DateUnit::DAY => DataType::Date32,
401                crate::DateUnit::MILLISECOND => DataType::Date64,
402                z => {
403                    return Err(ArrowError::ParseError(format!(
404                        "Date type with unit of {z:?} not supported"
405                    )));
406                }
407            }
408        }
409        crate::Type::Time => {
410            let time = field.type_as_time().ok_or_else(|| type_err("time"))?;
411            match (time.bitWidth(), time.unit()) {
412                (32, crate::TimeUnit::SECOND) => DataType::Time32(TimeUnit::Second),
413                (32, crate::TimeUnit::MILLISECOND) => DataType::Time32(TimeUnit::Millisecond),
414                (64, crate::TimeUnit::MICROSECOND) => DataType::Time64(TimeUnit::Microsecond),
415                (64, crate::TimeUnit::NANOSECOND) => DataType::Time64(TimeUnit::Nanosecond),
416                z => {
417                    return Err(ArrowError::ParseError(format!(
418                        "Time type with bit width of {} and unit of {:?} not supported",
419                        z.0, z.1
420                    )));
421                }
422            }
423        }
424        crate::Type::Timestamp => {
425            let timestamp = field
426                .type_as_timestamp()
427                .ok_or_else(|| type_err("timestamp"))?;
428            let timezone: Option<_> = timestamp.timezone().map(|tz| tz.into());
429            match timestamp.unit() {
430                crate::TimeUnit::SECOND => DataType::Timestamp(TimeUnit::Second, timezone),
431                crate::TimeUnit::MILLISECOND => {
432                    DataType::Timestamp(TimeUnit::Millisecond, timezone)
433                }
434                crate::TimeUnit::MICROSECOND => {
435                    DataType::Timestamp(TimeUnit::Microsecond, timezone)
436                }
437                crate::TimeUnit::NANOSECOND => DataType::Timestamp(TimeUnit::Nanosecond, timezone),
438                z => {
439                    return Err(ArrowError::ParseError(format!(
440                        "Timestamp type with unit of {z:?} not supported"
441                    )));
442                }
443            }
444        }
445        crate::Type::Interval => {
446            let interval = field
447                .type_as_interval()
448                .ok_or_else(|| type_err("interval"))?;
449            match interval.unit() {
450                crate::IntervalUnit::YEAR_MONTH => DataType::Interval(IntervalUnit::YearMonth),
451                crate::IntervalUnit::DAY_TIME => DataType::Interval(IntervalUnit::DayTime),
452                crate::IntervalUnit::MONTH_DAY_NANO => {
453                    DataType::Interval(IntervalUnit::MonthDayNano)
454                }
455                z => {
456                    return Err(ArrowError::ParseError(format!(
457                        "Interval type with unit of {z:?} unsupported"
458                    )));
459                }
460            }
461        }
462        crate::Type::Duration => {
463            let duration = field
464                .type_as_duration()
465                .ok_or_else(|| type_err("duration"))?;
466            match duration.unit() {
467                crate::TimeUnit::SECOND => DataType::Duration(TimeUnit::Second),
468                crate::TimeUnit::MILLISECOND => DataType::Duration(TimeUnit::Millisecond),
469                crate::TimeUnit::MICROSECOND => DataType::Duration(TimeUnit::Microsecond),
470                crate::TimeUnit::NANOSECOND => DataType::Duration(TimeUnit::Nanosecond),
471                z => {
472                    return Err(ArrowError::ParseError(format!(
473                        "Duration type with unit of {z:?} unsupported"
474                    )));
475                }
476            }
477        }
478        crate::Type::List => {
479            let child = single_child(&field, "list")?;
480            DataType::List(Arc::new(try_field_from(child)?))
481        }
482        crate::Type::LargeList => {
483            let child = single_child(&field, "large list")?;
484            DataType::LargeList(Arc::new(try_field_from(child)?))
485        }
486        crate::Type::ListView => {
487            let child = single_child(&field, "listview")?;
488            DataType::ListView(Arc::new(try_field_from(child)?))
489        }
490        crate::Type::LargeListView => {
491            let child = single_child(&field, "large listview")?;
492            DataType::LargeListView(Arc::new(try_field_from(child)?))
493        }
494        crate::Type::FixedSizeList => {
495            let child = single_child(&field, "list")?;
496            let fsl = field
497                .type_as_fixed_size_list()
498                .ok_or_else(|| type_err("fixed size list"))?;
499            DataType::FixedSizeList(Arc::new(try_field_from(child)?), fsl.listSize())
500        }
501        crate::Type::Struct_ => {
502            let fields = match field.children() {
503                Some(children) => children
504                    .iter()
505                    .map(try_field_from)
506                    .collect::<Result<Vec<_>, _>>()?
507                    .into(),
508                None => Fields::empty(),
509            };
510            DataType::Struct(fields)
511        }
512        crate::Type::RunEndEncoded => {
513            let children = field
514                .children()
515                .ok_or_else(|| ArrowError::ParseError("IPC list has no children".to_string()))?;
516            if children.len() != 2 {
517                return Err(ArrowError::ParseError(format!(
518                    "RunEndEncoded type should have exactly two children. Found {}",
519                    children.len()
520                )));
521            }
522            let run_ends_field = try_field_from(children.get(0))?;
523            let values_field = try_field_from(children.get(1))?;
524            DataType::RunEndEncoded(Arc::new(run_ends_field), Arc::new(values_field))
525        }
526        crate::Type::Map => {
527            let map = field.type_as_map().ok_or_else(|| type_err("map"))?;
528            let child = single_child(&field, "map")?;
529            DataType::Map(Arc::new(try_field_from(child)?), map.keysSorted())
530        }
531        crate::Type::Decimal => {
532            let fsb = field.type_as_decimal().ok_or_else(|| type_err("decimal"))?;
533            let bit_width = fsb.bitWidth();
534            let precision: u8 = fsb.precision().try_into().map_err(|_| {
535                ArrowError::ParseError(format!("Invalid decimal precision {}", fsb.precision()))
536            })?;
537            let scale: i8 = fsb.scale().try_into().map_err(|_| {
538                ArrowError::ParseError(format!("Invalid decimal scale {}", fsb.scale()))
539            })?;
540            match bit_width {
541                32 => DataType::Decimal32(precision, scale),
542                64 => DataType::Decimal64(precision, scale),
543                128 => DataType::Decimal128(precision, scale),
544                256 => DataType::Decimal256(precision, scale),
545                _ => {
546                    return Err(ArrowError::ParseError(format!(
547                        "Unexpected decimal bit width {bit_width}"
548                    )));
549                }
550            }
551        }
552        crate::Type::Union => {
553            let union = field.type_as_union().ok_or_else(|| type_err("union"))?;
554
555            let union_mode = match union.mode() {
556                crate::UnionMode::Dense => UnionMode::Dense,
557                crate::UnionMode::Sparse => UnionMode::Sparse,
558                mode => {
559                    return Err(ArrowError::ParseError(format!(
560                        "Unexpected union mode: {mode:?}"
561                    )));
562                }
563            };
564
565            let mut fields = vec![];
566            if let Some(children) = field.children() {
567                for i in 0..children.len() {
568                    fields.push(try_field_from(children.get(i))?);
569                }
570            }
571
572            let fields = match union.typeIds() {
573                None => UnionFields::try_from_fields(fields)?,
574                Some(ids) => UnionFields::try_new(ids.iter().map(|i| i as i8), fields)?,
575            };
576
577            DataType::Union(fields, union_mode)
578        }
579        t => return Err(ArrowError::ParseError(format!("Type {t:?} not supported"))),
580    };
581    Ok(data_type)
582}
583
584/// Return the single child of a flatbuffer Field, erroring if it does not have
585/// exactly one.
586fn single_child<'a>(field: &crate::Field<'a>, kind: &str) -> Result<crate::Field<'a>, ArrowError> {
587    let children = field
588        .children()
589        .ok_or_else(|| ArrowError::ParseError(format!("IPC {kind} has no children")))?;
590    if children.len() != 1 {
591        return Err(ArrowError::ParseError(format!(
592            "expect a {kind} to have one child"
593        )));
594    }
595    Ok(children.get(0))
596}
597
598pub(crate) struct FBFieldType<'b> {
599    pub(crate) type_type: crate::Type,
600    pub(crate) type_: WIPOffset<UnionWIPOffset>,
601    pub(crate) children: Option<WIPOffset<Vector<'b, ForwardsUOffset<crate::Field<'b>>>>>,
602}
603
604/// Create an IPC Field from an Arrow Field
605pub(crate) fn build_field<'a>(
606    fbb: &mut FlatBufferBuilder<'a>,
607    dictionary_tracker: &mut Option<&mut DictionaryTracker>,
608    field: &Field,
609) -> WIPOffset<crate::Field<'a>> {
610    // Optional custom metadata.
611    let fb_metadata = if field.metadata().is_empty() {
612        None
613    } else {
614        Some(metadata_to_fb(fbb, field.metadata()))
615    };
616
617    let fb_field_name = fbb.create_string(field.name().as_str());
618    let field_type = get_fb_field_type(field.data_type(), dictionary_tracker, fbb);
619
620    let fb_dictionary = if let Dictionary(index_type, _) = field.data_type() {
621        match dictionary_tracker {
622            Some(tracker) => Some(get_fb_dictionary(
623                index_type,
624                tracker.next_dict_id(),
625                field
626                    .dict_is_ordered()
627                    .expect("All Dictionary types have `dict_is_ordered`"),
628                fbb,
629            )),
630            None => panic!("IPC must no longer be used without dictionary tracker"),
631        }
632    } else {
633        None
634    };
635
636    let mut field_builder = crate::FieldBuilder::new(fbb);
637    field_builder.add_name(fb_field_name);
638    if let Some(dictionary) = fb_dictionary {
639        field_builder.add_dictionary(dictionary)
640    }
641    field_builder.add_type_type(field_type.type_type);
642    field_builder.add_nullable(field.is_nullable());
643    match field_type.children {
644        None => {}
645        Some(children) => field_builder.add_children(children),
646    }
647    field_builder.add_type_(field_type.type_);
648
649    if let Some(fb_metadata) = fb_metadata {
650        field_builder.add_custom_metadata(fb_metadata);
651    }
652
653    field_builder.finish()
654}
655
656/// Get the IPC type of a data type
657pub(crate) fn get_fb_field_type<'a>(
658    data_type: &DataType,
659    dictionary_tracker: &mut Option<&mut DictionaryTracker>,
660    fbb: &mut FlatBufferBuilder<'a>,
661) -> FBFieldType<'a> {
662    // some IPC implementations expect an empty list for child data, instead of a null value.
663    // An empty field list is thus returned for primitive types
664    let empty_fields: Vec<WIPOffset<crate::Field>> = vec![];
665    match data_type {
666        Null => FBFieldType {
667            type_type: crate::Type::Null,
668            type_: crate::NullBuilder::new(fbb).finish().as_union_value(),
669            children: Some(fbb.create_vector(&empty_fields[..])),
670        },
671        Boolean => FBFieldType {
672            type_type: crate::Type::Bool,
673            type_: crate::BoolBuilder::new(fbb).finish().as_union_value(),
674            children: Some(fbb.create_vector(&empty_fields[..])),
675        },
676        UInt8 | UInt16 | UInt32 | UInt64 => {
677            let children = fbb.create_vector(&empty_fields[..]);
678            let mut builder = crate::IntBuilder::new(fbb);
679            builder.add_is_signed(false);
680            match data_type {
681                UInt8 => builder.add_bitWidth(8),
682                UInt16 => builder.add_bitWidth(16),
683                UInt32 => builder.add_bitWidth(32),
684                UInt64 => builder.add_bitWidth(64),
685                _ => {}
686            }
687            FBFieldType {
688                type_type: crate::Type::Int,
689                type_: builder.finish().as_union_value(),
690                children: Some(children),
691            }
692        }
693        Int8 | Int16 | Int32 | Int64 => {
694            let children = fbb.create_vector(&empty_fields[..]);
695            let mut builder = crate::IntBuilder::new(fbb);
696            builder.add_is_signed(true);
697            match data_type {
698                Int8 => builder.add_bitWidth(8),
699                Int16 => builder.add_bitWidth(16),
700                Int32 => builder.add_bitWidth(32),
701                Int64 => builder.add_bitWidth(64),
702                _ => {}
703            }
704            FBFieldType {
705                type_type: crate::Type::Int,
706                type_: builder.finish().as_union_value(),
707                children: Some(children),
708            }
709        }
710        Float16 | Float32 | Float64 => {
711            let children = fbb.create_vector(&empty_fields[..]);
712            let mut builder = crate::FloatingPointBuilder::new(fbb);
713            match data_type {
714                Float16 => builder.add_precision(crate::Precision::HALF),
715                Float32 => builder.add_precision(crate::Precision::SINGLE),
716                Float64 => builder.add_precision(crate::Precision::DOUBLE),
717                _ => {}
718            }
719            FBFieldType {
720                type_type: crate::Type::FloatingPoint,
721                type_: builder.finish().as_union_value(),
722                children: Some(children),
723            }
724        }
725        Binary => FBFieldType {
726            type_type: crate::Type::Binary,
727            type_: crate::BinaryBuilder::new(fbb).finish().as_union_value(),
728            children: Some(fbb.create_vector(&empty_fields[..])),
729        },
730        LargeBinary => FBFieldType {
731            type_type: crate::Type::LargeBinary,
732            type_: crate::LargeBinaryBuilder::new(fbb)
733                .finish()
734                .as_union_value(),
735            children: Some(fbb.create_vector(&empty_fields[..])),
736        },
737        BinaryView => FBFieldType {
738            type_type: crate::Type::BinaryView,
739            type_: crate::BinaryViewBuilder::new(fbb).finish().as_union_value(),
740            children: Some(fbb.create_vector(&empty_fields[..])),
741        },
742        Utf8View => FBFieldType {
743            type_type: crate::Type::Utf8View,
744            type_: crate::Utf8ViewBuilder::new(fbb).finish().as_union_value(),
745            children: Some(fbb.create_vector(&empty_fields[..])),
746        },
747        Utf8 => FBFieldType {
748            type_type: crate::Type::Utf8,
749            type_: crate::Utf8Builder::new(fbb).finish().as_union_value(),
750            children: Some(fbb.create_vector(&empty_fields[..])),
751        },
752        LargeUtf8 => FBFieldType {
753            type_type: crate::Type::LargeUtf8,
754            type_: crate::LargeUtf8Builder::new(fbb).finish().as_union_value(),
755            children: Some(fbb.create_vector(&empty_fields[..])),
756        },
757        FixedSizeBinary(len) => {
758            let mut builder = crate::FixedSizeBinaryBuilder::new(fbb);
759            builder.add_byteWidth(*len);
760            FBFieldType {
761                type_type: crate::Type::FixedSizeBinary,
762                type_: builder.finish().as_union_value(),
763                children: Some(fbb.create_vector(&empty_fields[..])),
764            }
765        }
766        Date32 => {
767            let mut builder = crate::DateBuilder::new(fbb);
768            builder.add_unit(crate::DateUnit::DAY);
769            FBFieldType {
770                type_type: crate::Type::Date,
771                type_: builder.finish().as_union_value(),
772                children: Some(fbb.create_vector(&empty_fields[..])),
773            }
774        }
775        Date64 => {
776            let mut builder = crate::DateBuilder::new(fbb);
777            builder.add_unit(crate::DateUnit::MILLISECOND);
778            FBFieldType {
779                type_type: crate::Type::Date,
780                type_: builder.finish().as_union_value(),
781                children: Some(fbb.create_vector(&empty_fields[..])),
782            }
783        }
784        Time32(unit) | Time64(unit) => {
785            let mut builder = crate::TimeBuilder::new(fbb);
786            match unit {
787                TimeUnit::Second => {
788                    builder.add_bitWidth(32);
789                    builder.add_unit(crate::TimeUnit::SECOND);
790                }
791                TimeUnit::Millisecond => {
792                    builder.add_bitWidth(32);
793                    builder.add_unit(crate::TimeUnit::MILLISECOND);
794                }
795                TimeUnit::Microsecond => {
796                    builder.add_bitWidth(64);
797                    builder.add_unit(crate::TimeUnit::MICROSECOND);
798                }
799                TimeUnit::Nanosecond => {
800                    builder.add_bitWidth(64);
801                    builder.add_unit(crate::TimeUnit::NANOSECOND);
802                }
803            }
804            FBFieldType {
805                type_type: crate::Type::Time,
806                type_: builder.finish().as_union_value(),
807                children: Some(fbb.create_vector(&empty_fields[..])),
808            }
809        }
810        Timestamp(unit, tz) => {
811            let tz = tz.as_deref().unwrap_or_default();
812            let tz_str = fbb.create_string(tz);
813            let mut builder = crate::TimestampBuilder::new(fbb);
814            let time_unit = match unit {
815                TimeUnit::Second => crate::TimeUnit::SECOND,
816                TimeUnit::Millisecond => crate::TimeUnit::MILLISECOND,
817                TimeUnit::Microsecond => crate::TimeUnit::MICROSECOND,
818                TimeUnit::Nanosecond => crate::TimeUnit::NANOSECOND,
819            };
820            builder.add_unit(time_unit);
821            if !tz.is_empty() {
822                builder.add_timezone(tz_str);
823            }
824            FBFieldType {
825                type_type: crate::Type::Timestamp,
826                type_: builder.finish().as_union_value(),
827                children: Some(fbb.create_vector(&empty_fields[..])),
828            }
829        }
830        Interval(unit) => {
831            let mut builder = crate::IntervalBuilder::new(fbb);
832            let interval_unit = match unit {
833                IntervalUnit::YearMonth => crate::IntervalUnit::YEAR_MONTH,
834                IntervalUnit::DayTime => crate::IntervalUnit::DAY_TIME,
835                IntervalUnit::MonthDayNano => crate::IntervalUnit::MONTH_DAY_NANO,
836            };
837            builder.add_unit(interval_unit);
838            FBFieldType {
839                type_type: crate::Type::Interval,
840                type_: builder.finish().as_union_value(),
841                children: Some(fbb.create_vector(&empty_fields[..])),
842            }
843        }
844        Duration(unit) => {
845            let mut builder = crate::DurationBuilder::new(fbb);
846            let time_unit = match unit {
847                TimeUnit::Second => crate::TimeUnit::SECOND,
848                TimeUnit::Millisecond => crate::TimeUnit::MILLISECOND,
849                TimeUnit::Microsecond => crate::TimeUnit::MICROSECOND,
850                TimeUnit::Nanosecond => crate::TimeUnit::NANOSECOND,
851            };
852            builder.add_unit(time_unit);
853            FBFieldType {
854                type_type: crate::Type::Duration,
855                type_: builder.finish().as_union_value(),
856                children: Some(fbb.create_vector(&empty_fields[..])),
857            }
858        }
859        List(list_type) => {
860            let child = build_field(fbb, dictionary_tracker, list_type);
861            FBFieldType {
862                type_type: crate::Type::List,
863                type_: crate::ListBuilder::new(fbb).finish().as_union_value(),
864                children: Some(fbb.create_vector(&[child])),
865            }
866        }
867        ListView(list_type) => {
868            let child = build_field(fbb, dictionary_tracker, list_type);
869            FBFieldType {
870                type_type: crate::Type::ListView,
871                type_: crate::ListViewBuilder::new(fbb).finish().as_union_value(),
872                children: Some(fbb.create_vector(&[child])),
873            }
874        }
875        LargeListView(list_type) => {
876            let child = build_field(fbb, dictionary_tracker, list_type);
877            FBFieldType {
878                type_type: crate::Type::LargeListView,
879                type_: crate::LargeListViewBuilder::new(fbb)
880                    .finish()
881                    .as_union_value(),
882                children: Some(fbb.create_vector(&[child])),
883            }
884        }
885        LargeList(list_type) => {
886            let child = build_field(fbb, dictionary_tracker, list_type);
887            FBFieldType {
888                type_type: crate::Type::LargeList,
889                type_: crate::LargeListBuilder::new(fbb).finish().as_union_value(),
890                children: Some(fbb.create_vector(&[child])),
891            }
892        }
893        FixedSizeList(list_type, len) => {
894            let child = build_field(fbb, dictionary_tracker, list_type);
895            let mut builder = crate::FixedSizeListBuilder::new(fbb);
896            builder.add_listSize(*len);
897            FBFieldType {
898                type_type: crate::Type::FixedSizeList,
899                type_: builder.finish().as_union_value(),
900                children: Some(fbb.create_vector(&[child])),
901            }
902        }
903        Struct(fields) => {
904            // struct's fields are children
905            let mut children = vec![];
906            for field in fields {
907                children.push(build_field(fbb, dictionary_tracker, field));
908            }
909            FBFieldType {
910                type_type: crate::Type::Struct_,
911                type_: crate::Struct_Builder::new(fbb).finish().as_union_value(),
912                children: Some(fbb.create_vector(&children[..])),
913            }
914        }
915        RunEndEncoded(run_ends, values) => {
916            let run_ends_field = build_field(fbb, dictionary_tracker, run_ends);
917            let values_field = build_field(fbb, dictionary_tracker, values);
918            let children = [run_ends_field, values_field];
919            FBFieldType {
920                type_type: crate::Type::RunEndEncoded,
921                type_: crate::RunEndEncodedBuilder::new(fbb)
922                    .finish()
923                    .as_union_value(),
924                children: Some(fbb.create_vector(&children[..])),
925            }
926        }
927        Map(map_field, keys_sorted) => {
928            let child = build_field(fbb, dictionary_tracker, map_field);
929            let mut field_type = crate::MapBuilder::new(fbb);
930            field_type.add_keysSorted(*keys_sorted);
931            FBFieldType {
932                type_type: crate::Type::Map,
933                type_: field_type.finish().as_union_value(),
934                children: Some(fbb.create_vector(&[child])),
935            }
936        }
937        Dictionary(_, value_type) => {
938            // In this library, the dictionary "type" is a logical construct. Here we
939            // pass through to the value type, as we've already captured the index
940            // type in the DictionaryEncoding metadata in the parent field
941            get_fb_field_type(value_type, dictionary_tracker, fbb)
942        }
943        Decimal32(precision, scale) => {
944            let mut builder = crate::DecimalBuilder::new(fbb);
945            builder.add_precision(*precision as i32);
946            builder.add_scale(*scale as i32);
947            builder.add_bitWidth(32);
948            FBFieldType {
949                type_type: crate::Type::Decimal,
950                type_: builder.finish().as_union_value(),
951                children: Some(fbb.create_vector(&empty_fields[..])),
952            }
953        }
954        Decimal64(precision, scale) => {
955            let mut builder = crate::DecimalBuilder::new(fbb);
956            builder.add_precision(*precision as i32);
957            builder.add_scale(*scale as i32);
958            builder.add_bitWidth(64);
959            FBFieldType {
960                type_type: crate::Type::Decimal,
961                type_: builder.finish().as_union_value(),
962                children: Some(fbb.create_vector(&empty_fields[..])),
963            }
964        }
965        Decimal128(precision, scale) => {
966            let mut builder = crate::DecimalBuilder::new(fbb);
967            builder.add_precision(*precision as i32);
968            builder.add_scale(*scale as i32);
969            builder.add_bitWidth(128);
970            FBFieldType {
971                type_type: crate::Type::Decimal,
972                type_: builder.finish().as_union_value(),
973                children: Some(fbb.create_vector(&empty_fields[..])),
974            }
975        }
976        Decimal256(precision, scale) => {
977            let mut builder = crate::DecimalBuilder::new(fbb);
978            builder.add_precision(*precision as i32);
979            builder.add_scale(*scale as i32);
980            builder.add_bitWidth(256);
981            FBFieldType {
982                type_type: crate::Type::Decimal,
983                type_: builder.finish().as_union_value(),
984                children: Some(fbb.create_vector(&empty_fields[..])),
985            }
986        }
987        Union(fields, mode) => {
988            let mut children = vec![];
989            for (_, field) in fields.iter() {
990                children.push(build_field(fbb, dictionary_tracker, field));
991            }
992
993            let union_mode = match mode {
994                UnionMode::Sparse => crate::UnionMode::Sparse,
995                UnionMode::Dense => crate::UnionMode::Dense,
996            };
997
998            let fbb_type_ids =
999                fbb.create_vector(&fields.iter().map(|(t, _)| t as i32).collect::<Vec<_>>());
1000            let mut builder = crate::UnionBuilder::new(fbb);
1001            builder.add_mode(union_mode);
1002            builder.add_typeIds(fbb_type_ids);
1003
1004            FBFieldType {
1005                type_type: crate::Type::Union,
1006                type_: builder.finish().as_union_value(),
1007                children: Some(fbb.create_vector(&children[..])),
1008            }
1009        }
1010    }
1011}
1012
1013/// Create an IPC dictionary encoding
1014pub(crate) fn get_fb_dictionary<'a>(
1015    index_type: &DataType,
1016    dict_id: i64,
1017    dict_is_ordered: bool,
1018    fbb: &mut FlatBufferBuilder<'a>,
1019) -> WIPOffset<crate::DictionaryEncoding<'a>> {
1020    // We assume that the dictionary index type (as an integer) has already been
1021    // validated elsewhere, and can safely assume we are dealing with integers
1022    let mut index_builder = crate::IntBuilder::new(fbb);
1023
1024    match *index_type {
1025        Int8 | Int16 | Int32 | Int64 => index_builder.add_is_signed(true),
1026        UInt8 | UInt16 | UInt32 | UInt64 => index_builder.add_is_signed(false),
1027        _ => {}
1028    }
1029
1030    match *index_type {
1031        Int8 | UInt8 => index_builder.add_bitWidth(8),
1032        Int16 | UInt16 => index_builder.add_bitWidth(16),
1033        Int32 | UInt32 => index_builder.add_bitWidth(32),
1034        Int64 | UInt64 => index_builder.add_bitWidth(64),
1035        _ => {}
1036    }
1037
1038    let index_builder = index_builder.finish();
1039
1040    let mut builder = crate::DictionaryEncodingBuilder::new(fbb);
1041    builder.add_id(dict_id);
1042    builder.add_indexType(index_builder);
1043    builder.add_isOrdered(dict_is_ordered);
1044
1045    builder.finish()
1046}
1047
1048/// An owned container for a validated [`Message`]
1049///
1050/// Safely decoding a flatbuffer requires validating the various embedded offsets,
1051/// see [`Verifier`]. This is a potentially expensive operation, and it is therefore desirable
1052/// to only do this once. [`crate::root_as_message`] performs this validation on construction,
1053/// however, it returns a [`Message`] borrowing the provided byte slice. This prevents
1054/// storing this [`Message`] in the same data structure that owns the buffer, as this
1055/// would require self-referential borrows.
1056///
1057/// [`MessageBuffer`] solves this problem by providing a safe API for a [`Message`]
1058/// without a lifetime bound.
1059#[derive(Clone)]
1060pub struct MessageBuffer(Buffer);
1061
1062impl Debug for MessageBuffer {
1063    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1064        self.as_ref().fmt(f)
1065    }
1066}
1067
1068impl MessageBuffer {
1069    /// Try to create a [`MessageBuffer`] from the provided [`Buffer`]
1070    pub fn try_new(buf: Buffer) -> Result<Self, ArrowError> {
1071        let opts = VerifierOptions::default();
1072        let mut v = Verifier::new(&opts, &buf);
1073        <ForwardsUOffset<Message>>::run_verifier(&mut v, 0).map_err(|err| {
1074            ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1075        })?;
1076        Ok(Self(buf))
1077    }
1078
1079    /// Return the [`Message`]
1080    #[inline]
1081    pub fn as_ref(&self) -> Message<'_> {
1082        // SAFETY: Run verifier on construction
1083        unsafe { crate::root_as_message_unchecked(&self.0) }
1084    }
1085}
1086
1087#[cfg(test)]
1088mod tests {
1089    use super::*;
1090
1091    /// Build a schema flatbuffer with no `fields` vector, using `build_field` to
1092    /// optionally add a single field. Returns bytes that the flatbuffer verifier
1093    /// accepts, so `try_fb_to_schema` is the layer that must reject them.
1094    fn schema_bytes_without_fields() -> Vec<u8> {
1095        let mut fbb = FlatBufferBuilder::new();
1096        // Schema table with `fields` left unset -> fb.fields() == None.
1097        let schema = crate::SchemaBuilder::new(&mut fbb).finish();
1098        fbb.finish(schema, None);
1099        fbb.finished_data().to_vec()
1100    }
1101
1102    /// Finish `fbb` as a schema holding the single `field`, assert the verifier
1103    /// accepts the bytes but `try_fb_to_schema` rejects them (rather than panics).
1104    fn assert_single_field_schema_is_err<'a>(
1105        mut fbb: FlatBufferBuilder<'a>,
1106        field: WIPOffset<crate::Field<'a>>,
1107    ) {
1108        let fields = fbb.create_vector(&[field]);
1109        let schema = {
1110            let mut b = crate::SchemaBuilder::new(&mut fbb);
1111            b.add_fields(fields);
1112            b.finish()
1113        };
1114        fbb.finish(schema, None);
1115        let fb = crate::root_as_schema(fbb.finished_data()).expect("verifier should accept");
1116        assert!(try_fb_to_schema(fb).is_err());
1117    }
1118
1119    #[test]
1120    fn try_fb_to_schema_missing_fields_is_err() {
1121        let bytes = schema_bytes_without_fields();
1122        // the verifier accepts it...
1123        let fb = crate::root_as_schema(&bytes).expect("verifier should accept");
1124        // ...but conversion must return Err, not panic.
1125        assert!(try_fb_to_schema(fb).is_err());
1126    }
1127
1128    #[test]
1129    fn try_fb_to_schema_unknown_float_precision_is_err() {
1130        let mut fbb = FlatBufferBuilder::new();
1131        let name = fbb.create_string("f");
1132        let float = {
1133            let mut b = crate::FloatingPointBuilder::new(&mut fbb);
1134            // a precision the verifier does not range-check.
1135            b.add_precision(crate::Precision(99));
1136            b.finish()
1137        };
1138        let field = {
1139            let mut b = crate::FieldBuilder::new(&mut fbb);
1140            b.add_name(name);
1141            b.add_type_type(crate::Type::FloatingPoint);
1142            b.add_type_(float.as_union_value());
1143            b.finish()
1144        };
1145        assert_single_field_schema_is_err(fbb, field);
1146    }
1147
1148    #[test]
1149    fn try_fb_to_schema_negative_fixed_size_binary_width_is_err() {
1150        let mut fbb = FlatBufferBuilder::new();
1151        let fsb =
1152            crate::FixedSizeBinary::create(&mut fbb, &crate::FixedSizeBinaryArgs { byteWidth: -1 });
1153        let field = crate::Field::create(
1154            &mut fbb,
1155            &crate::FieldArgs {
1156                type_type: crate::Type::FixedSizeBinary,
1157                type_: Some(fsb.as_union_value()),
1158                ..Default::default()
1159            },
1160        );
1161        assert_single_field_schema_is_err(fbb, field);
1162    }
1163
1164    #[test]
1165    fn try_fb_to_schema_type_none_is_err() {
1166        let mut fbb = FlatBufferBuilder::new();
1167        let name = fbb.create_string("f");
1168        let field = {
1169            let mut b = crate::FieldBuilder::new(&mut fbb);
1170            b.add_name(name);
1171            // type_type left as the union default (Type::NONE).
1172            b.finish()
1173        };
1174        assert_single_field_schema_is_err(fbb, field);
1175    }
1176
1177    #[test]
1178    fn try_fb_to_schema_union_without_type_ids_too_many_fields_is_err() {
1179        let mut fbb = FlatBufferBuilder::new();
1180        let null = crate::Null::create(&mut fbb, &crate::NullArgs::default());
1181        let child = crate::Field::create(
1182            &mut fbb,
1183            &crate::FieldArgs {
1184                type_type: crate::Type::Null,
1185                type_: Some(null.as_union_value()),
1186                ..Default::default()
1187            },
1188        );
1189        let children = fbb.create_vector(&[child; 129]);
1190        let union = crate::Union::create(&mut fbb, &crate::UnionArgs::default());
1191        let field = crate::Field::create(
1192            &mut fbb,
1193            &crate::FieldArgs {
1194                type_type: crate::Type::Union,
1195                type_: Some(union.as_union_value()),
1196                children: Some(children),
1197                ..Default::default()
1198            },
1199        );
1200        assert_single_field_schema_is_err(fbb, field);
1201    }
1202
1203    #[test]
1204    fn convert_schema_round_trip() {
1205        let md = HashMap::from([("Key".to_string(), "value".to_string())]);
1206        let field_md = HashMap::from([("k".to_string(), "v".to_string())]);
1207        let schema = Schema::new_with_metadata(
1208            vec![
1209                Field::new("uint8", DataType::UInt8, false).with_metadata(field_md),
1210                Field::new("uint16", DataType::UInt16, true),
1211                Field::new("uint32", DataType::UInt32, false),
1212                Field::new("uint64", DataType::UInt64, true),
1213                Field::new("int8", DataType::Int8, true),
1214                Field::new("int16", DataType::Int16, false),
1215                Field::new("int32", DataType::Int32, true),
1216                Field::new("int64", DataType::Int64, false),
1217                Field::new("float16", DataType::Float16, true),
1218                Field::new("float32", DataType::Float32, false),
1219                Field::new("float64", DataType::Float64, true),
1220                Field::new("null", DataType::Null, false),
1221                Field::new("bool", DataType::Boolean, false),
1222                Field::new("date32", DataType::Date32, false),
1223                Field::new("date64", DataType::Date64, true),
1224                Field::new("time32[s]", DataType::Time32(TimeUnit::Second), true),
1225                Field::new("time32[ms]", DataType::Time32(TimeUnit::Millisecond), false),
1226                Field::new("time64[us]", DataType::Time64(TimeUnit::Microsecond), false),
1227                Field::new("time64[ns]", DataType::Time64(TimeUnit::Nanosecond), true),
1228                Field::new(
1229                    "timestamp[s]",
1230                    DataType::Timestamp(TimeUnit::Second, None),
1231                    false,
1232                ),
1233                Field::new(
1234                    "timestamp[ms]",
1235                    DataType::Timestamp(TimeUnit::Millisecond, None),
1236                    true,
1237                ),
1238                Field::new(
1239                    "timestamp[us]",
1240                    DataType::Timestamp(TimeUnit::Microsecond, Some("Africa/Johannesburg".into())),
1241                    false,
1242                ),
1243                Field::new(
1244                    "timestamp[ns]",
1245                    DataType::Timestamp(TimeUnit::Nanosecond, None),
1246                    true,
1247                ),
1248                Field::new(
1249                    "interval[ym]",
1250                    DataType::Interval(IntervalUnit::YearMonth),
1251                    true,
1252                ),
1253                Field::new(
1254                    "interval[dt]",
1255                    DataType::Interval(IntervalUnit::DayTime),
1256                    true,
1257                ),
1258                Field::new(
1259                    "interval[mdn]",
1260                    DataType::Interval(IntervalUnit::MonthDayNano),
1261                    true,
1262                ),
1263                Field::new("utf8", DataType::Utf8, false),
1264                Field::new("utf8_view", DataType::Utf8View, false),
1265                Field::new("binary", DataType::Binary, false),
1266                Field::new("binary_view", DataType::BinaryView, false),
1267                Field::new_list(
1268                    "list[u8]",
1269                    Field::new_list_field(DataType::UInt8, false),
1270                    true,
1271                ),
1272                Field::new_fixed_size_list(
1273                    "fixed_size_list[u8]",
1274                    Field::new_list_field(DataType::UInt8, false),
1275                    2,
1276                    true,
1277                ),
1278                Field::new_list(
1279                    "list[struct<float32, int32, bool>]",
1280                    Field::new_struct(
1281                        "struct",
1282                        vec![
1283                            Field::new("float32", UInt8, false),
1284                            Field::new("int32", Int32, true),
1285                            Field::new("bool", Boolean, true),
1286                        ],
1287                        true,
1288                    ),
1289                    false,
1290                ),
1291                Field::new_struct(
1292                    "struct<dictionary<int32, utf8>>",
1293                    vec![Field::new(
1294                        "dictionary<int32, utf8>",
1295                        Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1296                        false,
1297                    )],
1298                    false,
1299                ),
1300                Field::new_struct(
1301                    "struct<int64, list[struct<date32, list[struct<>]>]>",
1302                    vec![
1303                        Field::new("int64", DataType::Int64, true),
1304                        Field::new_list(
1305                            "list[struct<date32, list[struct<>]>]",
1306                            Field::new_struct(
1307                                "struct",
1308                                vec![
1309                                    Field::new("date32", DataType::Date32, true),
1310                                    Field::new_list(
1311                                        "list[struct<>]",
1312                                        Field::new(
1313                                            "struct",
1314                                            DataType::Struct(Fields::empty()),
1315                                            false,
1316                                        ),
1317                                        false,
1318                                    ),
1319                                ],
1320                                false,
1321                            ),
1322                            false,
1323                        ),
1324                    ],
1325                    false,
1326                ),
1327                Field::new_union(
1328                    "union<int64, list[union<date32, list[union<>]>]>",
1329                    vec![0, 1],
1330                    vec![
1331                        Field::new("int64", DataType::Int64, true),
1332                        Field::new_list(
1333                            "list[union<date32, list[union<>]>]",
1334                            Field::new_union(
1335                                "union<date32, list[union<>]>",
1336                                vec![0, 1],
1337                                vec![
1338                                    Field::new("date32", DataType::Date32, true),
1339                                    Field::new_list(
1340                                        "list[union<>]",
1341                                        Field::new(
1342                                            "union",
1343                                            DataType::Union(
1344                                                UnionFields::empty(),
1345                                                UnionMode::Sparse,
1346                                            ),
1347                                            false,
1348                                        ),
1349                                        false,
1350                                    ),
1351                                ],
1352                                UnionMode::Dense,
1353                            ),
1354                            false,
1355                        ),
1356                    ],
1357                    UnionMode::Sparse,
1358                ),
1359                Field::new("struct<>", DataType::Struct(Fields::empty()), true),
1360                Field::new(
1361                    "union<>",
1362                    DataType::Union(UnionFields::empty(), UnionMode::Dense),
1363                    true,
1364                ),
1365                Field::new(
1366                    "union<>",
1367                    DataType::Union(UnionFields::empty(), UnionMode::Sparse),
1368                    true,
1369                ),
1370                Field::new(
1371                    "union<int32, utf8>",
1372                    DataType::Union(
1373                        UnionFields::try_new(
1374                            vec![2, 3], // non-default type ids
1375                            vec![
1376                                Field::new("int32", DataType::Int32, true),
1377                                Field::new("utf8", DataType::Utf8, true),
1378                            ],
1379                        )
1380                        .unwrap(),
1381                        UnionMode::Dense,
1382                    ),
1383                    true,
1384                ),
1385                #[expect(deprecated)]
1386                Field::new_dict(
1387                    "dictionary<int32, utf8>",
1388                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1389                    true,
1390                    123,
1391                    true,
1392                ),
1393                #[expect(deprecated)]
1394                Field::new_dict(
1395                    "dictionary<uint8, uint32>",
1396                    DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
1397                    true,
1398                    123,
1399                    true,
1400                ),
1401                Field::new("decimal<usize, usize>", DataType::Decimal128(10, 6), false),
1402            ],
1403            md,
1404        );
1405
1406        let mut dictionary_tracker = DictionaryTracker::new(true);
1407        let fb = IpcSchemaEncoder::new()
1408            .with_dictionary_tracker(&mut dictionary_tracker)
1409            .schema_to_fb(&schema);
1410
1411        // read back fields
1412        let ipc = crate::root_as_schema(fb.finished_data()).unwrap();
1413        let schema2 = try_fb_to_schema(ipc).unwrap();
1414        assert_eq!(schema, schema2);
1415    }
1416
1417    #[test]
1418    fn schema_from_bytes() {
1419        // Bytes of a schema generated via following python code, using pyarrow 10.0.1:
1420        //
1421        // import pyarrow as pa
1422        // schema = pa.schema([pa.field('field1', pa.uint32(), nullable=False)])
1423        // sink = pa.BufferOutputStream()
1424        // with pa.ipc.new_stream(sink, schema) as writer:
1425        //     pass
1426        // # stripping continuation & length prefix & suffix bytes to get only schema bytes
1427        // [x for x in sink.getvalue().to_pybytes()][8:-8]
1428        let bytes: Vec<u8> = vec![
1429            16, 0, 0, 0, 0, 0, 10, 0, 12, 0, 6, 0, 5, 0, 8, 0, 10, 0, 0, 0, 0, 1, 4, 0, 12, 0, 0,
1430            0, 8, 0, 8, 0, 0, 0, 4, 0, 8, 0, 0, 0, 4, 0, 0, 0, 1, 0, 0, 0, 20, 0, 0, 0, 16, 0, 20,
1431            0, 8, 0, 0, 0, 7, 0, 12, 0, 0, 0, 16, 0, 16, 0, 0, 0, 0, 0, 0, 2, 16, 0, 0, 0, 32, 0,
1432            0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 102, 105, 101, 108, 100, 49, 0, 0, 0, 0, 6,
1433            0, 8, 0, 4, 0, 6, 0, 0, 0, 32, 0, 0, 0,
1434        ];
1435        let ipc = crate::root_as_message(&bytes).unwrap();
1436        let schema = ipc.header_as_schema().unwrap();
1437
1438        // generate same message with Rust
1439        let data_gen = crate::writer::IpcDataGenerator::default();
1440        let mut dictionary_tracker = DictionaryTracker::new(true);
1441        let arrow_schema = Schema::new(vec![Field::new("field1", DataType::UInt32, false)]);
1442        let bytes = data_gen
1443            .schema_to_bytes_with_dictionary_tracker(
1444                &arrow_schema,
1445                &mut dictionary_tracker,
1446                &crate::writer::IpcWriteOptions::default(),
1447            )
1448            .ipc_message;
1449
1450        let ipc2 = crate::root_as_message(&bytes).unwrap();
1451        let schema2 = ipc2.header_as_schema().unwrap();
1452
1453        // can't compare schema directly as it compares the underlying bytes, which can differ
1454        assert!(schema.custom_metadata().is_none());
1455        assert!(schema2.custom_metadata().is_none());
1456        assert_eq!(schema.endianness(), schema2.endianness());
1457        assert!(schema.features().is_none());
1458        assert!(schema2.features().is_none());
1459        assert_eq!(
1460            try_fb_to_schema(schema).unwrap(),
1461            try_fb_to_schema(schema2).unwrap()
1462        );
1463
1464        assert_eq!(ipc.version(), ipc2.version());
1465        assert_eq!(ipc.header_type(), ipc2.header_type());
1466        assert_eq!(ipc.bodyLength(), ipc2.bodyLength());
1467        assert!(ipc.custom_metadata().is_none());
1468        assert!(ipc2.custom_metadata().is_none());
1469    }
1470}
arrow_ipc/convert.rs

arrow_ipc/
convert.rs
