Helper for handling categorical codes.
(
array: "pa.Array",
)
| 360 | |
| 361 | |
| 362 | def _arrow_array_inf( |
| 363 | array: "pa.Array", |
| 364 | ) -> ArrayInf: |
| 365 | """Helper for handling categorical codes.""" |
| 366 | if not TYPE_CHECKING: |
| 367 | pa = import_pyarrow() |
| 368 | if not isinstance(array, pa.Array): # pylint: disable=E0606 |
| 369 | raise TypeError(f"Invalid input type: {type(array)}") |
| 370 | |
| 371 | mask, data = array.buffers() |
| 372 | jdata = make_array_interface( |
| 373 | data.address, |
| 374 | shape=(len(array),), |
| 375 | dtype=_arrow_npdtype()[array.type], |
| 376 | is_cuda=not data.is_cpu, |
| 377 | ) |
| 378 | |
| 379 | if mask is not None: |
| 380 | jmask: Optional[ArrayInf] = { |
| 381 | "data": (mask.address, True), |
| 382 | "typestr": "<t1", |
| 383 | "version": 3, |
| 384 | "strides": None, |
| 385 | "shape": (len(array),), |
| 386 | "mask": None, |
| 387 | } |
| 388 | if not mask.is_cpu: |
| 389 | jmask["stream"] = STREAM_PER_THREAD # type: ignore[index, typeddict-unknown-key] |
| 390 | else: |
| 391 | jmask = None |
| 392 | |
| 393 | jdata["mask"] = jmask |
| 394 | return jdata |
| 395 | |
| 396 | |
| 397 | def arrow_cat_inf( # pylint: disable=too-many-locals |
no test coverage detected