NVIDIA / NVIDIA/cudf

[FEA] Support canonical arrow extension types: FixedShapeTensorType and VariableShapeTensorType

Open
#14,589 1 comment 2 reactions 0 assignees View on GitHub
feature request Python question
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**
Feeding data from the CPU to the GPU is a bottleneck especially for computer vision. I'd like to store satellite images as parquet with georeferencing information in some columns and a column with Arrow's new FixedShapeTensorType extension array and then load it with cudf and get zero copy benefits when passing the tensor to the GPU and pytorch.

However it looks like cudf can't interpret this type.
```
import pyarrow as pa

tensor_type = pa.fixed_shape_tensor(pa.int32(), (2, 2))
arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
storage = pa.array(arr, pa.list_(pa.int32(), 4))
tensor_array = pa.ExtensionArray.from_storage(tensor_type, storage)

data = [
pa.array([1, 2, 3]),
pa.array(['foo', 'bar', None]),
tensor_array,
]
my_schema = pa.schema([('f0', pa.int8()),
('f1', pa.string()),
('tensors_int', tensor_type)])
table = pa.Table.from_arrays(data, schema=my_schema)

table.cast(table.schema)
```

```
import cudf

cudf.DataFrame.from_arrow(table)
```

```
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[2], line 3
1 import cudf
----> 3 cudf.DataFrame.from_arrow(table)

File ~/miniforge3/envs/rapids-23.10/lib/python3.10/site-packages/nvtx/nvtx.py:115, in annotate.__call__..inner(*args, **kwargs)
112 @wraps(func)
113 def inner(*args, **kwargs):
114 libnvtx_push_range(self.attributes, self.domain.handle)
--> 115 result = func(*args, **kwargs)
116 libnvtx_pop_range(self.domain.handle)
117 return result

File ~/miniforge3/envs/rapids-23.10/lib/python3.10/site-packages/cudf/core/dataframe.py:5322, in DataFrame.from_arrow(cls, table)
5319 for col_meta in table.schema.pandas_metadata["column_indexes"]:
5320 col_index_names.append(col_meta["name"])
-> 5322 out = super().from_arrow(table)
5323 if col_index_names is not None:
5324 out._data._level_names = col_index_names

File ~/miniforge3/envs/rapids-23.10/lib/python3.10/site-packages/nvtx/nvtx.py:115, in annotate.__call__..inner(*args, **kwargs)
112 @wraps(func)
113 def inner(*args, **kwargs):
114 libnvtx_push_range(self.attributes, self.domain.handle)
--> 115 result = func(*args, **kwargs)
116 libnvtx_pop_range(self.domain.handle)
117 return result

File ~/miniforge3/envs/rapids-23.10/lib/python3.10/site-packages/cudf/core/frame.py:1053, in Frame.from_arrow(cls, data)
1036 cudf_category_frame = {
1037 name: build_categorical_column(
1038 cudf_dictionaries_columns[name],
(...)
1046 )
1047 }
1049 # Handle non-dict arrays
1050 cudf_non_category_frame = {
1051 name: col
1052 for name, col in zip(
-> 1053 data.column_names, libcudf.interop.from_arrow(data)
1054 )
1055 }
1057 result = {**cudf_non_category_frame, **cudf_category_frame}
1059 # There are some special cases that need to be handled
1060 # based on metadata.

File ~/miniforge3/envs/rapids-23.10/lib/python3.10/contextlib.py:79, in ContextDecorator.__call__..inner(*args, **kwds)
76 @wraps(func)
77 def inner(*args, **kwds):
78 with self._recreate_cm():
---> 79 return func(*args, **kwds)

File interop.pyx:199, in cudf._lib.interop.from_arrow()

RuntimeError: CUDF failure at:/opt/conda/conda-bld/work/cpp/src/interop/from_arrow.cu:87: Unsupported type_id conversion to cudf
```

**Describe the solution you'd like**
I want to be able to load parquet files with cudf that contain a column with this tensor type and then easily hand it off from cudf to pytorch.

**Describe alternatives you've considered**
There are other formats like zarr for N-D arrays and direct loading to gpu, but I don't think zero copy between cpu and gpu is supported https://xarray.dev/blog/xarray-kvikio

Or I can continue loading parquet files with references to cloud optimized geotiff files, which is a lot slower.

**Additional context**
docs for the type are here: https://arrow.apache.org/docs/python/generated/pyarrow.FixedShapeTensorArray.html
and others are looking at implementing it for dataloading https://github.com/huggingface/datasets/issues/5272
there's also a variable shape equal dimension number type which would be very useful for efficiently loading satellite imagery time series where the time length can vary a lot depending on the sample, or the height and width can vary a lot depending on the sensor resolution: https://arrow.apache.org/docs/format/CanonicalExtensions.html#variable-shape-tensor

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.