Skip to main content

vortex_tensor/
matcher.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Matcher for tensor-like extension types.
5
6use vortex::dtype::extension::ExtDTypeRef;
7use vortex::dtype::extension::Matcher;
8
9use crate::fixed_shape::FixedShapeTensor;
10use crate::fixed_shape::FixedShapeTensorMetadata;
11use crate::vector::Vector;
12
13/// Matcher for any tensor-like extension type.
14///
15/// Currently the different kinds of tensors that are available are:
16///
17/// - `FixedShapeTensor`
18/// - `Vector`
19pub struct AnyTensor;
20
21/// The matched variant of a tensor-like extension type.
22#[derive(Debug, PartialEq, Eq)]
23pub enum TensorMatch<'a> {
24    /// A [`FixedShapeTensor`] extension type.
25    FixedShapeTensor(&'a FixedShapeTensorMetadata),
26    /// A [`Vector`] extension type.
27    Vector,
28}
29
30impl Matcher for AnyTensor {
31    type Match<'a> = TensorMatch<'a>;
32
33    fn try_match<'a>(item: &'a ExtDTypeRef) -> Option<Self::Match<'a>> {
34        if let Some(metadata) = item.metadata_opt::<FixedShapeTensor>() {
35            return Some(TensorMatch::FixedShapeTensor(metadata));
36        }
37        if item.metadata_opt::<Vector>().is_some() {
38            return Some(TensorMatch::Vector);
39        }
40        None
41    }
42}