GH-40062: [C++][Python] Conversion of Table to Arrow Tensor

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

@@ -906,10 +906,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) {
@@ -940,10 +941,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>

cpp/src/arrow/table.cc

@@ -36,10 +36,12 @@
 #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"
 #include "arrow/util/checked_cast.h"
+#include "arrow/util/logging.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/vector.h"
 
@@ -346,6 +348,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_; }