[Data] ParquetDatasink fails with tensor data and partition_cols

[schmidt-ai]

What happened + What you expected to happen

Writing data containing a tensor column to ParquetDatasink with partition column(s) fails because there is no pyarrow kernel for hash_list when called here.

This issue is solved by #50502.

Versions / Dependencies

MacOs 15.3
Python 3.10.16
ray==2.42.1
pyarrow==19.0.0

Reproduction script

from tempfile import TemporaryDirectory

import numpy as np
import ray.data

ctx = ray.data.DataContext.get_current()
ctx.use_arrow_tensor_v2 = False  # also fails if True

d = ray.data.from_items([{"partition_col": "a", "tensor": np.random.rand(5, 6)} for _ in range(10)])

with TemporaryDirectory() as tmpdir:
    d.write_parquet(tmpdir, partition_cols=["partition_col"])

pyarrow.lib.ArrowNotImplementedError: Function 'hash_list' has no kernel matching input types (extension<ray.data.arrow_tensor<ArrowTensorType>>, uint32)

This also fails if we explicitly cast to pyarrow.FixedShapeTensorArray:

from tempfile import TemporaryDirectory

import numpy as np
import pyarrow as pa
import ray.data


def to_arrow(batch):
    batch["tensor"] = pa.FixedShapeTensorArray.from_numpy_ndarray(batch["tensor"])
    return pa.Table.from_pydict(batch)


d = ray.data.from_items([{"partition_col": "a", "tensor": np.random.rand(5, 6).astype(np.float32)} for _ in range(10)])
d = d.map_batches(to_arrow)

with TemporaryDirectory() as tmpdir:
    d.write_parquet(tmpdir, partition_cols=["partition_col"])

pyarrow.lib.ArrowNotImplementedError: Function 'hash_list' has no kernel matching input types (extension<arrow.fixed_shape_tensor[value_type=float, shape=[5,6], permutation=[0,1]]>, uint32)

Implementing the idea from #50502 fixes things:

# imports from previous code snippet

class PartitionedParquetDatasink(_FileDatasink):
    def __init__(
        self,
        *args,
        partition_cols: list[str],
        schema: pa.Schema | None = None,
        write_to_dataset_kwargs: dict[str, Any] = {},
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.partition_cols = partition_cols
        self.schema = schema
        self.write_to_dataset_kwargs = write_to_dataset_kwargs

    def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None:
        write_to_dataset(
            table=pa.concat_tables(blocks),
            root_path=self.path,
            partition_cols=self.partition_cols,
            schema=self.schema,
            filesystem=self.filesystem,
            **self.write_to_dataset_kwargs,
        )

# stuff from previous snippet

with TemporaryDirectory() as tmpdir:
    d.write_datasink(PartitionedParquetDatasink(tmpdir, partition_cols=["partition_col"]))
    loaded = ray.data.read_parquet(tmpdir)
    loaded.show()

Issue Severity

Low: It annoys or frustrates me. I have a workaround using my implementation of PartitionedParquetDatasink.

[gvspraveen]

Thanks for reaching out. We are actively working on this issue.

Root cause is: underlying pyarrow.groupBy which doesnt support tensor types.

import pyarrow.parquet as pq

tensor_type = pa.fixed_shape_tensor(pa.int32(), [2, 2])

x = {"category": ["a", "b"] * 10, "tensor": list(np.random.random((20, 128)))}
schema = pa.schema(
    [
        ("category", pa.dictionary(pa.int32(), pa.string())),
        ("tensor", pa.fixed_shape_tensor(value_type=pa.float32(), shape=(128,))),
    ]
)
t = pa.Table.from_pydict(x, schema=schema)
t.group_by("category").aggregate([("tensor", "list")])