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 #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 <frim.alenka@gmail.com>
Co-authored-by: Alenka Frim <AlenkaF@users.noreply.github.com>
Co-authored-by: tadeja <tadeja@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Rok Mihevc <rok@mihevc.org>
Signed-off-by: AlenkaF <frim.alenka@gmail.com>

Author: 4 people

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>

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

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;
 

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) {

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_; }