Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GH-43352: [Docs][Python] Add all tensor classes documentation #45160

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
85 changes: 85 additions & 0 deletions docs/source/python/api/tables.rst
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,92 @@ Dataframe Interchange Protocol
Tensors
-------

PyArrow supports both dense and sparse tensors. Dense tensors store all data values explicitly, while sparse tensors represent only the non-zero elements and their locations, making them efficient for storage and computation.

Dense Tensors
-------------
ShaiviAgarwal2 marked this conversation as resolved.
Show resolved Hide resolved

.. autosummary::
:toctree: ../generated/

Tensor

Sparse Tensors
--------------

PyArrow supports the following sparse tensor formats:

.. autosummary::
:toctree: ../generated/

SparseCOOTensor
SparseCSRMatrix
SparseCSCMatrix
SparseCSFTensor

### SparseCOOTensor
ShaiviAgarwal2 marked this conversation as resolved.
Show resolved Hide resolved

The `SparseCOOTensor` represents a sparse tensor in Coordinate (COO) format, where non-zero elements are stored as tuples of row and column indices.

Example:
.. code-block:: python

import pyarrow as pa

indices = pa.array([[0, 0], [1, 2]])
data = pa.array([1, 2])
shape = (2, 3)

tensor = pa.SparseCOOTensor(indices, data, shape)
print(tensor.to_dense())

### SparseCSRMatrix

The `SparseCSRMatrix` represents a sparse matrix in Compressed Sparse Row (CSR) format. This format is useful for matrix-vector multiplication.

Example:
.. code-block:: python

import pyarrow as pa

data = pa.array([1, 2, 3])
indptr = pa.array([0, 2, 3])
indices = pa.array([0, 2, 1])
shape = (2, 3)

sparse_matrix = pa.SparseCSRMatrix.from_numpy(data, indptr, indices, shape)
print(sparse_matrix)

### SparseCSCMatrix

The `SparseCSCMatrix` represents a sparse matrix in Compressed Sparse Column (CSC) format, where data is stored by columns.

Example:
.. code-block:: python

import pyarrow as pa

data = pa.array([1, 2, 3])
indptr = pa.array([0, 1, 3])
indices = pa.array([0, 1, 2])
shape = (3, 2)

sparse_matrix = pa.SparseCSCMatrix.from_numpy(data, indptr, indices, shape)
print(sparse_matrix)

### SparseCSFTensor

The `SparseCSFTensor` represents a sparse tensor in Compressed Sparse Fiber (CSF) format, which is a generalization of the CSR format for higher dimensions.

Example:
.. code-block:: python

import pyarrow as pa

data = pa.array([1, 2, 3])
indptr = [pa.array([0, 1, 3]), pa.array([0, 2, 3])]
indices = [pa.array([0, 1]), pa.array([0, 1, 2])]
shape = (2, 3, 2)

sparse_tensor = pa.SparseCSFTensor.from_numpy(data, indptr, indices, shape)
print(sparse_tensor)
20 changes: 19 additions & 1 deletion python/pyarrow/tensor.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -595,7 +595,25 @@ shape: {0.shape}""".format(self)

cdef class SparseCSRMatrix(_Weakrefable):
"""
A sparse CSR matrix.
SparseCSRMatrix represents a sparse matrix in Compressed Sparse Row (CSR) format.

Attributes:
indptr : array
Index pointer array.
indices : array
Column indices of the corresponding non-zero values.
shape : tuple
Shape of the matrix.
dim_names : list, optional
Names of the dimensions.

Example:
>>> import pyarrow as pa
>>> indptr = pa.array([0, 2, 3])
>>> indices = pa.array([0, 2, 1])
>>> shape = (2, 3)
>>> tensor = pa.SparseCSRMatrix(indptr, indices, shape)
>>> print(tensor)
ShaiviAgarwal2 marked this conversation as resolved.
Show resolved Hide resolved
"""

def __init__(self):
Expand Down
Loading