arrow/extension/variable_shape_tensor.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <cstdint>
0021 #include <memory>
0022 #include <optional>
0023 #include <string>
0024 #include <vector>
0025
0026 #include "arrow/extension_type.h"
0027
0028 namespace arrow::extension {
0029
0030 class ARROW_EXPORT VariableShapeTensorArray : public ExtensionArray {
0031  public:
0032   using ExtensionArray::ExtensionArray;
0033 };
0034
0035 /// \brief Concrete type class for variable-shape Tensor data.
0036 /// This is a canonical arrow extension type.
0037 /// See: https://arrow.apache.org/docs/format/CanonicalExtensions.html
0038 class ARROW_EXPORT VariableShapeTensorType : public ExtensionType {
0039  public:
0040   VariableShapeTensorType(const std::shared_ptr<DataType>& value_type, int32_t ndim,
0041                           std::vector<int64_t> permutation = {},
0042                           std::vector<std::string> dim_names = {},
0043                           std::vector<std::optional<int64_t>> uniform_shape = {})
0044       : ExtensionType(struct_({::arrow::field("data", list(value_type)),
0045                                ::arrow::field("shape", fixed_size_list(int32(), ndim))})),
0046         value_type_(value_type),
0047         ndim_(ndim),
0048         permutation_(std::move(permutation)),
0049         dim_names_(std::move(dim_names)),
0050         uniform_shape_(std::move(uniform_shape)) {}
0051
0052   std::string extension_name() const override { return "arrow.variable_shape_tensor"; }
0053   std::string ToString(bool show_metadata = false) const override;
0054
0055   /// Number of dimensions of tensor elements
0056   int32_t ndim() const { return ndim_; }
0057
0058   /// Value type of tensor elements
0059   const std::shared_ptr<DataType>& value_type() const { return value_type_; }
0060
0061   /// Permutation mapping from logical to physical memory layout of tensor elements
0062   const std::vector<int64_t>& permutation() const { return permutation_; }
0063
0064   /// Dimension names of tensor elements. Dimensions are ordered physically.
0065   const std::vector<std::string>& dim_names() const { return dim_names_; }
0066
0067   /// Shape of uniform dimensions.
0068   const std::vector<std::optional<int64_t>>& uniform_shape() const {
0069     return uniform_shape_;
0070   }
0071
0072   bool ExtensionEquals(const ExtensionType& other) const override;
0073
0074   std::string Serialize() const override;
0075
0076   Result<std::shared_ptr<DataType>> Deserialize(
0077       std::shared_ptr<DataType> storage_type,
0078       const std::string& serialized_data) const override;
0079
0080   /// Create a VariableShapeTensorArray from ArrayData
0081   std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
0082
0083   /// \brief Convert an ExtensionScalar to a Tensor
0084   ///
0085   /// This method will return a Tensor from ExtensionScalar with strides derived
0086   /// from shape and permutation stored. Shape and dim_names will be permuted
0087   /// according to permutation stored in the VariableShapeTensorType.
0088   static Result<std::shared_ptr<Tensor>> MakeTensor(
0089       const std::shared_ptr<ExtensionScalar>&);
0090
0091   /// \brief Create a VariableShapeTensorType instance
0092   static Result<std::shared_ptr<DataType>> Make(
0093       const std::shared_ptr<DataType>& value_type, int32_t ndim,
0094       std::vector<int64_t> permutation = {}, std::vector<std::string> dim_names = {},
0095       std::vector<std::optional<int64_t>> uniform_shape = {});
0096
0097  private:
0098   std::shared_ptr<DataType> value_type_;
0099   int32_t ndim_;
0100   std::vector<int64_t> permutation_;
0101   std::vector<std::string> dim_names_;
0102   std::vector<std::optional<int64_t>> uniform_shape_;
0103 };
0104
0105 /// \brief Return a VariableShapeTensorType instance.
0106 ARROW_EXPORT std::shared_ptr<DataType> variable_shape_tensor(
0107     const std::shared_ptr<DataType>& value_type, int32_t ndim,
0108     std::vector<int64_t> permutation = {}, std::vector<std::string> dim_names = {},
0109     std::vector<std::optional<int64_t>> uniform_shape = {});
0110
0111 }  // namespace arrow::extension