arrow_schema/extension/canonical/
fixed_shape_tensor.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//! FixedShapeTensor
19//!
20//! <https://arrow.apache.org/docs/format/CanonicalExtensions.html#fixed-shape-tensor>
21
22use serde_core::de::{self, MapAccess, Visitor};
23use serde_core::ser::SerializeStruct;
24use serde_core::{Deserialize, Deserializer, Serialize, Serializer};
25use std::fmt;
26
27use crate::{ArrowError, DataType, extension::ExtensionType};
28
29/// The extension type for fixed shape tensor.
30///
31/// Extension name: `arrow.fixed_shape_tensor`.
32///
33/// The storage type of the extension: `FixedSizeList` where:
34/// - `value_type` is the data type of individual tensor elements.
35/// - `list_size` is the product of all the elements in tensor shape.
36///
37/// Extension type parameters:
38/// - `value_type`: the Arrow data type of individual tensor elements.
39/// - `shape`: the physical shape of the contained tensors as an array.
40///
41/// Optional parameters describing the logical layout:
42/// - `dim_names`: explicit names to tensor dimensions as an array. The
43///   length of it should be equal to the shape length and equal to the
44///   number of dimensions.
45///   `dim_names` can be used if the dimensions have
46///   well-known names and they map to the physical layout (row-major).
47/// - `permutation`: indices of the desired ordering of the original
48///   dimensions, defined as an array.
49///   The indices contain a permutation of the values `[0, 1, .., N-1]`
50///   where `N` is the number of dimensions. The permutation indicates
51///   which dimension of the logical layout corresponds to which dimension
52///   of the physical tensor (the i-th dimension of the logical view
53///   corresponds to the dimension with number `permutations[i]` of the
54///   physical tensor).
55///   Permutation can be useful in case the logical order of the tensor is
56///   a permutation of the physical order (row-major).
57///   When logical and physical layout are equal, the permutation will
58///   always be `([0, 1, .., N-1])` and can therefore be left out.
59///
60/// Description of the serialization:
61/// The metadata must be a valid JSON object including shape of the
62/// contained tensors as an array with key `shape` plus optional
63/// dimension names with keys `dim_names` and ordering of the
64/// dimensions with key `permutation`.
65/// Example: `{ "shape": [2, 5]}`
66/// Example with `dim_names` metadata for NCHW ordered data:
67/// `{ "shape": [100, 200, 500], "dim_names": ["C", "H", "W"]}`
68/// Example of permuted 3-dimensional tensor:
69/// `{ "shape": [100, 200, 500], "permutation": [2, 0, 1]}`
70///
71/// This is the physical layout shape and the shape of the logical layout
72/// would in this case be `[500, 100, 200]`.
73///
74/// <https://arrow.apache.org/docs/format/CanonicalExtensions.html#fixed-shape-tensor>
75#[derive(Debug, Clone, PartialEq)]
76pub struct FixedShapeTensor {
77    /// The data type of individual tensor elements.
78    value_type: DataType,
79
80    /// The metadata of this extension type.
81    metadata: FixedShapeTensorMetadata,
82}
83
84impl FixedShapeTensor {
85    /// Returns a new fixed shape tensor extension type.
86    ///
87    /// # Error
88    ///
89    /// Return an error if the provided dimension names or permutations are
90    /// invalid.
91    pub fn try_new(
92        value_type: DataType,
93        shape: impl IntoIterator<Item = usize>,
94        dimension_names: Option<Vec<String>>,
95        permutations: Option<Vec<usize>>,
96    ) -> Result<Self, ArrowError> {
97        // TODO: are all data types are suitable as value type?
98        FixedShapeTensorMetadata::try_new(shape, dimension_names, permutations).map(|metadata| {
99            Self {
100                value_type,
101                metadata,
102            }
103        })
104    }
105
106    /// Returns the value type of the individual tensor elements.
107    pub fn value_type(&self) -> &DataType {
108        &self.value_type
109    }
110
111    /// Returns the product of all the elements in tensor shape.
112    pub fn list_size(&self) -> usize {
113        self.metadata.list_size()
114    }
115
116    /// Returns the physical shape of the contained tensors.
117    ///
118    /// [`Self::dimensions`] and [`Self::list_size`] do not identify this
119    /// shape. `[2, 6]` and `[3, 4]` are both 2-dimensional and both have
120    /// 12 elements.
121    ///
122    /// ```
123    /// # use arrow_schema::extension::FixedShapeTensor;
124    /// # use arrow_schema::DataType;
125    /// let tensor = FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, None).unwrap();
126    /// assert_eq!(tensor.shape(), &[2, 6]);
127    /// assert_eq!(tensor.list_size(), 12);
128    /// ```
129    pub fn shape(&self) -> &[usize] {
130        self.metadata.shape()
131    }
132
133    /// Returns the number of dimensions in this fixed shape tensor.
134    pub fn dimensions(&self) -> usize {
135        self.metadata.dimensions()
136    }
137
138    /// Returns the names of the dimensions in this fixed shape tensor, if
139    /// set.
140    pub fn dimension_names(&self) -> Option<&[String]> {
141        self.metadata.dimension_names()
142    }
143
144    /// Returns the indices of the desired ordering of the original
145    /// dimensions, if set.
146    pub fn permutations(&self) -> Option<&[usize]> {
147        self.metadata.permutations()
148    }
149
150    /// Returns the logical shape.
151    ///
152    /// Logical dimension `i` has the size of physical dimension
153    /// `permutations[i]`. Without a permutation, this is the physical shape.
154    ///
155    /// Physical shape `[100, 200, 500]` and permutation `[2, 0, 1]` have
156    /// logical shape `[500, 100, 200]`.
157    ///
158    /// ```
159    /// # use arrow_schema::extension::FixedShapeTensor;
160    /// # use arrow_schema::DataType;
161    /// let tensor = FixedShapeTensor::try_new(
162    ///     DataType::Float32,
163    ///     [100, 200, 500],
164    ///     None,
165    ///     Some(vec![2, 0, 1]),
166    /// )
167    /// .unwrap();
168    /// assert_eq!(tensor.logical_shape(), vec![500, 100, 200]);
169    /// ```
170    pub fn logical_shape(&self) -> Vec<usize> {
171        self.metadata.logical_shape()
172    }
173}
174
175/// Extension type metadata for [`FixedShapeTensor`].
176#[derive(Debug, Clone, PartialEq)]
177pub struct FixedShapeTensorMetadata {
178    /// The physical shape of the contained tensors.
179    shape: Vec<usize>,
180
181    /// Explicit names to tensor dimensions.
182    dim_names: Option<Vec<String>>,
183
184    /// Indices of the desired ordering of the original dimensions.
185    permutations: Option<Vec<usize>>,
186}
187
188impl Serialize for FixedShapeTensorMetadata {
189    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
190    where
191        S: Serializer,
192    {
193        let mut state = serializer.serialize_struct("FixedShapeTensorMetadata", 3)?;
194        state.serialize_field("shape", &self.shape)?;
195        state.serialize_field("dim_names", &self.dim_names)?;
196        state.serialize_field("permutations", &self.permutations)?;
197        state.end()
198    }
199}
200
201#[derive(Debug)]
202enum MetadataField {
203    Shape,
204    DimNames,
205    Permutations,
206}
207
208struct MetadataFieldVisitor;
209
210impl Visitor<'_> for MetadataFieldVisitor {
211    type Value = MetadataField;
212
213    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
214        formatter.write_str("`shape`, `dim_names`, or `permutations`")
215    }
216
217    fn visit_str<E>(self, value: &str) -> Result<MetadataField, E>
218    where
219        E: de::Error,
220    {
221        match value {
222            "shape" => Ok(MetadataField::Shape),
223            "dim_names" => Ok(MetadataField::DimNames),
224            "permutations" => Ok(MetadataField::Permutations),
225            _ => Err(de::Error::unknown_field(
226                value,
227                &["shape", "dim_names", "permutations"],
228            )),
229        }
230    }
231}
232
233impl<'de> Deserialize<'de> for MetadataField {
234    fn deserialize<D>(deserializer: D) -> Result<MetadataField, D::Error>
235    where
236        D: Deserializer<'de>,
237    {
238        deserializer.deserialize_identifier(MetadataFieldVisitor)
239    }
240}
241
242struct FixedShapeTensorMetadataVisitor;
243
244impl<'de> Visitor<'de> for FixedShapeTensorMetadataVisitor {
245    type Value = FixedShapeTensorMetadata;
246
247    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
248        formatter.write_str("struct FixedShapeTensorMetadata")
249    }
250
251    fn visit_seq<V>(self, mut seq: V) -> Result<FixedShapeTensorMetadata, V::Error>
252    where
253        V: de::SeqAccess<'de>,
254    {
255        let shape = seq
256            .next_element()?
257            .ok_or_else(|| de::Error::invalid_length(0, &self))?;
258        let dim_names = seq
259            .next_element()?
260            .ok_or_else(|| de::Error::invalid_length(1, &self))?;
261        let permutations = seq
262            .next_element()?
263            .ok_or_else(|| de::Error::invalid_length(2, &self))?;
264        Ok(FixedShapeTensorMetadata {
265            shape,
266            dim_names,
267            permutations,
268        })
269    }
270
271    fn visit_map<V>(self, mut map: V) -> Result<FixedShapeTensorMetadata, V::Error>
272    where
273        V: MapAccess<'de>,
274    {
275        let mut shape = None;
276        let mut dim_names = None;
277        let mut permutations = None;
278
279        while let Some(key) = map.next_key()? {
280            match key {
281                MetadataField::Shape => {
282                    if shape.is_some() {
283                        return Err(de::Error::duplicate_field("shape"));
284                    }
285                    shape = Some(map.next_value()?);
286                }
287                MetadataField::DimNames => {
288                    if dim_names.is_some() {
289                        return Err(de::Error::duplicate_field("dim_names"));
290                    }
291                    dim_names = Some(map.next_value()?);
292                }
293                MetadataField::Permutations => {
294                    if permutations.is_some() {
295                        return Err(de::Error::duplicate_field("permutations"));
296                    }
297                    permutations = Some(map.next_value()?);
298                }
299            }
300        }
301
302        let shape = shape.ok_or_else(|| de::Error::missing_field("shape"))?;
303
304        Ok(FixedShapeTensorMetadata {
305            shape,
306            dim_names,
307            permutations,
308        })
309    }
310}
311
312impl<'de> Deserialize<'de> for FixedShapeTensorMetadata {
313    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
314    where
315        D: Deserializer<'de>,
316    {
317        deserializer.deserialize_struct(
318            "FixedShapeTensorMetadata",
319            &["shape", "dim_names", "permutations"],
320            FixedShapeTensorMetadataVisitor,
321        )
322    }
323}
324
325impl FixedShapeTensorMetadata {
326    /// Returns metadata for a fixed shape tensor extension type.
327    ///
328    /// # Error
329    ///
330    /// Return an error if the provided dimension names or permutations are
331    /// invalid.
332    pub fn try_new(
333        shape: impl IntoIterator<Item = usize>,
334        dimension_names: Option<Vec<String>>,
335        permutations: Option<Vec<usize>>,
336    ) -> Result<Self, ArrowError> {
337        let shape = shape.into_iter().collect::<Vec<_>>();
338        let dimensions = shape.len();
339
340        let dim_names = dimension_names.map(|dimension_names| {
341            if dimension_names.len() != dimensions {
342                Err(ArrowError::InvalidArgumentError(format!(
343                    "FixedShapeTensor dimension names size mismatch, expected {dimensions}, found {}", dimension_names.len()
344                )))
345            } else {
346                Ok(dimension_names)
347            }
348        }).transpose()?;
349
350        let permutations = permutations
351            .map(|permutations| {
352                if permutations.len() != dimensions {
353                    Err(ArrowError::InvalidArgumentError(format!(
354                        "FixedShapeTensor permutations size mismatch, expected {dimensions}, found {}",
355                        permutations.len()
356                    )))
357                } else {
358                    let mut sorted_permutations = permutations.clone();
359                    sorted_permutations.sort_unstable();
360                    if (0..dimensions).zip(sorted_permutations).any(|(a, b)| a != b) {
361                        Err(ArrowError::InvalidArgumentError(format!(
362                            "FixedShapeTensor permutations invalid, expected a permutation of [0, 1, .., N-1], where N is the number of dimensions: {dimensions}"
363                        )))
364                    } else {
365                        Ok(permutations)
366                    }
367                }
368            })
369            .transpose()?;
370
371        Ok(Self {
372            shape,
373            dim_names,
374            permutations,
375        })
376    }
377
378    /// Returns the physical shape of the contained tensors.
379    pub fn shape(&self) -> &[usize] {
380        &self.shape
381    }
382
383    /// Returns the product of all the elements in tensor shape.
384    pub fn list_size(&self) -> usize {
385        self.shape.iter().product()
386    }
387
388    /// Returns the number of dimensions in this fixed shape tensor.
389    pub fn dimensions(&self) -> usize {
390        self.shape.len()
391    }
392
393    /// Returns the names of the dimensions in this fixed shape tensor, if
394    /// set.
395    pub fn dimension_names(&self) -> Option<&[String]> {
396        self.dim_names.as_ref().map(AsRef::as_ref)
397    }
398
399    /// Returns the indices of the desired ordering of the original
400    /// dimensions, if set.
401    pub fn permutations(&self) -> Option<&[usize]> {
402        self.permutations.as_ref().map(AsRef::as_ref)
403    }
404
405    /// Returns the logical shape.
406    ///
407    /// See [`FixedShapeTensor::logical_shape`].
408    pub fn logical_shape(&self) -> Vec<usize> {
409        match &self.permutations {
410            Some(permutation) => permutation.iter().map(|&index| self.shape[index]).collect(),
411            None => self.shape.clone(),
412        }
413    }
414}
415
416impl ExtensionType for FixedShapeTensor {
417    const NAME: &'static str = "arrow.fixed_shape_tensor";
418
419    type Metadata = FixedShapeTensorMetadata;
420
421    fn metadata(&self) -> &Self::Metadata {
422        &self.metadata
423    }
424
425    fn serialize_metadata(&self) -> Option<String> {
426        Some(serde_json::to_string(&self.metadata).expect("metadata serialization"))
427    }
428
429    fn deserialize_metadata(metadata: Option<&str>) -> Result<Self::Metadata, ArrowError> {
430        metadata.map_or_else(
431            || {
432                Err(ArrowError::InvalidArgumentError(
433                    "FixedShapeTensor extension types requires metadata".to_owned(),
434                ))
435            },
436            |value| {
437                serde_json::from_str(value).map_err(|e| {
438                    ArrowError::InvalidArgumentError(format!(
439                        "FixedShapeTensor metadata deserialization failed: {e}"
440                    ))
441                })
442            },
443        )
444    }
445
446    fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> {
447        let expected = DataType::new_fixed_size_list(
448            self.value_type.clone(),
449            i32::try_from(self.list_size()).expect("overflow"),
450            false,
451        );
452        data_type
453            .equals_datatype(&expected)
454            .then_some(())
455            .ok_or_else(|| {
456                ArrowError::InvalidArgumentError(format!(
457                    "FixedShapeTensor data type mismatch, expected {expected}, found {data_type}"
458                ))
459            })
460    }
461
462    fn try_new(data_type: &DataType, metadata: Self::Metadata) -> Result<Self, ArrowError> {
463        match data_type {
464            DataType::FixedSizeList(field, list_size) if !field.is_nullable() => {
465                // Make sure the metadata is valid.
466                let metadata = FixedShapeTensorMetadata::try_new(
467                    metadata.shape,
468                    metadata.dim_names,
469                    metadata.permutations,
470                )?;
471                // Make sure it is compatible with this data type.
472                let expected_size = i32::try_from(metadata.list_size()).expect("overflow");
473                if *list_size != expected_size {
474                    Err(ArrowError::InvalidArgumentError(format!(
475                        "FixedShapeTensor list size mismatch, expected {expected_size} (metadata), found {list_size} (data type)"
476                    )))
477                } else {
478                    Ok(Self {
479                        value_type: field.data_type().clone(),
480                        metadata,
481                    })
482                }
483            }
484            data_type => Err(ArrowError::InvalidArgumentError(format!(
485                "FixedShapeTensor data type mismatch, expected FixedSizeList with non-nullable field, found {data_type}"
486            ))),
487        }
488    }
489}
490
491#[cfg(test)]
492mod tests {
493    #[cfg(feature = "canonical_extension_types")]
494    use crate::extension::CanonicalExtensionType;
495    use crate::{
496        Field,
497        extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType},
498    };
499
500    use super::*;
501
502    #[test]
503    fn valid() -> Result<(), ArrowError> {
504        let fixed_shape_tensor = FixedShapeTensor::try_new(
505            DataType::Float32,
506            [100, 200, 500],
507            Some(vec!["C".to_owned(), "H".to_owned(), "W".to_owned()]),
508            Some(vec![2, 0, 1]),
509        )?;
510        let mut field = Field::new_fixed_size_list(
511            "",
512            Field::new("", DataType::Float32, false),
513            i32::try_from(fixed_shape_tensor.list_size()).expect("overflow"),
514            false,
515        );
516        field.try_with_extension_type(fixed_shape_tensor.clone())?;
517        assert_eq!(
518            field.try_extension_type::<FixedShapeTensor>()?,
519            fixed_shape_tensor
520        );
521        assert_eq!(fixed_shape_tensor.shape(), &[100, 200, 500]);
522        assert_eq!(fixed_shape_tensor.metadata().shape(), &[100, 200, 500]);
523        #[cfg(feature = "canonical_extension_types")]
524        assert_eq!(
525            field.try_canonical_extension_type()?,
526            CanonicalExtensionType::FixedShapeTensor(fixed_shape_tensor)
527        );
528        Ok(())
529    }
530
531    #[test]
532    fn logical_shape_follows_permutation() -> Result<(), ArrowError> {
533        let physical = FixedShapeTensor::try_new(DataType::Float32, [100, 200, 500], None, None)?;
534        assert_eq!(physical.logical_shape(), vec![100, 200, 500]);
535        assert_eq!(physical.metadata().logical_shape(), vec![100, 200, 500]);
536
537        let permuted = FixedShapeTensor::try_new(
538            DataType::Float32,
539            [100, 200, 500],
540            Some(vec!["C".to_owned(), "H".to_owned(), "W".to_owned()]),
541            Some(vec![2, 0, 1]),
542        )?;
543        assert_eq!(permuted.logical_shape(), vec![500, 100, 200]);
544        assert_eq!(permuted.metadata().logical_shape(), vec![500, 100, 200]);
545
546        let identity =
547            FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, Some(vec![0, 1]))?;
548        assert_eq!(identity.logical_shape(), vec![2, 6]);
549        Ok(())
550    }
551
552    #[test]
553    fn shape_distinguishes_equal_sizes() -> Result<(), ArrowError> {
554        let wide = FixedShapeTensor::try_new(DataType::Float32, [2, 6], None, None)?;
555        let tall = FixedShapeTensor::try_new(DataType::Float32, [3, 4], None, None)?;
556        assert_eq!(wide.dimensions(), tall.dimensions());
557        assert_eq!(wide.list_size(), tall.list_size());
558        assert_eq!(wide.shape(), &[2, 6]);
559        assert_eq!(tall.shape(), &[3, 4]);
560        assert_ne!(wide.shape(), tall.shape());
561        Ok(())
562    }
563
564    #[test]
565    #[should_panic(expected = "Extension type name missing")]
566    fn missing_name() {
567        let field =
568            Field::new_fixed_size_list("", Field::new("", DataType::Float32, false), 3, false)
569                .with_metadata([(
570                    EXTENSION_TYPE_METADATA_KEY,
571                    r#"{ "shape": [100, 200, 500], }"#,
572                )]);
573        field.extension_type::<FixedShapeTensor>();
574    }
575
576    #[test]
577    #[should_panic(expected = "FixedShapeTensor data type mismatch, expected FixedSizeList")]
578    fn invalid_type() {
579        let fixed_shape_tensor =
580            FixedShapeTensor::try_new(DataType::Int32, [100, 200, 500], None, None).unwrap();
581        let field = Field::new_fixed_size_list(
582            "",
583            Field::new("", DataType::Float32, false),
584            i32::try_from(fixed_shape_tensor.list_size()).expect("overflow"),
585            false,
586        );
587        field.with_extension_type(fixed_shape_tensor);
588    }
589
590    #[test]
591    #[should_panic(expected = "FixedShapeTensor extension types requires metadata")]
592    fn missing_metadata() {
593        let field =
594            Field::new_fixed_size_list("", Field::new("", DataType::Float32, false), 3, false)
595                .with_metadata([(EXTENSION_TYPE_NAME_KEY, FixedShapeTensor::NAME)]);
596        field.extension_type::<FixedShapeTensor>();
597    }
598
599    #[test]
600    #[should_panic(expected = "FixedShapeTensor metadata deserialization failed: \
601        unknown field `not-shape`, expected one of `shape`, `dim_names`, `permutations`")]
602    fn invalid_metadata() {
603        let fixed_shape_tensor =
604            FixedShapeTensor::try_new(DataType::Float32, [100, 200, 500], None, None).unwrap();
605        let field = Field::new_fixed_size_list(
606            "",
607            Field::new("", DataType::Float32, false),
608            i32::try_from(fixed_shape_tensor.list_size()).expect("overflow"),
609            false,
610        )
611        .with_metadata([
612            (EXTENSION_TYPE_NAME_KEY, FixedShapeTensor::NAME),
613            (EXTENSION_TYPE_METADATA_KEY, r#"{ "not-shape": [] }"#),
614        ]);
615        field.extension_type::<FixedShapeTensor>();
616    }
617
618    #[test]
619    #[should_panic(
620        expected = "FixedShapeTensor dimension names size mismatch, expected 3, found 2"
621    )]
622    fn invalid_metadata_dimension_names() {
623        FixedShapeTensor::try_new(
624            DataType::Float32,
625            [100, 200, 500],
626            Some(vec!["a".to_owned(), "b".to_owned()]),
627            None,
628        )
629        .unwrap();
630    }
631
632    #[test]
633    #[should_panic(expected = "FixedShapeTensor permutations size mismatch, expected 3, found 2")]
634    fn invalid_metadata_permutations_len() {
635        FixedShapeTensor::try_new(DataType::Float32, [100, 200, 500], None, Some(vec![1, 0]))
636            .unwrap();
637    }
638
639    #[test]
640    #[should_panic(
641        expected = "FixedShapeTensor permutations invalid, expected a permutation of [0, 1, .., N-1], where N is the number of dimensions: 3"
642    )]
643    fn invalid_metadata_permutations_values() {
644        FixedShapeTensor::try_new(
645            DataType::Float32,
646            [100, 200, 500],
647            None,
648            Some(vec![4, 3, 2]),
649        )
650        .unwrap();
651    }
652}
arrow_schema/extension/canonical/fixed_shape_tensor.rs

arrow_schema/extension/canonical/
fixed_shape_tensor.rs
