Support for fixed-size list in conversion of range tuple

apache · Dec 12, 2024 · 9996bc9 · 9996bc9
1 parent 70343d3
commit 9996bc9
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 0 deletions.
diff --git a/cpp/src/arrow/stl.h b/cpp/src/arrow/stl.h
@@ -187,6 +187,34 @@ struct ConversionTraits<std::vector<ValueCType>>
   }
 };
 
+template <class CType_, std::size_t N>
+struct ConversionTraits<std::array<CType_, N>>
+    : public CTypeTraits<std::array<CType_, N>> {
+  static arrow::Status AppendRow(FixedSizeListBuilder& builder,
+                                 const std::array<CType_, N>& values) {
+    auto vb = ::arrow::internal::checked_cast<typename CTypeTraits<CType_>::BuilderType*>(
+        builder.value_builder());
+    ARROW_RETURN_NOT_OK(builder.Append());
+    return vb->AppendValues(values.data(), N);
+  }
+
+  static std::array<CType_, N> GetEntry(const ::arrow::FixedSizeListArray& array,
+                                        size_t j) {
+    using ElementArrayType =
+        typename TypeTraits<typename stl::ConversionTraits<CType_>::ArrowType>::ArrayType;
+
+    const ElementArrayType& value_array =
+        ::arrow::internal::checked_cast<const ElementArrayType&>(*array.values());
+
+    std::array<CType_, N> arr;
+    for (int64_t i = 0; i < N; i++) {
+      arr[i] =
+          stl::ConversionTraits<CType_>::GetEntry(value_array, array.value_offset(j) + i);
+    }
+    return arr;
+  }
+};
+
 template <typename Optional>
 struct ConversionTraits<Optional, enable_if_optional_like<Optional>>
     : public CTypeTraits<typename std::decay<decltype(*std::declval<Optional>())>::type> {

diff --git a/cpp/src/arrow/stl_test.cc b/cpp/src/arrow/stl_test.cc
@@ -245,6 +245,26 @@ TEST(TestTableFromTupleVector, ListType) {
   ASSERT_TRUE(expected_table->Equals(*table));
 }
 
+TEST(TestTableFromTupleVector, FixedSizeListType) {
+  using tuple_type = std::tuple<std::array<int64_t, 4>>;
+
+  auto expected_schema = std::make_shared<Schema>(
+      FieldVector{field("column1", fixed_size_list(int64(), 4), false)});
+  std::shared_ptr<Array> expected_array =
+      ArrayFromJSON(fixed_size_list(int64(), 4), "[[1, 1, 2, 34], [2, -4, 1, 1]]");
+  std::shared_ptr<Table> expected_table = Table::Make(expected_schema, {expected_array});
+  std::cout << expected_table->ToString() << std::endl;
+
+  std::vector<tuple_type> rows{tuple_type(std::array<int64_t, 4>{1, 1, 2, 34}),
+                               tuple_type(std::array<int64_t, 4>{2, -4, 1, 1})};
+  std::vector<std::string> names{"column1"};
+
+  std::shared_ptr<Table> table;
+  ASSERT_OK(TableFromTupleRange(default_memory_pool(), rows, names, &table));
+
+  ASSERT_TRUE(expected_table->Equals(*table));
+}
+
 TEST(TestTableFromTupleVector, ReferenceTuple) {
   std::vector<std::string> names{"column1", "column2", "column3", "column4", "column5",
                                  "column6", "column7", "column8", "column9", "column10"};

diff --git a/cpp/src/arrow/type_traits.h b/cpp/src/arrow/type_traits.h
@@ -540,6 +540,15 @@ struct CTypeTraits<std::vector<CType>> : public TypeTraits<ListType> {
   }
 };
 
+template <typename CType, std::size_t N>
+struct CTypeTraits<std::array<CType, N>> : public TypeTraits<FixedSizeListType> {
+  using ArrowType = FixedSizeListType;
+
+  static auto type_singleton() {
+    return fixed_size_list(CTypeTraits<CType>::type_singleton(), N);
+  }
+};
+
 /// \addtogroup type-traits
 /// @{
 template <>