This is an automated email from the ASF dual-hosted git repository.
AlenkaF pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new ca37093fc2 GH-40062: [C++][Python] Conversion of Table to Arrow Tensor
(#41870)
ca37093fc2 is described below
commit ca37093fc2433b98d45a44fd4146d87ea4c3755e
Author: Alenka Frim <[email protected]>
AuthorDate: Wed Jul 1 14:23:39 2026 +0200
GH-40062: [C++][Python] Conversion of Table to Arrow Tensor (#41870)
### Rationale for this change
There is currently no method to convert Arrow Table to Arrow Tensor
(conversion from columnar format to a contiguous block of memory). This work is
a continuation of `RecordBatch::ToTensor` work, see
https://github.com/apache/arrow/issues/40058.
### What changes are included in this PR?
This PR:
- implements `Table::ToTensor` conversion
- adds bindings to Python
- adds benchmarks in C++
- removes the code in `RecordBatch::ToTensor` and uses the Table
implementation (`RecordBatch::ToTensor` benchmarks checked)
### Are these changes tested?
Yes, in C++ and Python.
### Are there any user-facing changes?
No, it is a new feature.
* GitHub Issue: #40062
Lead-authored-by: AlenkaF <[email protected]>
Co-authored-by: Alenka Frim <[email protected]>
Co-authored-by: tadeja <[email protected]>
Co-authored-by: Copilot Autofix powered by AI
<[email protected]>
Co-authored-by: Rok Mihevc <[email protected]>
Signed-off-by: AlenkaF <[email protected]>
---
cpp/src/arrow/record_batch.cc | 1 -
cpp/src/arrow/record_batch.h | 6 +-
cpp/src/arrow/record_batch_test.cc | 56 +++-
cpp/src/arrow/table.cc | 9 +
cpp/src/arrow/table.h | 12 +
cpp/src/arrow/table_test.cc | 570 +++++++++++++++++++++++++++++++++++
cpp/src/arrow/tensor.cc | 179 +++++++----
cpp/src/arrow/tensor.h | 4 +
cpp/src/arrow/tensor_benchmark.cc | 41 +++
python/pyarrow/includes/libarrow.pxd | 3 +
python/pyarrow/table.pxi | 87 +++++-
python/pyarrow/tests/test_table.py | 113 +++++--
12 files changed, 967 insertions(+), 114 deletions(-)
diff --git a/cpp/src/arrow/record_batch.cc b/cpp/src/arrow/record_batch.cc
index 12e0f553b7..bc2612f92a 100644
--- a/cpp/src/arrow/record_batch.cc
+++ b/cpp/src/arrow/record_batch.cc
@@ -18,7 +18,6 @@
#include "arrow/record_batch.h"
#include <algorithm>
-#include <cmath>
#include <cstdlib>
#include <memory>
#include <mutex>
diff --git a/cpp/src/arrow/record_batch.h b/cpp/src/arrow/record_batch.h
index 17d7f9857a..d7c6de1ed3 100644
--- a/cpp/src/arrow/record_batch.h
+++ b/cpp/src/arrow/record_batch.h
@@ -90,11 +90,9 @@ class ARROW_EXPORT RecordBatch {
/// in the resulting struct array.
Result<std::shared_ptr<StructArray>> ToStructArray() const;
- /// \brief Convert record batch with one data type to Tensor
+ /// \brief Convert RecordBatch to Tensor
///
- /// Create a Tensor object with shape (number of rows, number of columns) and
- /// strides (type size in bytes, type size in bytes * number of rows).
- /// Generated Tensor will have column-major layout.
+ /// Create a Tensor object.
///
/// \param[in] null_to_nan if true, convert nulls to NaN
/// \param[in] row_major if true, create row-major Tensor else column-major
Tensor
diff --git a/cpp/src/arrow/record_batch_test.cc
b/cpp/src/arrow/record_batch_test.cc
index 904285fd1c..fea47244da 100644
--- a/cpp/src/arrow/record_batch_test.cc
+++ b/cpp/src/arrow/record_batch_test.cc
@@ -910,10 +910,11 @@ TEST_F(TestRecordBatch, ToTensorUnsupportedMissing) {
auto batch = RecordBatch::Make(schema, length, {a0, a1});
- ASSERT_RAISES_WITH_MESSAGE(TypeError,
- "Type error: Can only convert a RecordBatch with
no nulls. "
- "Set null_to_nan to true to convert nulls to NaN",
- batch->ToTensor());
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError,
+ "Type error: Can only convert a Table or RecordBatch with no "
+ "nulls. Set null_to_nan to true to convert nulls to NaN",
+ batch->ToTensor());
}
TEST_F(TestRecordBatch, ToTensorEmptyBatch) {
@@ -944,10 +945,11 @@ TEST_F(TestRecordBatch, ToTensorEmptyBatch) {
auto batch_no_columns =
RecordBatch::Make(::arrow::schema({}), 10,
std::vector<std::shared_ptr<Array>>{});
- ASSERT_RAISES_WITH_MESSAGE(TypeError,
- "Type error: Conversion to Tensor for
RecordBatches without "
- "columns/schema is not supported.",
- batch_no_columns->ToTensor());
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError,
+ "Type error: Conversion to Tensor for Tables or RecordBatches "
+ "without columns/schema is not supported.",
+ batch_no_columns->ToTensor());
}
template <typename DataType>
@@ -1116,6 +1118,44 @@ TEST_F(TestRecordBatch, ToTensorSupportedNullToNan) {
CheckTensorRowMajor<FloatType>(tensor2_row, 18, shape, strides_2);
}
+TEST_F(TestRecordBatch, ToTensorNullToNanFloat16) {
+ // Tensor::Equals does not yet support NaN-aware comparison for float16, so
+ // null slots are verified by inspecting the raw buffer directly.
+ const int length = 9;
+
+ auto f0 = field("f0", float16());
+ auto f1 = field("f1", float16());
+ auto schema = ::arrow::schema({f0, f1});
+
+ auto a0 = ArrayFromJSON(float16(), "[null, 2, 3, 4, 5, 6, 7, 8, 9]");
+ auto a1 = ArrayFromJSON(float16(), "[10, 20, 30, 40, null, 60, 70, 80, 90]");
+ auto batch = RecordBatch::Make(schema, length, {a0, a1});
+
+ // Column-major
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ batch->ToTensor(/*null_to_nan=*/true,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 2};
+ const int64_t f16_size = sizeof(uint16_t);
+ CheckTensor<HalfFloatType>(tensor, 18, shape, {f16_size, f16_size *
shape[0]});
+
+ const auto* buf = reinterpret_cast<const uint16_t*>(tensor->raw_data());
+ EXPECT_TRUE(util::Float16::FromBits(buf[0]).is_nan());
+ EXPECT_TRUE(util::Float16::FromBits(buf[13]).is_nan());
+
+ // Row-major
+ ASSERT_OK_AND_ASSIGN(auto tensor_row, batch->ToTensor(/*null_to_nan=*/true));
+ ASSERT_OK(tensor_row->Validate());
+
+ CheckTensorRowMajor<HalfFloatType>(tensor_row, 18, shape,
+ {f16_size * shape[1], f16_size});
+
+ const auto* buf_row = reinterpret_cast<const
uint16_t*>(tensor_row->raw_data());
+ EXPECT_TRUE(util::Float16::FromBits(buf_row[0]).is_nan());
+ EXPECT_TRUE(util::Float16::FromBits(buf_row[9]).is_nan());
+}
+
TEST_F(TestRecordBatch, ToTensorSupportedTypesMixed) {
const int length = 9;
diff --git a/cpp/src/arrow/table.cc b/cpp/src/arrow/table.cc
index 68a8a1951f..7a7168e931 100644
--- a/cpp/src/arrow/table.cc
+++ b/cpp/src/arrow/table.cc
@@ -36,6 +36,7 @@
#include "arrow/record_batch.h"
#include "arrow/result.h"
#include "arrow/status.h"
+#include "arrow/tensor.h"
#include "arrow/type.h"
#include "arrow/type_fwd.h"
#include "arrow/type_traits.h"
@@ -346,6 +347,14 @@ Result<std::shared_ptr<Table>>
Table::FromChunkedStructArray(
array->length());
}
+Result<std::shared_ptr<Tensor>> Table::ToTensor(bool null_to_nan, bool
row_major,
+ MemoryPool* pool) const {
+ std::shared_ptr<Tensor> tensor;
+ ARROW_RETURN_NOT_OK(
+ internal::TableToTensor(*this, null_to_nan, row_major, pool, &tensor));
+ return tensor;
+}
+
std::vector<std::string> Table::ColumnNames() const {
std::vector<std::string> names(num_columns());
for (int i = 0; i < num_columns(); ++i) {
diff --git a/cpp/src/arrow/table.h b/cpp/src/arrow/table.h
index dee6f6fdd3..051060a52c 100644
--- a/cpp/src/arrow/table.h
+++ b/cpp/src/arrow/table.h
@@ -102,6 +102,18 @@ class ARROW_EXPORT Table {
static Result<std::shared_ptr<Table>> FromChunkedStructArray(
const std::shared_ptr<ChunkedArray>& array);
+ /// \brief Convert Table to Tensor
+ ///
+ /// Create a Tensor object.
+ ///
+ /// \param[in] null_to_nan if true, convert nulls to NaN
+ /// \param[in] row_major if true, create row-major Tensor else column-major
Tensor
+ /// \param[in] pool the memory pool to allocate the tensor buffer
+ /// \return the resulting Tensor
+ Result<std::shared_ptr<Tensor>> ToTensor(
+ bool null_to_nan = false, bool row_major = true,
+ MemoryPool* pool = default_memory_pool()) const;
+
/// \brief Return the table schema
const std::shared_ptr<Schema>& schema() const { return schema_; }
diff --git a/cpp/src/arrow/table_test.cc b/cpp/src/arrow/table_test.cc
index 4182bf020a..6d60b6bda5 100644
--- a/cpp/src/arrow/table_test.cc
+++ b/cpp/src/arrow/table_test.cc
@@ -33,9 +33,11 @@
#include "arrow/compute/cast.h"
#include "arrow/record_batch.h"
#include "arrow/status.h"
+#include "arrow/tensor.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/testing/random.h"
#include "arrow/type.h"
+#include "arrow/util/float16.h"
#include "arrow/util/key_value_metadata.h"
namespace arrow {
@@ -528,6 +530,574 @@ TEST_F(TestTable, ConcatenateTables) {
ASSERT_RAISES(Invalid, ConcatenateTables({t1, t3}));
}
+TEST_F(TestTable, ToTensorUnsupportedType) {
+ auto f0 = field("f0", int32());
+ // Unsupported data type
+ auto f1 = field("f1", utf8());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(int32(), {"[1, 2, 3]", "[4, 5, 6, 7, 8, 9]"});
+ auto a1 = ChunkedArrayFromJSON(
+ utf8(), {R"(["a", "b", "c", "a", "b"])", R"(["c", "a", "b", "c"])"});
+
+ auto table = Table::Make(schema, {a0, a1});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError, "Type error: DataType is not supported: " +
a1->type()->ToString(),
+ table->ToTensor());
+
+ // Unsupported boolean data type
+ auto f2 = field("f2", boolean());
+
+ std::vector<std::shared_ptr<Field>> fields2 = {f0, f2};
+ auto schema2 = ::arrow::schema(fields2);
+ auto a2 = ChunkedArrayFromJSON(
+ boolean(), {"[true, false, true, true, false, true, false, true,
true]"});
+ auto table2 = Table::Make(schema2, {a0, a2});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError, "Type error: DataType is not supported: " +
a2->type()->ToString(),
+ table2->ToTensor());
+}
+
+TEST_F(TestTable, ToTensorUnsupportedMissing) {
+ auto f0 = field("f0", int32());
+ auto f1 = field("f1", int32());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(int32(), {"[1, 2, 3]", "[4, 5, 6, 7, 8, 9]"});
+ auto a1 = ChunkedArrayFromJSON(int32(), {"[10, 20]", "[30, 40, null, 60, 70,
80, 90]"});
+
+ auto table = Table::Make(schema, {a0, a1});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError,
+ "Type error: Can only convert a Table or RecordBatch with no "
+ "nulls. Set null_to_nan to true to convert nulls to NaN",
+ table->ToTensor());
+}
+
+TEST_F(TestTable, ToTensorEmptyTable) {
+ auto f0 = field("f0", int32());
+ auto f1 = field("f1", int32());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<Table> empty, Table::MakeEmpty(schema));
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_column,
+ empty->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor_column->Validate());
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_row, empty->ToTensor());
+ ASSERT_OK(tensor_row->Validate());
+
+ const std::vector<int64_t> strides = {4, 4};
+ const std::vector<int64_t> shape = {0, 2};
+
+ EXPECT_EQ(strides, tensor_column->strides());
+ EXPECT_EQ(shape, tensor_column->shape());
+ EXPECT_EQ(strides, tensor_row->strides());
+ EXPECT_EQ(shape, tensor_row->shape());
+
+ auto table_no_columns =
+ Table::Make(::arrow::schema({}), std::vector<std::shared_ptr<Array>>{});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ TypeError,
+ "Type error: Conversion to Tensor for Tables or RecordBatches "
+ "without columns/schema is not supported.",
+ table_no_columns->ToTensor());
+}
+
+template <typename DataType>
+void CheckTableToTensor(const std::shared_ptr<Tensor>& tensor, const int size,
+ const std::vector<int64_t> shape,
+ const std::vector<int64_t> f_strides) {
+ EXPECT_EQ(size, tensor->size());
+ EXPECT_EQ(TypeTraits<DataType>::type_singleton(), tensor->type());
+ EXPECT_EQ(shape, tensor->shape());
+ EXPECT_EQ(f_strides, tensor->strides());
+ EXPECT_FALSE(tensor->is_row_major());
+ EXPECT_TRUE(tensor->is_column_major());
+ EXPECT_TRUE(tensor->is_contiguous());
+}
+
+template <typename DataType>
+void CheckTableToTensorRowMajor(const std::shared_ptr<Tensor>& tensor, const
int size,
+ const std::vector<int64_t> shape,
+ const std::vector<int64_t> strides) {
+ EXPECT_EQ(size, tensor->size());
+ EXPECT_EQ(TypeTraits<DataType>::type_singleton(), tensor->type());
+ EXPECT_EQ(shape, tensor->shape());
+ EXPECT_EQ(strides, tensor->strides());
+ EXPECT_TRUE(tensor->is_row_major());
+ EXPECT_FALSE(tensor->is_column_major());
+ EXPECT_TRUE(tensor->is_contiguous());
+}
+
+TEST_F(TestTable, ToTensorSupportedNaN) {
+ auto f0 = field("f0", float32());
+ auto f1 = field("f1", float32());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(float32(), {"[NaN, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a1 =
+ ChunkedArrayFromJSON(float32(), {"[10, 20]", "[30, 40, NaN, 60, 70, 80,
90]"});
+
+ auto table = Table::Make(schema, {a0, a1});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 2};
+ const int64_t f32_size = sizeof(float);
+ std::vector<int64_t> f_strides = {f32_size, f32_size * shape[0]};
+ std::shared_ptr<Tensor> tensor_expected = TensorFromJSON(
+ float32(), "[NaN, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, NaN, 60,
70, 80, 90]",
+ shape, f_strides);
+
+ EXPECT_FALSE(tensor_expected->Equals(*tensor));
+ EXPECT_TRUE(tensor_expected->Equals(*tensor,
EqualOptions().nans_equal(true)));
+ CheckTableToTensor<FloatType>(tensor, 18, shape, f_strides);
+}
+
+TEST_F(TestTable, ToTensorSupportedNullToNan) {
+ // int32 + float32 = float64
+ auto f0 = field("f0", int32());
+ auto f1 = field("f1", float32());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(int32(), {"[null, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a1 =
+ ChunkedArrayFromJSON(float32(), {"[10, 20]", "[30, 40, null, 60, 70, 80,
90]"});
+
+ auto table = Table::Make(schema, {a0, a1});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/true,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 2};
+ const int64_t f64_size = sizeof(double);
+ std::vector<int64_t> f_strides = {f64_size, f64_size * shape[0]};
+ std::shared_ptr<Tensor> tensor_expected = TensorFromJSON(
+ float64(), "[NaN, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, NaN, 60,
70, 80, 90]",
+ shape, f_strides);
+
+ EXPECT_FALSE(tensor_expected->Equals(*tensor));
+ EXPECT_TRUE(tensor_expected->Equals(*tensor,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensor<DoubleType>(tensor, 18, shape, f_strides);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_row, table->ToTensor(/*null_to_nan=*/true));
+ ASSERT_OK(tensor_row->Validate());
+
+ std::vector<int64_t> strides = {f64_size * shape[1], f64_size};
+ std::shared_ptr<Tensor> tensor_expected_row = TensorFromJSON(
+ float64(), "[NaN, 10, 2, 20, 3, 30, 4, 40, 5, NaN, 6, 60, 7, 70, 8,
80, 9, 90]",
+ shape, strides);
+
+ EXPECT_FALSE(tensor_expected_row->Equals(*tensor_row));
+ EXPECT_TRUE(tensor_expected_row->Equals(*tensor_row,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensorRowMajor<DoubleType>(tensor_row, 18, shape, strides);
+
+ // int32 -> float64
+ auto f2 = field("f2", int32());
+
+ std::vector<std::shared_ptr<Field>> fields1 = {f0, f2};
+ auto schema1 = ::arrow::schema(fields1);
+
+ auto a2 = ChunkedArrayFromJSON(int32(), {"[10, 20]", "[30, 40, null, 60, 70,
80, 90]"});
+ auto table1 = Table::Make(schema1, {a0, a2});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor1,
+ table1->ToTensor(/*null_to_nan=*/true,
/*row_major=*/false));
+ ASSERT_OK(tensor1->Validate());
+
+ EXPECT_FALSE(tensor_expected->Equals(*tensor1));
+ EXPECT_TRUE(tensor_expected->Equals(*tensor1,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensor<DoubleType>(tensor1, 18, shape, f_strides);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor1_row,
table1->ToTensor(/*null_to_nan=*/true));
+ ASSERT_OK(tensor1_row->Validate());
+
+ EXPECT_FALSE(tensor_expected_row->Equals(*tensor1_row));
+ EXPECT_TRUE(tensor_expected_row->Equals(*tensor1_row,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensorRowMajor<DoubleType>(tensor1_row, 18, shape, strides);
+
+ // int8 -> float32
+ auto f3 = field("f3", int8());
+ auto f4 = field("f4", int8());
+
+ std::vector<std::shared_ptr<Field>> fields2 = {f3, f4};
+ auto schema2 = ::arrow::schema(fields2);
+
+ auto a3 = ChunkedArrayFromJSON(int8(), {"[null, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a4 = ChunkedArrayFromJSON(int8(), {"[10, 20]", "[30, 40, null, 60, 70,
80, 90]"});
+ auto table2 = Table::Make(schema2, {a3, a4});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor2,
+ table2->ToTensor(/*null_to_nan=*/true,
/*row_major=*/false));
+ ASSERT_OK(tensor2->Validate());
+
+ const int64_t f32_size = sizeof(float);
+ std::vector<int64_t> f_strides_2 = {f32_size, f32_size * shape[0]};
+ std::shared_ptr<Tensor> tensor_expected_2 = TensorFromJSON(
+ float32(), "[NaN, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, NaN, 60,
70, 80, 90]",
+ shape, f_strides_2);
+
+ EXPECT_FALSE(tensor_expected_2->Equals(*tensor2));
+ EXPECT_TRUE(tensor_expected_2->Equals(*tensor2,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensor<FloatType>(tensor2, 18, shape, f_strides_2);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor2_row,
table2->ToTensor(/*null_to_nan=*/true));
+ ASSERT_OK(tensor2_row->Validate());
+
+ std::vector<int64_t> strides_2 = {f32_size * shape[1], f32_size};
+ std::shared_ptr<Tensor> tensor2_expected_row = TensorFromJSON(
+ float32(), "[NaN, 10, 2, 20, 3, 30, 4, 40, 5, NaN, 6, 60, 7, 70, 8,
80, 9, 90]",
+ shape, strides_2);
+
+ EXPECT_FALSE(tensor2_expected_row->Equals(*tensor2_row));
+ EXPECT_TRUE(
+ tensor2_expected_row->Equals(*tensor2_row,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensorRowMajor<FloatType>(tensor2_row, 18, shape, strides_2);
+}
+
+TEST_F(TestTable, ToTensorNullToNanFloat16) {
+ // Tensor::Equals does not yet support NaN-aware comparison for float16, so
+ // null slots are verified by inspecting the raw buffer directly.
+ auto f0 = field("f0", float16());
+ auto f1 = field("f1", float16());
+ auto schema = ::arrow::schema({f0, f1});
+
+ auto a0 = ChunkedArrayFromJSON(float16(), {"[null, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a1 =
+ ChunkedArrayFromJSON(float16(), {"[10, 20]", "[30, 40, null, 60, 70, 80,
90]"});
+ auto table = Table::Make(schema, {a0, a1});
+
+ // Column-major
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/true,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 2};
+ const int64_t f16_size = sizeof(uint16_t);
+ CheckTableToTensor<HalfFloatType>(tensor, 18, shape, {f16_size, f16_size *
shape[0]});
+
+ const auto* buf = reinterpret_cast<const uint16_t*>(tensor->raw_data());
+ EXPECT_TRUE(util::Float16::FromBits(buf[0]).is_nan());
+ EXPECT_TRUE(util::Float16::FromBits(buf[13]).is_nan());
+
+ // Row-major
+ ASSERT_OK_AND_ASSIGN(auto tensor_row, table->ToTensor(/*null_to_nan=*/true));
+ ASSERT_OK(tensor_row->Validate());
+
+ CheckTableToTensorRowMajor<HalfFloatType>(tensor_row, 18, shape,
+ {f16_size * shape[1], f16_size});
+
+ const auto* buf_row = reinterpret_cast<const
uint16_t*>(tensor_row->raw_data());
+ EXPECT_TRUE(util::Float16::FromBits(buf_row[0]).is_nan());
+ EXPECT_TRUE(util::Float16::FromBits(buf_row[9]).is_nan());
+}
+
+TEST_F(TestTable, ToTensorSupportedTypesMixed) {
+ auto f0 = field("f0", uint16());
+ auto f1 = field("f1", int16());
+ auto f2 = field("f2", float32());
+
+ auto a0 = ChunkedArrayFromJSON(uint16(), {"[1, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a1 = ChunkedArrayFromJSON(int16(), {"[10, 20]", "[30, 40, 50, 60, 70,
80, 90]"});
+ auto a2 = ChunkedArrayFromJSON(float32(),
+ {"[100, 200, 300, NaN, 500, 600]", "[700,
800, 900]"});
+
+ // Single column
+ std::vector<std::shared_ptr<Field>> fields = {f0};
+ auto schema = ::arrow::schema(fields);
+ auto table = Table::Make(schema, {a0});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 1};
+ const int64_t uint16_size = sizeof(uint16_t);
+ std::vector<int64_t> f_strides = {uint16_size, uint16_size * shape[0]};
+ std::shared_ptr<Tensor> tensor_expected =
+ TensorFromJSON(uint16(), "[1, 2, 3, 4, 5, 6, 7, 8, 9]", shape,
f_strides);
+
+ EXPECT_TRUE(tensor_expected->Equals(*tensor));
+ CheckTableToTensor<UInt16Type>(tensor, 9, shape, f_strides);
+
+ // uint16 + int16 = int32
+ std::vector<std::shared_ptr<Field>> fields1 = {f0, f1};
+ auto schema1 = ::arrow::schema(fields1);
+ auto table1 = Table::Make(schema1, {a0, a1});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor1,
+ table1->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor1->Validate());
+
+ std::vector<int64_t> shape1 = {9, 2};
+ const int64_t int32_size = sizeof(int32_t);
+ std::vector<int64_t> f_strides_1 = {int32_size, int32_size * shape1[0]};
+ std::shared_ptr<Tensor> tensor_expected_1 = TensorFromJSON(
+ int32(), "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 60,
70, 80, 90]",
+ shape1, f_strides_1);
+
+ EXPECT_TRUE(tensor_expected_1->Equals(*tensor1));
+
+ CheckTableToTensor<Int32Type>(tensor1, 18, shape1, f_strides_1);
+
+ ASSERT_EQ(tensor1->type()->bit_width(),
tensor_expected_1->type()->bit_width());
+
+ ASSERT_EQ(1, tensor_expected_1->Value<Int32Type>({0, 0}));
+ ASSERT_EQ(2, tensor_expected_1->Value<Int32Type>({1, 0}));
+ ASSERT_EQ(10, tensor_expected_1->Value<Int32Type>({0, 1}));
+
+ // uint16 + int16 + float32 = float64
+ std::vector<std::shared_ptr<Field>> fields2 = {f0, f1, f2};
+ auto schema2 = ::arrow::schema(fields2);
+ auto table2 = Table::Make(schema2, {a0, a1, a2});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor2,
+ table2->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor2->Validate());
+
+ std::vector<int64_t> shape2 = {9, 3};
+ const int64_t f64_size = sizeof(double);
+ std::vector<int64_t> f_strides_2 = {f64_size, f64_size * shape2[0]};
+ std::shared_ptr<Tensor> tensor_expected_2 =
+ TensorFromJSON(float64(),
+ "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30,
40, 50,"
+ "60, 70, 80, 90, 100, 200, 300, NaN, 500, 600, 700, 800,
900]",
+ shape2, f_strides_2);
+
+ EXPECT_FALSE(tensor_expected_2->Equals(*tensor2));
+ EXPECT_TRUE(tensor_expected_2->Equals(*tensor2,
EqualOptions().nans_equal(true)));
+
+ CheckTableToTensor<DoubleType>(tensor2, 27, shape2, f_strides_2);
+}
+
+TEST_F(TestTable, ToTensorUnsupportedMixedFloat16) {
+ auto f0 = field("f0", float16());
+ auto f1 = field("f1", float64());
+
+ auto a0 = ChunkedArrayFromJSON(float16(), {"[1, 2, 3]", "[4, 5, 6, 7, 8,
9]"});
+ auto a1 = ChunkedArrayFromJSON(float64(), {"[10, 20]", "[30, 40, 50, 60, 70,
80, 90]"});
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1};
+ auto schema = ::arrow::schema(fields);
+ auto table = Table::Make(schema, {a0, a1});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ NotImplemented, "NotImplemented: Casting from or to halffloat is not
supported.",
+ table->ToTensor());
+
+ std::vector<std::shared_ptr<Field>> fields1 = {f1, f0};
+ auto schema1 = ::arrow::schema(fields1);
+ auto table1 = Table::Make(schema1, {a1, a0});
+
+ ASSERT_RAISES_WITH_MESSAGE(
+ NotImplemented, "NotImplemented: Casting from or to halffloat is not
supported.",
+ table1->ToTensor());
+}
+
+template <typename DataType>
+class TestTableToTensorColumnMajor : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(TestTableToTensorColumnMajor);
+
+TYPED_TEST_P(TestTableToTensorColumnMajor, SupportedTypes) {
+ using DataType = TypeParam;
+ using c_data_type = typename DataType::c_type;
+ const int unit_size = sizeof(c_data_type);
+
+ auto f0 = field("f0", TypeTraits<DataType>::type_singleton());
+ auto f1 = field("f1", TypeTraits<DataType>::type_singleton());
+ auto f2 = field("f2", TypeTraits<DataType>::type_singleton());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1, f2};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[1, 2, 3]", "[4, 5, 6, 7, 8, 9]"});
+ auto a1 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[10, 20]", "[30, 40, 50, 60, 70, 80, 90]"});
+ auto a2 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[100, 100, 100, 100, 100, 100]", "[100,
100, 100]"});
+
+ auto table = Table::Make(schema, {a0, a1, a2});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/false,
/*row_major=*/false));
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 3};
+ std::vector<int64_t> f_strides = {unit_size, unit_size * shape[0]};
+ std::shared_ptr<Tensor> tensor_expected = TensorFromJSON(
+ TypeTraits<DataType>::type_singleton(),
+ "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50,
60, 70, "
+ "80, 90, 100, 100, 100, 100, 100, 100, 100, 100, 100]",
+ shape, f_strides);
+
+ EXPECT_TRUE(tensor_expected->Equals(*tensor));
+ CheckTableToTensor<DataType>(tensor, 27, shape, f_strides);
+
+ // Test offsets
+ auto table_slice = table->Slice(1);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_sliced,
table_slice->ToTensor(/*null_to_nan=*/false,
+
/*row_major=*/false));
+ ASSERT_OK(tensor_sliced->Validate());
+
+ std::vector<int64_t> shape_sliced = {8, 3};
+ std::vector<int64_t> f_strides_sliced = {unit_size, unit_size *
shape_sliced[0]};
+ std::shared_ptr<Tensor> tensor_expected_sliced =
+ TensorFromJSON(TypeTraits<DataType>::type_singleton(),
+ "[2, 3, 4, 5, 6, 7, 8, 9, 20, 30, 40,
50, 60, "
+ "70, 80, 90, 100, 100, 100, 100, 100, 100, 100, 100]",
+ shape_sliced, f_strides_sliced);
+
+ EXPECT_TRUE(tensor_expected_sliced->Equals(*tensor_sliced));
+ CheckTableToTensor<DataType>(tensor_sliced, 24, shape_sliced,
f_strides_sliced);
+
+ auto table_slice_1 = table->Slice(1, 5);
+
+ ASSERT_OK_AND_ASSIGN(
+ auto tensor_sliced_1,
+ table_slice_1->ToTensor(/*null_to_nan=*/false, /*row_major=*/false));
+ ASSERT_OK(tensor_sliced_1->Validate());
+
+ std::vector<int64_t> shape_sliced_1 = {5, 3};
+ std::vector<int64_t> f_strides_sliced_1 = {unit_size, unit_size *
shape_sliced_1[0]};
+ std::shared_ptr<Tensor> tensor_expected_sliced_1 =
+ TensorFromJSON(TypeTraits<DataType>::type_singleton(),
+ "[2, 3, 4, 5, 6, 20, 30, 40, 50, 60, 100, 100, 100, 100,
100]",
+ shape_sliced_1, f_strides_sliced_1);
+
+ EXPECT_TRUE(tensor_expected_sliced_1->Equals(*tensor_sliced_1));
+ CheckTableToTensor<DataType>(tensor_sliced_1, 15, shape_sliced_1,
f_strides_sliced_1);
+}
+
+REGISTER_TYPED_TEST_SUITE_P(TestTableToTensorColumnMajor, SupportedTypes);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt8, TestTableToTensorColumnMajor, UInt8Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt16, TestTableToTensorColumnMajor,
UInt16Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt32, TestTableToTensorColumnMajor,
UInt32Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt64, TestTableToTensorColumnMajor,
UInt64Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int8, TestTableToTensorColumnMajor, Int8Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int16, TestTableToTensorColumnMajor, Int16Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int32, TestTableToTensorColumnMajor, Int32Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int64, TestTableToTensorColumnMajor, Int64Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float16, TestTableToTensorColumnMajor,
HalfFloatType);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float32, TestTableToTensorColumnMajor,
FloatType);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float64, TestTableToTensorColumnMajor,
DoubleType);
+
+template <typename DataType>
+class TestTableToTensorRowMajor : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(TestTableToTensorRowMajor);
+
+TYPED_TEST_P(TestTableToTensorRowMajor, SupportedTypes) {
+ using DataType = TypeParam;
+ using c_data_type = typename DataType::c_type;
+ const int unit_size = sizeof(c_data_type);
+
+ auto f0 = field("f0", TypeTraits<DataType>::type_singleton());
+ auto f1 = field("f1", TypeTraits<DataType>::type_singleton());
+ auto f2 = field("f2", TypeTraits<DataType>::type_singleton());
+
+ std::vector<std::shared_ptr<Field>> fields = {f0, f1, f2};
+ auto schema = ::arrow::schema(fields);
+
+ auto a0 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[1, 2, 3]", "[4, 5, 6, 7, 8, 9]"});
+ auto a1 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[10, 20]", "[30, 40, 50, 60, 70, 80, 90]"});
+ auto a2 = ChunkedArrayFromJSON(TypeTraits<DataType>::type_singleton(),
+ {"[100, 100, 100, 100, 100, 100]", "[100,
100, 100]"});
+
+ auto table = Table::Make(schema, {a0, a1, a2});
+
+ ASSERT_OK_AND_ASSIGN(auto tensor, table->ToTensor());
+ ASSERT_OK(tensor->Validate());
+
+ std::vector<int64_t> shape = {9, 3};
+ std::vector<int64_t> strides = {unit_size * shape[1], unit_size};
+ std::shared_ptr<Tensor> tensor_expected =
+ TensorFromJSON(TypeTraits<DataType>::type_singleton(),
+ "[1, 10, 100, 2, 20, 100, 3, 30, 100, 4, 40, 100, 5,
50, 100, 6, "
+ "60, 100, 7, 70, 100, 8, 80, 100, 9, 90, 100]",
+ shape, strides);
+
+ EXPECT_TRUE(tensor_expected->Equals(*tensor));
+ CheckTableToTensorRowMajor<DataType>(tensor, 27, shape, strides);
+
+ // Test offsets
+ auto table_slice = table->Slice(1);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_sliced, table_slice->ToTensor());
+ ASSERT_OK(tensor_sliced->Validate());
+
+ std::vector<int64_t> shape_sliced = {8, 3};
+ std::vector<int64_t> strides_sliced = {unit_size * shape[1], unit_size};
+ std::shared_ptr<Tensor> tensor_expected_sliced =
+ TensorFromJSON(TypeTraits<DataType>::type_singleton(),
+ "[2, 20, 100, 3, 30, 100, 4, 40, 100, 5, 50, 100, 6, "
+ "60, 100, 7, 70, 100, 8, 80, 100, 9, 90, 100]",
+ shape_sliced, strides_sliced);
+
+ EXPECT_TRUE(tensor_expected_sliced->Equals(*tensor_sliced));
+ CheckTableToTensorRowMajor<DataType>(tensor_sliced, 24, shape_sliced,
strides_sliced);
+
+ auto table_slice_1 = table->Slice(1, 5);
+
+ ASSERT_OK_AND_ASSIGN(auto tensor_sliced_1, table_slice_1->ToTensor());
+ ASSERT_OK(tensor_sliced_1->Validate());
+
+ std::vector<int64_t> shape_sliced_1 = {5, 3};
+ std::vector<int64_t> strides_sliced_1 = {unit_size * shape_sliced_1[1],
unit_size};
+ std::shared_ptr<Tensor> tensor_expected_sliced_1 =
+ TensorFromJSON(TypeTraits<DataType>::type_singleton(),
+ "[2, 20, 100, 3, 30, 100, 4, 40, 100, 5, 50, 100, 6, 60,
100]",
+ shape_sliced_1, strides_sliced_1);
+
+ EXPECT_TRUE(tensor_expected_sliced_1->Equals(*tensor_sliced_1));
+ CheckTableToTensorRowMajor<DataType>(tensor_sliced_1, 15, shape_sliced_1,
+ strides_sliced_1);
+}
+
+REGISTER_TYPED_TEST_SUITE_P(TestTableToTensorRowMajor, SupportedTypes);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt8, TestTableToTensorRowMajor, UInt8Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt16, TestTableToTensorRowMajor, UInt16Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt32, TestTableToTensorRowMajor, UInt32Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(UInt64, TestTableToTensorRowMajor, UInt64Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int8, TestTableToTensorRowMajor, Int8Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int16, TestTableToTensorRowMajor, Int16Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int32, TestTableToTensorRowMajor, Int32Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Int64, TestTableToTensorRowMajor, Int64Type);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float16, TestTableToTensorRowMajor,
HalfFloatType);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float32, TestTableToTensorRowMajor, FloatType);
+INSTANTIATE_TYPED_TEST_SUITE_P(Float64, TestTableToTensorRowMajor, DoubleType);
+
std::shared_ptr<Table> MakeTableWithOneNullFilledColumn(
const std::string& column_name, const std::shared_ptr<DataType>& data_type,
const int length) {
diff --git a/cpp/src/arrow/tensor.cc b/cpp/src/arrow/tensor.cc
index 49b7a0b28e..b5988d7810 100644
--- a/cpp/src/arrow/tensor.cc
+++ b/cpp/src/arrow/tensor.cc
@@ -28,11 +28,12 @@
#include <type_traits>
#include <vector>
-#include "arrow/record_batch.h"
#include "arrow/status.h"
+#include "arrow/table.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/checked_cast.h"
+#include "arrow/util/float16.h"
#include "arrow/util/int_util_overflow.h"
#include "arrow/util/logging_internal.h"
#include "arrow/util/unreachable.h"
@@ -225,7 +226,7 @@ Status ValidateTensorParameters(const
std::shared_ptr<DataType>& type,
}
template <typename Out>
-struct ConvertColumnsToTensorVisitor {
+struct ConvertArrayToTensorVisitor {
Out*& out_values;
const ArrayData& in_data;
@@ -246,8 +247,14 @@ struct ConvertColumnsToTensorVisitor {
}
} else {
for (int64_t i = 0; i < in_data.length; ++i) {
- *out_values++ =
- in_data.IsNull(i) ? static_cast<Out>(NAN) :
static_cast<Out>(in_values[i]);
+ if constexpr (T::type_id == Type::HALF_FLOAT && std::is_same_v<Out,
uint16_t>) {
+ *out_values++ = in_data.IsNull(i)
+ ?
std::numeric_limits<util::Float16>::quiet_NaN().bits()
+ : static_cast<Out>(in_values[i]);
+ } else {
+ *out_values++ = in_data.IsNull(i) ? static_cast<Out>(NAN)
+ : static_cast<Out>(in_values[i]);
+ }
}
}
return Status::OK();
@@ -257,11 +264,12 @@ struct ConvertColumnsToTensorVisitor {
};
template <typename Out>
-struct ConvertColumnsToTensorRowMajorVisitor {
+struct ConvertArrayToTensorRowMajorVisitor {
Out*& out_values;
const ArrayData& in_data;
- int num_cols;
- int col_idx;
+ int64_t num_cols;
+ int64_t col_idx;
+ int64_t chunk_idx;
template <typename T>
Status Visit(const T&) {
@@ -269,14 +277,23 @@ struct ConvertColumnsToTensorRowMajorVisitor {
using In = typename T::c_type;
auto in_values = ArraySpan(in_data).GetSpan<In>(1, in_data.length);
+ const int64_t base = chunk_idx * num_cols + col_idx;
+
if (in_data.null_count == 0) {
for (int64_t i = 0; i < in_data.length; ++i) {
- out_values[i * num_cols + col_idx] = static_cast<Out>(in_values[i]);
+ out_values[base + i * num_cols] = static_cast<Out>(in_values[i]);
}
} else {
for (int64_t i = 0; i < in_data.length; ++i) {
- out_values[i * num_cols + col_idx] =
- in_data.IsNull(i) ? static_cast<Out>(NAN) :
static_cast<Out>(in_values[i]);
+ if constexpr (T::type_id == Type::HALF_FLOAT && std::is_same_v<Out,
uint16_t>) {
+ out_values[base + i * num_cols] =
+ in_data.IsNull(i) ?
std::numeric_limits<util::Float16>::quiet_NaN().bits()
+ : static_cast<Out>(in_values[i]);
+ } else {
+ out_values[base + i * num_cols] = in_data.IsNull(i)
+ ? static_cast<Out>(NAN)
+ :
static_cast<Out>(in_values[i]);
+ }
}
}
return Status::OK();
@@ -285,50 +302,75 @@ struct ConvertColumnsToTensorRowMajorVisitor {
}
};
-template <typename DataType>
-inline void ConvertColumnsToTensor(const RecordBatch& batch, uint8_t* out,
+template <typename DataType, typename Container>
+inline void ConvertColumnsToTensor(const Container& container, uint8_t* out,
bool row_major) {
using CType = typename arrow::TypeTraits<DataType>::CType;
auto* out_values = reinterpret_cast<CType*>(out);
- int i = 0;
- for (const auto& column : batch.columns()) {
- if (row_major) {
- ConvertColumnsToTensorRowMajorVisitor<CType> visitor{out_values,
*column->data(),
-
batch.num_columns(), i++};
- DCHECK_OK(VisitTypeInline(*column->type(), &visitor));
- } else {
- ConvertColumnsToTensorVisitor<CType> visitor{out_values,
*column->data()};
- DCHECK_OK(VisitTypeInline(*column->type(), &visitor));
+ const int num_columns = container.num_columns();
+
+ for (int col_idx = 0; col_idx < num_columns; ++col_idx) {
+ if constexpr (std::is_same_v<Container, Table>) {
+ int64_t chunk_idx = 0;
+
+ for (const auto& chunk : container.columns()[col_idx]->chunks()) {
+ if (row_major) {
+ ConvertArrayToTensorRowMajorVisitor<CType> visitor{
+ out_values, *chunk->data(), num_columns, col_idx, chunk_idx};
+ DCHECK_OK(VisitTypeInline(*chunk->type(), &visitor));
+ chunk_idx += chunk->length();
+ } else {
+ ConvertArrayToTensorVisitor<CType> visitor{out_values,
*chunk->data()};
+ DCHECK_OK(VisitTypeInline(*chunk->type(), &visitor));
+ }
+ }
+ } else if constexpr (std::is_same_v<Container, RecordBatch>) {
+ const auto& array_data = container.column_data()[col_idx];
+
+ if (row_major) {
+ ConvertArrayToTensorRowMajorVisitor<CType> visitor{out_values,
*array_data,
+ num_columns,
col_idx, 0};
+ DCHECK_OK(VisitTypeInline(*array_data->type, &visitor));
+ } else {
+ ConvertArrayToTensorVisitor<CType> visitor{out_values, *array_data};
+ DCHECK_OK(VisitTypeInline(*array_data->type, &visitor));
+ }
}
}
}
-Status RecordBatchToTensor(const RecordBatch& batch, bool null_to_nan, bool
row_major,
- MemoryPool* pool, std::shared_ptr<Tensor>* tensor) {
- if (batch.num_columns() == 0) {
+template <typename Container>
+Status ToTensorImpl(const Container& container, bool null_to_nan, bool
row_major,
+ MemoryPool* pool, std::shared_ptr<Tensor>* tensor) {
+ if (container.num_columns() == 0) {
return Status::TypeError(
- "Conversion to Tensor for RecordBatches without columns/schema is not "
+ "Conversion to Tensor for Tables or RecordBatches without
columns/schema is not "
"supported.");
}
// Check for no validity bitmap of each field
// if null_to_nan conversion is set to false
- for (int i = 0; i < batch.num_columns(); ++i) {
- if (batch.column(i)->null_count() > 0 && !null_to_nan) {
+ for (int i = 0; i < container.num_columns(); ++i) {
+ int64_t null_count = 0;
+ if constexpr (std::is_same_v<Container, Table>) {
+ null_count = container.column(i)->null_count();
+ } else if constexpr (std::is_same_v<Container, RecordBatch>) {
+ null_count = container.column_data(i)->GetNullCount();
+ }
+ if (null_count > 0 && !null_to_nan) {
return Status::TypeError(
- "Can only convert a RecordBatch with no nulls. Set null_to_nan to
true to "
- "convert nulls to NaN");
+ "Can only convert a Table or RecordBatch with no nulls. Set
null_to_nan to "
+ "true to convert nulls to NaN");
}
}
// Check for supported data types and merge fields
// to get the resulting uniform data type
- if (!is_integer(batch.column(0)->type()->id()) &&
- !is_floating(batch.column(0)->type()->id())) {
- return Status::TypeError("DataType is not supported: ",
- batch.column(0)->type()->ToString());
+ const auto& col_0_type = container.schema()->field(0)->type();
+ if (!is_integer(col_0_type->id()) && !is_floating(col_0_type->id())) {
+ return Status::TypeError("DataType is not supported: ",
col_0_type->ToString());
}
- std::shared_ptr<Field> result_field = batch.schema()->field(0);
+ std::shared_ptr<Field> result_field = container.schema()->field(0);
std::shared_ptr<DataType> result_type = result_field->type();
Field::MergeOptions options;
@@ -336,24 +378,27 @@ Status RecordBatchToTensor(const RecordBatch& batch, bool
null_to_nan, bool row_
options.promote_integer_sign = true;
options.promote_numeric_width = true;
- if (batch.num_columns() > 1) {
- for (int i = 1; i < batch.num_columns(); ++i) {
- if (!is_numeric(batch.column(i)->type()->id())) {
- return Status::TypeError("DataType is not supported: ",
- batch.column(i)->type()->ToString());
+ if (container.num_columns() > 1) {
+ for (int i = 1; i < container.num_columns(); ++i) {
+ const auto& col_type = container.schema()->field(i)->type();
+
+ if (!is_numeric(col_type->id())) {
+ return Status::TypeError("DataType is not supported: ",
col_type->ToString());
}
// Casting of float16 is not supported, throw an error in this case
- if ((batch.column(i)->type()->id() == Type::HALF_FLOAT ||
+ if ((col_type->id() == Type::HALF_FLOAT ||
result_field->type()->id() == Type::HALF_FLOAT) &&
- batch.column(i)->type()->id() != result_field->type()->id()) {
+ col_type->id() != result_field->type()->id()) {
return Status::NotImplemented("Casting from or to halffloat is not
supported.");
}
- ARROW_ASSIGN_OR_RAISE(
- result_field,
- result_field->MergeWith(
- batch.schema()->field(i)->WithName(result_field->name()),
options));
+ if (!col_type->Equals(result_field->type())) {
+ ARROW_ASSIGN_OR_RAISE(
+ result_field,
+ result_field->MergeWith(
+ container.schema()->field(i)->WithName(result_field->name()),
options));
+ }
}
result_type = result_field->type();
}
@@ -368,42 +413,46 @@ Status RecordBatchToTensor(const RecordBatch& batch, bool
null_to_nan, bool row_
}
// Allocate memory
- ARROW_ASSIGN_OR_RAISE(
- std::shared_ptr<Buffer> result,
- AllocateBuffer(result_type->bit_width() * batch.num_columns() *
batch.num_rows(),
- pool));
+ int64_t buffer_size = result_type->byte_width();
+ if (internal::MultiplyWithOverflow(
+ buffer_size, static_cast<int64_t>(container.num_columns()),
&buffer_size) ||
+ internal::MultiplyWithOverflow(buffer_size, container.num_rows(),
&buffer_size)) {
+ return Status::Invalid("Buffer size for tensor would not fit in 64-bit
integer");
+ }
+ ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> result,
+ AllocateBuffer(buffer_size, pool));
// Copy data
switch (result_type->id()) {
case Type::UINT8:
- ConvertColumnsToTensor<UInt8Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<UInt8Type>(container, result->mutable_data(),
row_major);
break;
case Type::UINT16:
case Type::HALF_FLOAT:
- ConvertColumnsToTensor<UInt16Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<UInt16Type>(container, result->mutable_data(),
row_major);
break;
case Type::UINT32:
- ConvertColumnsToTensor<UInt32Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<UInt32Type>(container, result->mutable_data(),
row_major);
break;
case Type::UINT64:
- ConvertColumnsToTensor<UInt64Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<UInt64Type>(container, result->mutable_data(),
row_major);
break;
case Type::INT8:
- ConvertColumnsToTensor<Int8Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<Int8Type>(container, result->mutable_data(),
row_major);
break;
case Type::INT16:
- ConvertColumnsToTensor<Int16Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<Int16Type>(container, result->mutable_data(),
row_major);
break;
case Type::INT32:
- ConvertColumnsToTensor<Int32Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<Int32Type>(container, result->mutable_data(),
row_major);
break;
case Type::INT64:
- ConvertColumnsToTensor<Int64Type>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<Int64Type>(container, result->mutable_data(),
row_major);
break;
case Type::FLOAT:
- ConvertColumnsToTensor<FloatType>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<FloatType>(container, result->mutable_data(),
row_major);
break;
case Type::DOUBLE:
- ConvertColumnsToTensor<DoubleType>(batch, result->mutable_data(),
row_major);
+ ConvertColumnsToTensor<DoubleType>(container, result->mutable_data(),
row_major);
break;
default:
return Status::TypeError("DataType is not supported: ",
result_type->ToString());
@@ -412,7 +461,7 @@ Status RecordBatchToTensor(const RecordBatch& batch, bool
null_to_nan, bool row_
// Construct Tensor object
const auto& fixed_width_type =
internal::checked_cast<const FixedWidthType&>(*result_type);
- std::vector<int64_t> shape = {batch.num_rows(), batch.num_columns()};
+ std::vector<int64_t> shape = {container.num_rows(), container.num_columns()};
std::vector<int64_t> strides;
if (row_major) {
@@ -427,6 +476,16 @@ Status RecordBatchToTensor(const RecordBatch& batch, bool
null_to_nan, bool row_
return Status::OK();
}
+Status TableToTensor(const Table& table, bool null_to_nan, bool row_major,
+ MemoryPool* pool, std::shared_ptr<Tensor>* tensor) {
+ return ToTensorImpl(table, null_to_nan, row_major, pool, tensor);
+}
+
+Status RecordBatchToTensor(const RecordBatch& batch, bool null_to_nan, bool
row_major,
+ MemoryPool* pool, std::shared_ptr<Tensor>* tensor) {
+ return ToTensorImpl(batch, null_to_nan, row_major, pool, tensor);
+}
+
} // namespace internal
/// Constructor with strides and dimension names
diff --git a/cpp/src/arrow/tensor.h b/cpp/src/arrow/tensor.h
index beb62a11bd..1300003c29 100644
--- a/cpp/src/arrow/tensor.h
+++ b/cpp/src/arrow/tensor.h
@@ -77,6 +77,10 @@ Status ValidateTensorParameters(const
std::shared_ptr<DataType>& type,
const std::vector<int64_t>& strides,
const std::vector<std::string>& dim_names);
+ARROW_EXPORT
+Status TableToTensor(const Table& table, bool null_to_nan, bool row_major,
+ MemoryPool* pool, std::shared_ptr<Tensor>* tensor);
+
ARROW_EXPORT
Status RecordBatchToTensor(const RecordBatch& batch, bool null_to_nan, bool
row_major,
MemoryPool* pool, std::shared_ptr<Tensor>* tensor);
diff --git a/cpp/src/arrow/tensor_benchmark.cc
b/cpp/src/arrow/tensor_benchmark.cc
index 91a9270ef3..f064abf612 100644
--- a/cpp/src/arrow/tensor_benchmark.cc
+++ b/cpp/src/arrow/tensor_benchmark.cc
@@ -18,6 +18,7 @@
#include "benchmark/benchmark.h"
#include "arrow/record_batch.h"
+#include "arrow/table.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/testing/random.h"
#include "arrow/type.h"
@@ -51,6 +52,37 @@ static void BatchToTensorSimple(benchmark::State& state) {
state.SetBytesProcessed(state.iterations() * ty->byte_width() * num_rows *
num_cols);
}
+template <typename ValueType, bool row_major>
+static void TableToTensorSimple(benchmark::State& state) {
+ using CType = typename ValueType::c_type;
+ std::shared_ptr<DataType> ty = TypeTraits<ValueType>::type_singleton();
+
+ const int64_t num_cols = state.range(1);
+ const int64_t num_rows = state.range(0) / num_cols / sizeof(CType);
+ arrow::random::RandomArrayGenerator gen_{42};
+
+ std::vector<std::shared_ptr<Field>> fields = {};
+ std::vector<std::shared_ptr<ChunkedArray>> columns = {};
+
+ for (int64_t i = 0; i < num_cols; ++i) {
+ fields.push_back(field("f" + std::to_string(i), ty));
+ const int64_t chunk1_len = num_rows / 2;
+ const int64_t chunk2_len = num_rows - chunk1_len;
+ ArrayVector arrays = {gen_.ArrayOf(ty, chunk1_len), gen_.ArrayOf(ty,
chunk2_len)};
+ auto chunks = std::make_shared<ChunkedArray>(arrays, ty);
+ columns.push_back(chunks);
+ }
+ auto schema = std::make_shared<Schema>(std::move(fields));
+ auto table = Table::Make(schema, columns);
+
+ for (auto _ : state) {
+ ASSERT_OK_AND_ASSIGN(auto tensor,
+ table->ToTensor(/*null_to_nan=*/false,
/*row_major=*/row_major));
+ }
+ state.SetItemsProcessed(state.iterations() * num_rows * num_cols);
+ state.SetBytesProcessed(state.iterations() * ty->byte_width() * num_rows *
num_cols);
+}
+
void SetArgs(benchmark::internal::Benchmark* bench) {
for (int64_t size : {kL1Size, kL2Size}) {
for (int64_t num_columns : {3, 30, 300}) {
@@ -65,4 +97,13 @@ BENCHMARK_TEMPLATE(BatchToTensorSimple,
Int16Type)->Apply(SetArgs);
BENCHMARK_TEMPLATE(BatchToTensorSimple, Int32Type)->Apply(SetArgs);
BENCHMARK_TEMPLATE(BatchToTensorSimple, Int64Type)->Apply(SetArgs);
+#define DECLARE_TABLE_TO_TENSOR_BENCHMARKS(row_major)
\
+ BENCHMARK_TEMPLATE(TableToTensorSimple, Int8Type,
row_major)->Apply(SetArgs); \
+ BENCHMARK_TEMPLATE(TableToTensorSimple, Int16Type,
row_major)->Apply(SetArgs); \
+ BENCHMARK_TEMPLATE(TableToTensorSimple, Int32Type,
row_major)->Apply(SetArgs); \
+ BENCHMARK_TEMPLATE(TableToTensorSimple, Int64Type,
row_major)->Apply(SetArgs);
+
+DECLARE_TABLE_TO_TENSOR_BENCHMARKS(false);
+DECLARE_TABLE_TO_TENSOR_BENCHMARKS(true);
+
} // namespace arrow
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index f4ac1fd5ef..8b4786ecbf 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -1139,6 +1139,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
const shared_ptr[CSchema]& schema,
const vector[shared_ptr[CRecordBatch]]& batches)
+ CResult[shared_ptr[CTensor]] ToTensor(c_bool null_to_nan, c_bool
row_major,
+ CMemoryPool* pool) const
+
int num_columns()
int64_t num_rows()
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
index 5b3872bd26..fc7c4fcfc8 100644
--- a/python/pyarrow/table.pxi
+++ b/python/pyarrow/table.pxi
@@ -3634,7 +3634,7 @@ cdef class RecordBatch(_Tabular):
b: [10,20,30,40,null]
Convert a RecordBatch to row-major Tensor with null values
- written as NaN values
+ written as ``NaN``:
>>> batch.to_tensor(null_to_nan=True)
<pyarrow.Tensor>
@@ -3648,7 +3648,7 @@ cdef class RecordBatch(_Tabular):
[ 4., 40.],
[nan, nan]])
- Convert a RecordBatch to column-major Tensor
+ Convert a RecordBatch to column-major Tensor:
>>> batch.to_tensor(null_to_nan=True, row_major=False)
<pyarrow.Tensor>
@@ -3664,15 +3664,11 @@ cdef class RecordBatch(_Tabular):
"""
self._assert_cpu()
cdef:
- shared_ptr[CRecordBatch] c_record_batch
shared_ptr[CTensor] c_tensor
CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
- c_record_batch = pyarrow_unwrap_batch(self)
with nogil:
- c_tensor = GetResultValue(
-
<CResult[shared_ptr[CTensor]]>deref(c_record_batch).ToTensor(null_to_nan,
-
row_major, pool))
+ c_tensor = GetResultValue(self.batch.ToTensor(null_to_nan,
row_major, pool))
return pyarrow_wrap_tensor(c_tensor)
def copy_to(self, destination):
@@ -4994,7 +4990,7 @@ cdef class Table(_Tabular):
animals: string
----
n_legs: [[2,4,5,100],[2,4,5,100]]
- animals: [["Flamingo","Horse","Brittle
stars","Centipede"],["Flamingo","Horse","Brittle stars","Centipede"]]
+ animals: [["Flamingo",...,"Centipede"],["Flamingo",...,"Centipede"]]
"""
cdef:
vector[shared_ptr[CRecordBatch]] c_batches
@@ -5089,6 +5085,81 @@ cdef class Table(_Tabular):
return result
+ def to_tensor(self, c_bool null_to_nan=False, c_bool row_major=True,
MemoryPool memory_pool=None):
+ """
+ Convert to a :class:`~pyarrow.Tensor`.
+
+ Tables that can be converted have fields of type signed or unsigned
integer or float,
+ including all bit-widths.
+
+ ``null_to_nan`` is ``False`` by default and this method will raise an
error in case
+ any nulls are present. Tables with nulls can be converted with
``null_to_nan`` set to
+ ``True``. In this case null values are converted to ``NaN`` and
integer type arrays are
+ promoted to the appropriate float type.
+
+ Parameters
+ ----------
+ null_to_nan : bool, default False
+ Whether to write null values in the result as ``NaN``.
+ row_major : bool, default True
+ Whether resulting Tensor is row-major or column-major
+ memory_pool : MemoryPool, default None
+ For memory allocations, if required, otherwise use default pool
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> table = pa.table(
+ ... [
+ ... pa.chunked_array([[1, 2], [3, 4, None]], type=pa.int32()),
+ ... pa.chunked_array([[10, 20, 30], [40, None]],
type=pa.float32()),
+ ... ], names = ["a", "b"]
+ ... )
+
+ >>> table
+ pyarrow.Table
+ a: int32
+ b: float
+ ----
+ a: [[1,2],[3,4,null]]
+ b: [[10,20,30],[40,null]]
+
+ Convert a Table to row-major Tensor with null values written as
``NaN``:
+
+ >>> table.to_tensor(null_to_nan=True)
+ <pyarrow.Tensor>
+ type: double
+ shape: (5, 2)
+ strides: (16, 8)
+ >>> table.to_tensor(null_to_nan=True).to_numpy()
+ array([[ 1., 10.],
+ [ 2., 20.],
+ [ 3., 30.],
+ [ 4., 40.],
+ [nan, nan]])
+
+ Convert a Table to column-major Tensor
+
+ >>> table.to_tensor(null_to_nan=True, row_major=False)
+ <pyarrow.Tensor>
+ type: double
+ shape: (5, 2)
+ strides: (8, 40)
+ >>> table.to_tensor(null_to_nan=True, row_major=False).to_numpy()
+ array([[ 1., 10.],
+ [ 2., 20.],
+ [ 3., 30.],
+ [ 4., 40.],
+ [nan, nan]])
+ """
+ self._assert_cpu()
+ cdef:
+ shared_ptr[CTensor] c_tensor
+ CMemoryPool* pool = maybe_unbox_memory_pool(memory_pool)
+ with nogil:
+ c_tensor = GetResultValue(self.table.ToTensor(null_to_nan,
row_major, pool))
+ return pyarrow_wrap_tensor(c_tensor)
+
def to_reader(self, max_chunksize=None):
"""
Convert the Table to a RecordBatchReader.
diff --git a/python/pyarrow/tests/test_table.py
b/python/pyarrow/tests/test_table.py
index b267486fca..cb010f4387 100644
--- a/python/pyarrow/tests/test_table.py
+++ b/python/pyarrow/tests/test_table.py
@@ -998,7 +998,7 @@ def check_tensors(tensor, expected_tensor, type, size):
@pytest.mark.parametrize('typ_str', [
"uint8", "uint16", "uint32", "uint64",
"int8", "int16", "int32", "int64",
- "float32", "float64",
+ "float16", "float32", "float64",
])
def test_recordbatch_to_tensor_uniform_type(typ_str):
typ = np.dtype(typ_str)
@@ -1056,61 +1056,44 @@ def test_recordbatch_to_tensor_uniform_type(typ_str):
@pytest.mark.numpy
-def test_recordbatch_to_tensor_uniform_float_16():
- arr1 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
- arr2 = [10, 20, 30, 40, 50, 60, 70, 80, 90]
- arr3 = [100, 100, 100, 100, 100, 100, 100, 100, 100]
- batch = pa.RecordBatch.from_arrays(
- [
- pa.array(np.array(arr1, dtype=np.float16), type=pa.float16()),
- pa.array(np.array(arr2, dtype=np.float16), type=pa.float16()),
- pa.array(np.array(arr3, dtype=np.float16), type=pa.float16()),
- ], ["a", "b", "c"]
- )
-
- result = batch.to_tensor(row_major=False)
- x = np.column_stack([arr1, arr2, arr3]).astype(np.float16, order="F")
- expected = pa.Tensor.from_numpy(x)
- check_tensors(result, expected, pa.float16(), 27)
-
- result = batch.to_tensor()
- x = np.column_stack([arr1, arr2, arr3]).astype(np.float16, order="C")
- expected = pa.Tensor.from_numpy(x)
- check_tensors(result, expected, pa.float16(), 27)
-
-
[email protected]
-def test_recordbatch_to_tensor_mixed_type():
[email protected](
+ ('cls'),
+ [
+ (pa.Table),
+ (pa.RecordBatch)
+ ]
+)
+def test_to_tensor_mixed_type(cls):
# uint16 + int16 = int32
arr1 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
arr2 = [10, 20, 30, 40, 50, 60, 70, 80, 90]
arr3 = [100, 200, 300, np.nan, 500, 600, 700, 800, 900]
- batch = pa.RecordBatch.from_arrays(
+ tabular = cls.from_arrays(
[
pa.array(arr1, type=pa.uint16()),
pa.array(arr2, type=pa.int16()),
], ["a", "b"]
)
- result = batch.to_tensor(row_major=False)
+ result = tabular.to_tensor(row_major=False)
x = np.column_stack([arr1, arr2]).astype(np.int32, order="F")
expected = pa.Tensor.from_numpy(x)
check_tensors(result, expected, pa.int32(), 18)
- result = batch.to_tensor()
+ result = tabular.to_tensor()
x = np.column_stack([arr1, arr2]).astype(np.int32, order="C")
expected = pa.Tensor.from_numpy(x)
check_tensors(result, expected, pa.int32(), 18)
# uint16 + int16 + float32 = float64
- batch = pa.RecordBatch.from_arrays(
+ tabular = cls.from_arrays(
[
pa.array(arr1, type=pa.uint16()),
pa.array(arr2, type=pa.int16()),
pa.array(arr3, type=pa.float32()),
], ["a", "b", "c"]
)
- result = batch.to_tensor(row_major=False)
+ result = tabular.to_tensor(row_major=False)
x = np.column_stack([arr1, arr2, arr3]).astype(np.float64, order="F")
expected = pa.Tensor.from_numpy(x)
@@ -1120,7 +1103,7 @@ def test_recordbatch_to_tensor_mixed_type():
assert result.shape == expected.shape
assert result.strides == expected.strides
- result = batch.to_tensor()
+ result = tabular.to_tensor()
x = np.column_stack([arr1, arr2, arr3]).astype(np.float64, order="C")
expected = pa.Tensor.from_numpy(x)
@@ -1184,7 +1167,7 @@ def test_recordbatch_to_tensor_null():
)
with pytest.raises(
pa.ArrowTypeError,
- match="Can only convert a RecordBatch with no nulls."
+ match="Can only convert a Table or RecordBatch with no nulls."
):
batch.to_tensor()
@@ -1269,6 +1252,70 @@ def test_recordbatch_to_tensor_unsupported():
batch.to_tensor()
[email protected]
[email protected]('typ_str', [
+ "uint8", "uint16", "uint32", "uint64",
+ "int8", "int16", "int32", "int64",
+ "float16", "float32", "float64",
+])
+def test_table_to_tensor_uniform_type(typ_str):
+ arr1 = [[1, 2, 3], [4, 5, 6, 7, 8, 9]]
+ arr2 = [[10, 20], [30, 40, 50, 60, 70, 80, 90]]
+ arr3 = [[100, 100, 100, 100, 100, 100], [100, 100, 100]]
+ table = pa.Table.from_arrays(
+ [
+ pa.chunked_array(arr1, type=pa.from_numpy_dtype(typ_str)),
+ pa.chunked_array(arr2, type=pa.from_numpy_dtype(typ_str)),
+ pa.chunked_array(arr3, type=pa.from_numpy_dtype(typ_str)),
+ ], ["a", "b", "c"]
+ )
+
+ arr1_f = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+ arr2_f = [10, 20, 30, 40, 50, 60, 70, 80, 90]
+ arr3_f = [100, 100, 100, 100, 100, 100, 100, 100, 100]
+
+ result = table.to_tensor(row_major=False)
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="F")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 27)
+
+ result = table.to_tensor()
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="C")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 27)
+
+ # Test offset
+ table1 = table.slice(1)
+ arr1_f = [2, 3, 4, 5, 6, 7, 8, 9]
+ arr2_f = [20, 30, 40, 50, 60, 70, 80, 90]
+ arr3_f = [100, 100, 100, 100, 100, 100, 100, 100]
+
+ result = table1.to_tensor(row_major=False)
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="F")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 24)
+
+ result = table1.to_tensor()
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="C")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 24)
+
+ table2 = table.slice(1, 5)
+ arr1_f = [2, 3, 4, 5, 6]
+ arr2_f = [20, 30, 40, 50, 60]
+ arr3_f = [100, 100, 100, 100, 100]
+
+ result = table2.to_tensor(row_major=False)
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="F")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 15)
+
+ result = table2.to_tensor()
+ x = np.column_stack([arr1_f, arr2_f, arr3_f]).astype(typ_str, order="C")
+ expected = pa.Tensor.from_numpy(x)
+ check_tensors(result, expected, pa.from_numpy_dtype(typ_str), 15)
+
+
def _table_like_slice_tests(factory):
data = [
pa.array(range(5)),
