Describe the enhancement requested

Background

There is currently one way to convert tabular PyArrow data structure into an array with contiguous memory layout (_Tabular to numpy ndarray) and that is through the array protocol using the __array__ method.

In the PyArrow implementation of the array protocol we use NumPy to convert each of the table columns into a numpy array and then to stack the columns together into a 2-dimensional array. The method produces a row-major layout as we stack the columns by axis 1.

There is no method in Arrow C++ to convert Table/RecordBatch to a Tensor.

There exists another conversion producing a similar result. This conversion is used to convert PyArrow _Tabluar object to a pandas DataFrame. Pandas DataFrame has a 2D array layout structure defined for each data type present. This 2D data structure is called a block. Therefore a dataframe with 3 int8 columns and 2 float32 columns will have a structure of 2 blocks, first being an int8 block with size (3, n) and second being a float32 block of size (2, n).

If all columns in the PyArrow table were of the same data type we would get one 2D block that would have a contiguous memory layout.

The resulting tensor memory layout is in this case column-major. For row-major memory layout the design will need a bit more research and discussion.

In [1]: import pyarrow as pa
   ...: import numpy as np
   ...: 
   ...: # Construct a pyarrow.Table
   ...: arr_1 = pa.array([2, 4, 5, 100], type=pa.int8())
   ...: arr_2 = pa.array([1, 2, 3, 4], type=pa.int16())
   ...: names = ["arr_1", "arr_2"]
   ...: t = pa.table([arr_1, arr_2], names=names)
   ...: # Define common data type and cast all the columns
   ...: common_dtype = pa.int16()
   ...: column_arrays = [
   ...:     t.column(i).cast(target_type=common_dtype) for i in range(t.num_columns)
   ...: ]
   ...: 
   ...: options = dict(
   ...:     pool=None,
   ...:     strings_to_categorical=False,
   ...:     zero_copy_only=False,
   ...:     integer_object_nulls=False,
   ...:     date_as_object=True,
   ...:     timestamp_as_object=False,
   ...:     use_threads=True,
   ...:     deduplicate_objects=True,
   ...:     safe=True,
   ...:     split_blocks=False,
   ...:     self_destruct=False,
   ...:     maps_as_pydicts=None,
   ...:     coerce_temporal_nanoseconds=False
   ...: )
   ...: 
   ...: # CONVERT TABLE TO 2 DIMENSIONAL NDARRAY
   ...:
   ...: # Using pyarrow -> pandas conversion
   ...: r1 = pa.lib.table_to_blocks(options, pa.table(column_arrays, names=names), None, None)[0]["block"]
   ...: # Using array protocol
   ...: r2 = t.__array__()

In [2]: t
Out[2]: 
pyarrow.Table
arr_1: int8
arr_2: int16
----
arr_1: [[2,4,5,100]]
arr_2: [[1,2,3,4]]

In [3]: r1
Out[3]: 
array([[  2,   4,   5, 100],
       [  1,   2,   3,   4]], dtype=int16)

In [4]: r2
Out[4]: 
array([[  2,   1],
       [  4,   2],
       [  5,   3],
       [100,   4]], dtype=int16)

In [5]: r1.flags
Out[5]: 
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : False
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False

In [6]: r2.flags
Out[6]: 
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False

Feature proposal

Add a new feature in Arrow C++ to convert Arrow Table and RecordBatch to a Tensor. The conversion should:

• return an Arrow Tensor (not FixedShapeTensorArray!),
• support following data types: uint (8/16/32/64), int (8/16/32/64) and float (16/32/64)
• support NaNs,
• add an option to convert NULL to NaN (validity bitmaps not supported in the resulting Tensor),
• support column and row-major layout.

We also plan to implement the DLPack protocol for the Tensor class separately. With these features the ML libraries will have the option to consume matrices (2D tensors) directly or via the DLPack protocol without the need to do their own gymnastics.

Umbrella issue

This issue is an umbrella issue for smaller tasks created for easier reviews:

• [C++][Python] Basic conversion of RecordBatch to Arrow Tensor #40059
• [C++] Add benchmark for ToTensor conversions #40357
• [C++] Add TensorFromJSON helper function #40297
• [C++][Python] Basic conversion of RecordBatch to Arrow Tensor - add support for different data types #40060
• [C++][Python] Basic conversion of RecordBatch to Arrow Tensor - add option to cast NULL to NaN #40061
• [C++][Python] Basic conversion of RecordBatch to Arrow Tensor - add support for row-major #40866
• Add benchmark with mixed types and varying percentage of null values
• [C++][Python] Conversion of Table to Arrow Tensor #40062
• [C++][Python] Row-major conversion of Table/RecordBatch to Arrow Tensor #40063

Documentation issues:

• [Docs][C++][Python] Add initial documentation for RecordBatch::Tensor conversion #40841
• Extra examples, tutorials
• Cookbook

Component(s)

C++, Python