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

arrow_schema/extension/canonical/
variable_shape_tensor.rs
