GH-15483: [C++] Add a Fixed Shape Tensor canonical ExtensionType

cpp/src/arrow/CMakeLists.txt

@@ -551,6 +551,7 @@ endif()
 if(ARROW_JSON)
  list(APPEND
  ARROW_SRCS
+ extension/fixed_shape_tensor.cc
  json/options.cc
  json/chunked_builder.cc
  json/chunker.cc
@@ -883,6 +884,7 @@ endif()
 
 if(ARROW_JSON)
  add_subdirectory(json)
+ add_subdirectory(extension)
 endif()
 
 if(ARROW_ORC)

cpp/src/arrow/extension/CMakeLists.txt

@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+add_arrow_test(test
+ SOURCES
+ fixed_shape_tensor_test.cc
+ PREFIX
+ "arrow-fixed-shape-tensor")
+
+arrow_install_all_headers("arrow/extension")

cpp/src/arrow/extension/fixed_shape_tensor.cc

@@ -0,0 +1,170 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <numeric>
+#include <sstream>
+
+#include "arrow/extension/fixed_shape_tensor.h"
+
+#include "arrow/array/array_nested.h"
+#include "arrow/array/array_primitive.h"
+#include "arrow/json/rapidjson_defs.h" // IWYU pragma: keep
+#include "arrow/util/int_util_overflow.h"
+#include "arrow/util/logging.h"
+#include "arrow/util/sort.h"
+
+#include <rapidjson/document.h>
+#include <rapidjson/writer.h>
+
+namespace rj = arrow::rapidjson;
+
+namespace arrow {
+namespace extension {
+
+bool FixedShapeTensorType::ExtensionEquals(const ExtensionType& other) const {
+ if (extension_name() != other.extension_name()) {
+ return false;
+ }
+ const auto& other_ext = static_cast<const FixedShapeTensorType&>(other);
+
+ auto is_permutation_trivial = [](const std::vector<int64_t>& permutation) {
+ for (size_t i = 1; i < permutation.size(); ++i) {
+ if (permutation[i - 1] + 1 != permutation[i]) {
+ return false;
+ }
+ }
+ return true;
+ };
+ const bool permutation_equivalent =
+ ((permutation_ == other_ext.permutation()) ||
+ (permutation_.empty() && is_permutation_trivial(other_ext.permutation())) ||
+ (is_permutation_trivial(permutation_) && other_ext.permutation().empty()));
+
+ return (storage_type()->Equals(other_ext.storage_type())) &&
+ (this->shape() == other_ext.shape()) && (dim_names_ == other_ext.dim_names()) &&
+ permutation_equivalent;
+}
+
+std::string FixedShapeTensorType::Serialize() const {
+ rj::Document document;
+ document.SetObject();
+ rj::Document::AllocatorType& allocator = document.GetAllocator();
+
+ rj::Value shape(rj::kArrayType);
+ for (auto v : shape_) {
+ shape.PushBack(v, allocator);
+ }
+ document.AddMember(rj::Value("shape", allocator), shape, allocator);
+
+ if (!permutation_.empty()) {
+ rj::Value permutation(rj::kArrayType);
+ for (auto v : permutation_) {
+ permutation.PushBack(v, allocator);
+ }
+ document.AddMember(rj::Value("permutation", allocator), permutation, allocator);
+ }
+
+ if (!dim_names_.empty()) {
+ rj::Value dim_names(rj::kArrayType);
+ for (std::string v : dim_names_) {
+ dim_names.PushBack(rj::Value{}.SetString(v.c_str(), allocator), allocator);
+ }
+ document.AddMember(rj::Value("dim_names", allocator), dim_names, allocator);
+ }
+
+ rj::StringBuffer buffer;
+ rj::Writer<rj::StringBuffer> writer(buffer);
+ document.Accept(writer);
+ return buffer.GetString();
+}
+
+Result<std::shared_ptr<DataType>> FixedShapeTensorType::Deserialize(
+ std::shared_ptr<DataType> storage_type, const std::string& serialized_data) const {
+ if (storage_type->id() != Type::FIXED_SIZE_LIST) {
+ return Status::Invalid("Expected FixedSizeList storage type, got ",
+ storage_type->ToString());
+ }
+ auto value_type =
+ internal::checked_pointer_cast<FixedSizeListType>(storage_type)->value_type();
+ rj::Document document;
+ if (document.Parse(serialized_data.data(), serialized_data.length()).HasParseError() ||
+ !document.HasMember("shape") || !document["shape"].IsArray()) {
+ return Status::Invalid("Invalid serialized JSON data: ", serialized_data);
+ }
+
+ std::vector<int64_t> shape;
+ for (auto& x : document["shape"].GetArray()) {
+ shape.emplace_back(x.GetInt64());
+ }
+ std::vector<int64_t> permutation;
+ if (document.HasMember("permutation")) {
+ for (auto& x : document["permutation"].GetArray()) {
+ permutation.emplace_back(x.GetInt64());
+ }
+ if (shape.size() != permutation.size()) {
+ return Status::Invalid("Invalid permutation");
+ }
+ }
+ std::vector<std::string> dim_names;
+ if (document.HasMember("dim_names")) {
+ for (auto& x : document["dim_names"].GetArray()) {
+ dim_names.emplace_back(x.GetString());
+ }
+ if (shape.size() != dim_names.size()) {
+ return Status::Invalid("Invalid dim_names");
+ }
+ }
+
+ return fixed_shape_tensor(value_type, shape, permutation, dim_names);
+}
+
+std::shared_ptr<Array> FixedShapeTensorType::MakeArray(
+ std::shared_ptr<ArrayData> data) const {
+ DCHECK_EQ(data->type->id(), Type::EXTENSION);
+ DCHECK_EQ("arrow.fixed_shape_tensor",
+ static_cast<const ExtensionType&>(*data->type).extension_name());
+ return std::make_shared<ExtensionArray>(data);
+}
+
+Result<std::shared_ptr<DataType>> FixedShapeTensorType::Make(
+ const std::shared_ptr<DataType>& value_type, const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& permutation, const std::vector<std::string>& dim_names) {
+ if (!permutation.empty() && shape.size() != permutation.size()) {
+ return Status::Invalid("permutation size must match shape size. Expected: ",
+ shape.size(), " Got: ", permutation.size());
+ }
+ if (!dim_names.empty() && shape.size() != dim_names.size()) {
+ return Status::Invalid("dim_names size must match shape size. Expected: ",
+ shape.size(), " Got: ", dim_names.size());
+ }
+ const auto size = std::accumulate(shape.begin(), shape.end(), static_cast<int64_t>(1),
+ std::multiplies<>());
+ return std::make_shared<FixedShapeTensorType>(value_type, static_cast<int32_t>(size),
+ shape, permutation, dim_names);
+}
+
+std::shared_ptr<DataType> fixed_shape_tensor(const std::shared_ptr<DataType>& value_type,
+ const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& permutation,
+ const std::vector<std::string>& dim_names) {
+ auto maybe_type = FixedShapeTensorType::Make(value_type, shape, permutation, dim_names);
+ ARROW_DCHECK_OK(maybe_type.status());
+ return maybe_type.MoveValueUnsafe();
+}
+
+} // namespace extension
+} // namespace arrow

cpp/src/arrow/extension/fixed_shape_tensor.h

@@ -0,0 +1,92 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/extension_type.h"
+
+namespace arrow {
+namespace extension {
+
+class ARROW_EXPORT FixedShapeTensorArray : public ExtensionArray {
+ public:
+ using ExtensionArray::ExtensionArray;
+};
+
+/// \brief Concrete type class for constant-size Tensor data.
+/// This is a canonical arrow extension type.
+/// See: https://arrow.apache.org/docs/format/CanonicalExtensions.html
+class ARROW_EXPORT FixedShapeTensorType : public ExtensionType {
+ public:
+ FixedShapeTensorType(const std::shared_ptr<DataType>& value_type, const int32_t& size,
+ const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& permutation = {},
+ const std::vector<std::string>& dim_names = {})
+ : ExtensionType(fixed_size_list(value_type, size)),
+ value_type_(value_type),
+ shape_(shape),
+ permutation_(permutation),
+ dim_names_(dim_names) {}
+
+ std::string extension_name() const override { return "arrow.fixed_shape_tensor"; }
+
+ /// Number of dimensions of tensor elements
+ size_t ndim() { return shape_.size(); }
+
+ /// Shape of tensor elements
+ const std::vector<int64_t> shape() const { return shape_; }
+
+ /// Value type of tensor elements
+ const std::shared_ptr<DataType> value_type() const { return value_type_; }
+
+ /// Permutation mapping from logical to physical memory layout of tensor elements
+ const std::vector<int64_t>& permutation() const { return permutation_; }
+
+ /// Dimension names of tensor elements. Dimensions are ordered physically.
+ const std::vector<std::string>& dim_names() const { return dim_names_; }
+
+ bool ExtensionEquals(const ExtensionType& other) const override;
+
+ std::string Serialize() const override;
+
+ Result<std::shared_ptr<DataType>> Deserialize(
+ std::shared_ptr<DataType> storage_type,
+ const std::string& serialized_data) const override;
+
+ /// Create a FixedShapeTensorArray from ArrayData
+ std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const override;
+
+ /// \brief Create a FixedShapeTensorType instance
+ static Result<std::shared_ptr<DataType>> Make(
+ const std::shared_ptr<DataType>& value_type, const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& permutation = {},
+ const std::vector<std::string>& dim_names = {});
+
+ private:
+ std::shared_ptr<DataType> storage_type_;
+ std::shared_ptr<DataType> value_type_;
+ std::vector<int64_t> shape_;
+ std::vector<int64_t> permutation_;
+ std::vector<std::string> dim_names_;
+};
+
+/// \brief Return a FixedShapeTensorType instance.
+ARROW_EXPORT std::shared_ptr<DataType> fixed_shape_tensor(
+ const std::shared_ptr<DataType>& storage_type, const std::vector<int64_t>& shape,
+ const std::vector<int64_t>& permutation = {},
+ const std::vector<std::string>& dim_names = {});
+
+} // namespace extension
+} // namespace arrow