[FEA] Support canonical arrow extension types: FixedShapeTensorType and VariableShapeTensorType
- 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
Assessment
This issue has not been assessed yet.