seroverse.core¶

class ScalarView:
    r"""Concrete base class for immutable scalar views extracted from a [`BatchedContainer`][seroverse.core.collections.BatchedContainer].

    All scalar views are `@dataclass(frozen=True, slots=True)` records representing a single
    element of a Structure-of-Arrays (SoA) batch container. Inheriting from `ScalarView`
    automatically provides:
    - NumPy-safe structural equality (`__eq__`) using `np.array_equal` for any `np.ndarray` fields
      while preserving `@dataclass`-generated `__hash__` on scalar-only views.
    - Functional copy-update via [`.replace(**changes)`][seroverse.core.collections.ScalarView.replace].
    - Shallow dictionary serialization via [`.to_dict()`][seroverse.core.collections.ScalarView.to_dict].
    """

    __slots__ = ()
    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        if "__dataclass_fields__" in cls.__dict__:
            cls.__eq__ = ScalarView._dataclass_eq  # type: ignore[method-assign]

    def _dataclass_eq(self, other: object) -> bool:
        r"""Compare two ScalarView instances, handling both scalar and `np.ndarray` fields safely."""
        if type(self) is not type(other):
            return False
        for f in fields(self):
            if not f.compare:
                continue
            a = getattr(self, f.name)
            b = getattr(other, f.name)
            if isinstance(a, np.ndarray) or isinstance(b, np.ndarray):
                if not np.array_equal(a, b):
                    return False
            elif a != b:
                return False
        return True

    def replace(self, **changes: Any) -> Self:
        r"""Return a new scalar view instance with specified fields replaced."""
        return dataclass_replace(self, **changes)

    def to_dict(self) -> dict[str, Any]:
        r"""Return a shallow dictionary mapping field names to their values."""
        return {f.name: getattr(self, f.name) for f in fields(self)}

class BatchedContainer[T: ScalarView, S: BatchedContainer[Any, Any]]:
    r"""Concrete generic base class for Structure-of-Arrays (SoA) batch containers.

    Subclasses must be `@dataclass(frozen=True, slots=True)` classes whose stored fields are all
    annotated as 1D NumPy arrays (e.g. `Int32Array`, `UInt32Array`, `Float64Array`). Upon class
    creation, `BatchedContainer.__init_subclass__` validates and registers each field's `np.dtype`
    into `_field_dtypes` and automatically provides:
    - `__len__`, `__iter__`, and NumPy-safe `__eq__`
    - `__getitem__` with `@overload` typing, negative-index bounds checking, and automatic dispatch
      to `_get_scalar(idx)` (for integer indices) and `_get_slice(item)` (for slices/masks/index arrays)
    - `filter(mask)`, `replace(**changes)`, `empty()`, `concat(batches)`, and `from_list(items)`

    Type Parameters:
        T: The [`ScalarView`][seroverse.core.collections.ScalarView] dataclass returned when indexing a single element.
        S: The batched container type returned when slicing or filtering.
    """

    __slots__ = ()
    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
    _field_dtypes: ClassVar[dict[str, np.dtype[Any]]] = {}
    _field_ndims: ClassVar[dict[str, int]] = {}
    _meta_fields: ClassVar[tuple[str, ...]] = ()
    _primary_field: ClassVar[str] = ""
    _scalar_type: ClassVar[type[Any] | None] = None
    _scalar_source_fields: ClassVar[tuple[str, ...]] = ()
    _scalar_converters: ClassVar[tuple[Callable[[Any], Any] | None, ...]] = ()

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        # Extract scalar view type parameter T from generic bases if concrete
        for base in getattr(cls, "__orig_bases__", ()):
            origin = get_origin(base)
            if origin is not None and isinstance(origin, type) and issubclass(origin, BatchedContainer):
                args = get_args(base)
                if args and isinstance(args[0], type) and issubclass(args[0], ScalarView):
                    cls._scalar_type = args[0]
                    break

        if "__dataclass_fields__" not in cls.__dict__:
            return

        base_dtypes: dict[str, np.dtype[Any]] = {}
        base_ndims: dict[str, int] = {}
        for mro_cls in reversed(cls.__mro__[1:]):
            if issubclass(mro_cls, BatchedContainer) and "_field_dtypes" in mro_cls.__dict__:
                base_dtypes.update(mro_cls._field_dtypes)
                base_ndims.update(mro_cls._field_ndims)

        try:
            hints = annotationlib.get_annotations(cls, format=annotationlib.Format.VALUE)
        except Exception:
            hints = annotationlib.get_annotations(cls, format=annotationlib.Format.FORWARDREF)

        dtypes: dict[str, np.dtype[Any]] = {}
        ndims: dict[str, int] = {}
        meta: list[str] = []
        for f in fields(cls):
            if not f.init:
                continue
            if f.name in base_dtypes and f.name not in hints:
                dtypes[f.name] = base_dtypes[f.name]
                ndims[f.name] = base_ndims.get(f.name, 1)
                continue
            ann = hints.get(f.name, f.type)
            has_default = f.default is not MISSING or f.default_factory is not MISSING
            try:
                dtypes[f.name] = _resolve_array_dtype(ann, cls.__name__, f.name)
                ndims[f.name] = _resolve_array_ndim(ann)
            except TypeError:
                if has_default or (isinstance(ann, type) and issubclass(ann, BatchedContainer)):
                    meta.append(f.name)
                else:
                    raise

        cls._field_dtypes = dtypes
        cls._field_ndims = ndims
        cls._meta_fields = tuple(meta)
        cls._primary_field = next(iter(dtypes), "")

        scalar_cls = cls._scalar_type
        if scalar_cls is not None and "__dataclass_fields__" in scalar_cls.__dict__:
            s_fields = fields(scalar_cls)
            cls._scalar_source_fields = tuple(dtypes)[: len(s_fields)]
            s_hints: dict[str, Any] = {}
            for s_mro in reversed(scalar_cls.__mro__):
                if "__dataclass_fields__" in s_mro.__dict__:
                    try:
                        s_hints.update(annotationlib.get_annotations(s_mro, format=annotationlib.Format.VALUE))
                    except Exception:
                        pass
            convs: list[Callable[[Any], Any] | None] = []
            has_conv = False
            for sf in s_fields:
                sf_ann = s_hints.get(sf.name, sf.type)
                if isinstance(sf_ann, type) and issubclass(sf_ann, Enum):
                    convs.append(sf_ann)
                    has_conv = True
                else:
                    convs.append(None)
            cls._scalar_converters = tuple(convs) if has_conv else ()
        else:
            cls._scalar_source_fields = tuple(dtypes)
            cls._scalar_converters = ()

        cls.__eq__ = BatchedContainer._container_eq  # type: ignore[method-assign]

    def _container_eq(self, other: object) -> bool:
        r"""Compare two containers for exact array equality across all registered SoA fields."""
        if type(self) is not type(other):
            return False
        return all(np.array_equal(getattr(self, name), getattr(other, name)) for name in self._field_dtypes) and all(
            getattr(self, m) == getattr(other, m) for m in self._meta_fields
        )

    def __len__(self) -> int:
        r"""Return the number of records in the batch."""
        return len(getattr(self, self._primary_field))

    def _resolve_index(self, item: int | np.integer[Any]) -> int:
        r"""Normalize a potentially negative integer index and validate bounds."""
        idx = int(item)
        n = len(self)
        if idx < 0:
            idx += n
        if idx < 0 or idx >= n:
            raise IndexError(f"{type(self).__name__} index {item} out of range for batch of length {n}")
        return idx

    def _get_scalar(self, idx: int) -> T:
        r"""Materialize the scalar view `T` at normalized 0-based index `idx`."""
        scalar_cls = self._scalar_type
        if scalar_cls is None:
            raise NotImplementedError(f"{type(self).__name__} must implement _get_scalar(idx)")
        convs = self._scalar_converters
        if convs:
            return cast(
                T,
                scalar_cls(
                    *(
                        conv(getattr(self, f)[idx].item()) if conv is not None else getattr(self, f)[idx].item()
                        for f, conv in zip(self._scalar_source_fields, convs, strict=True)
                    )
                ),
            )
        return cast(T, scalar_cls(*(getattr(self, f)[idx].item() for f in self._scalar_source_fields)))

    def _get_slice(self, item: slice | np.ndarray[Any, Any] | list[Any]) -> S:
        r"""Slice all 1D SoA arrays by a slice, boolean mask, or integer index sequence."""
        indices: Any = item if isinstance(item, slice) else self._normalize_indices(item)
        return cast(S, type(self)(**{f: getattr(self, f)[indices] for f in self._field_dtypes}))

    @overload
    def __getitem__(self, item: int | np.integer[Any]) -> T: ...

    @overload
    def __getitem__(self, item: slice | np.ndarray[tuple[int], Any] | list[Any]) -> S: ...

    @overload
    def __getitem__(self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]) -> T | S: ...

    def __getitem__(self, item: int | np.integer[Any] | slice | np.ndarray[tuple[int], Any] | list[Any]) -> T | S:
        r"""Access a single scalar record `T` by integer index, or a sliced container `S` by slice/mask."""
        if isinstance(item, (int, np.integer)):
            return self._get_scalar(self._resolve_index(item))
        return self._get_slice(item)

    def __iter__(self) -> Iterator[T]:
        r"""Iterate over the container, yielding each scalar view `T` in order."""
        for i in range(len(self)):
            yield self._get_scalar(i)

    def _normalize_indices(self, item: slice | np.ndarray[Any, Any] | list[Any]) -> Int32Array:
        r"""Resolve a slice, boolean mask, or index list into a 1D `int32` index array."""
        return normalize_indices(item, len(self))

    def filter(self, mask: BoolArray | list[bool]) -> S:
        r"""Return a filtered container containing only records where `mask` is True."""
        return self._get_slice(mask)

    def replace(self, **changes: Any) -> Self:
        r"""Return a new container instance with specified array fields replaced."""
        return dataclass_replace(self, **changes)

    def to_dict(self) -> dict[str, Any]:
        r"""Return a shallow dictionary mapping container field names to their underlying NumPy arrays."""
        out: dict[str, Any] = {f: getattr(self, f) for f in self._field_dtypes}
        for m in self._meta_fields:
            out[m] = getattr(self, m)
        return out

    def chunks(self, size: int) -> Iterator[S]:
        r"""Yield contiguous $O(1)$ zero-copy mini-batch slices of at most `size` records.

        Args:
            size: Maximum number of records per chunk (must be `> 0`).

        Yields:
            S: Contiguous mini-batch container slice.

        Raises:
            ValueError: If `size <= 0`.
        """
        if size <= 0:
            raise ValueError(f"Chunk size must be positive, got {size}")
        n = len(self)
        for start in range(0, n, size):
            yield self._get_slice(slice(start, min(start + size, n)))

    def argsort(
        self,
        *by: str,
        descending: bool | Sequence[bool] = False,
    ) -> Int32Array:
        r"""Compute the stable 0-based `int32` permutation indices that sort this container by `*by`.

        Supports sorting by one or more registered 1D array columns or computed 1D array properties
        (in primary-to-secondary priority order), with either a single `descending` boolean or a
        per-key sequence of booleans.

        Args:
            *by: Field or property names to sort by (defaults to the primary container field).
            descending: Sort direction (`bool` applied to all keys, or a `Sequence[bool]` per key).

        Returns:
            Int32Array: 1D `int32` index permutation of length `len(self)`.

        Raises:
            ValueError: If a key does not resolve to a 1D NumPy array of length `len(self)` or
                if `descending` length does not match `len(by)`.
        """
        n = len(self)
        keys = by if by else (self._primary_field,)
        if isinstance(descending, bool):
            desc_flags = (descending,) * len(keys)
        else:
            desc_flags = tuple(bool(d) for d in descending)
            if len(desc_flags) != len(keys):
                raise ValueError(
                    f"Length of descending ({len(desc_flags)}) must match number of sort keys ({len(keys)})"
                )

        arrays: list[np.ndarray[Any, Any]] = []
        for k in keys:
            if not hasattr(self, k):
                raise ValueError(f"{type(self).__name__} has no field or property {k!r}")
            arr = getattr(self, k)
            if not isinstance(arr, np.ndarray) or arr.ndim != 1 or len(arr) != n:
                raise ValueError(f"Sort key {k!r} on {type(self).__name__} must be a 1D NumPy array of length {n}")
            arrays.append(arr)

        if n <= 1:
            return np.arange(n, dtype=np.int32)

        order = np.arange(n, dtype=np.int32)
        for arr, desc in zip(reversed(arrays), reversed(desc_flags), strict=True):
            sub = arr[order]
            if desc:
                step = (n - 1) - np.argsort(sub[::-1], kind="stable")[::-1]
            else:
                step = np.argsort(sub, kind="stable")
            order = order[step]
        return order.astype(np.int32, copy=False)

    def sort_by(
        self,
        *by: str,
        descending: bool | Sequence[bool] = False,
    ) -> S:
        r"""Return a new container sorted stably by one or more 1D array fields or properties."""
        if len(self) <= 1 and not by:
            return cast(S, self)
        return self._get_slice(self.argsort(*by, descending=descending))

    @classmethod
    def empty(cls) -> Self:
        r"""Create an empty 0-length container with registered `np.dtype` arrays."""
        ndims = cls._field_ndims
        return cls(
            **{
                f: np.empty((0,) * ndims.get(f, 1) if ndims.get(f, 1) > 1 else 0, dtype=dt)
                for f, dt in cls._field_dtypes.items()
            }
        )

    @classmethod
    def concat(cls, batches: Iterable[Self]) -> Self:
        r"""Concatenate multiple containers of the same type into a single batch."""
        batches_list = [b for b in batches if len(b) > 0]
        if not batches_list:
            return cls.empty()
        if len(batches_list) == 1:
            return batches_list[0]
        return cls(**{f: np.concatenate([getattr(b, f) for b in batches_list]) for f in cls._field_dtypes})

    @classmethod
    def from_list(cls, items: Sequence[T]) -> Self:
        r"""Pack a sequence of [`ScalarView`][seroverse.core.collections.ScalarView] instances into a batch container."""
        if not items:
            return cls.empty()
        scalar_cls = cls._scalar_type or type(items[0])
        s_fields = tuple(f.name for f in fields(scalar_cls))
        return cls(
            **{
                b_f: np.array([getattr(it, s_f) for it in items], dtype=dt)
                for (b_f, dt), s_f in zip(cls._field_dtypes.items(), s_fields, strict=True)
            }
        )

class RaggedArrayContainer[T: ScalarView, S: BatchedContainer[Any, Any]](BatchedContainer[T, S]):
    r"""Concrete generic base class for ragged Structure-of-Arrays (SoA) containers.

    Ragged containers store variable-length per-record arrays (declared in `_ragged_fields`)
    in flat, contiguous 1D NumPy buffers indexed by `offsets` and `lengths` (`Int32Array`),
    alongside any fixed-length per-record 1D arrays (`_record_fields`).

    Subclasses only need to declare `_ragged_fields: ClassVar[tuple[str, ...]]` (e.g. `("seqs",)`
    or `("members",)`), and inherit `__len__`, `__getitem__`, `_get_scalar`, `_get_slice`,
    `empty`, `concat`, `from_list`, `filter`, and `indptr` automatically.
    """

    __slots__ = ()
    offsets: Int32Array
    lengths: Int32Array
    _ragged_fields: ClassVar[tuple[str, ...]] = ()
    _ragged_fields_set: ClassVar[frozenset[str]] = frozenset()
    _record_fields: ClassVar[tuple[str, ...]] = ()
    _data_fields: ClassVar[tuple[str, ...]] = ()

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        if "__dataclass_fields__" not in cls.__dict__:
            return
        ragged_set = frozenset(cls._ragged_fields)
        cls._ragged_fields_set = ragged_set
        cls._data_fields = tuple(f for f in cls._field_dtypes if f not in ("offsets", "lengths"))
        cls._record_fields = tuple(f for f in cls._data_fields if f not in ragged_set)

    def __len__(self) -> int:
        r"""Return the number of records in the ragged container."""
        return len(self.offsets)

    @property
    def indptr(self) -> Int32Array:
        r"""CSR-compatible $(N + 1)$-length index pointer array (`[0, *cumsum(lengths)]`)."""
        n = len(self.offsets)
        if n == 0:
            return np.zeros(1, dtype=np.int32)
        out = np.empty(n + 1, dtype=np.int32)
        out[:n] = self.offsets
        out[n] = self.offsets[-1] + self.lengths[-1]
        return out

    @staticmethod
    def lengths_to_offsets(
        lengths: Int32Array,
    ) -> Int32Array:
        r"""Compute 0-based starting offsets from a 1D `lengths` array."""
        return lengths_to_offsets(lengths)

    def _get_scalar(self, idx: int) -> T:
        r"""Materialize the scalar view `T` at normalized 0-based index `idx`."""
        scalar_cls = self._scalar_type
        if scalar_cls is None:
            raise NotImplementedError(f"{type(self).__name__} must implement _get_scalar(idx)")
        off = int(self.offsets[idx])
        end = off + int(self.lengths[idx])
        ragged = self._ragged_fields_set
        return cast(
            T,
            scalar_cls(
                *(getattr(self, f)[off:end] if f in ragged else getattr(self, f)[idx].item() for f in self._data_fields)
            ),
        )

    def _get_slice(self, item: slice | np.ndarray[Any, Any] | list[Any]) -> S:
        r"""Slice both per-record arrays and variable-length ragged arrays in parallel."""
        ragged_arrays = tuple(getattr(self, f) for f in self._ragged_fields)
        indices, new_offsets, new_lengths, sliced_ragged = self._slice_ragged(item, *ragged_arrays)
        if len(indices) == 0:
            return cast(S, self.empty())
        kwargs: dict[str, Any] = {"offsets": new_offsets, "lengths": new_lengths}
        for f, arr in zip(self._ragged_fields, sliced_ragged, strict=True):
            kwargs[f] = arr
        for f in self._record_fields:
            kwargs[f] = getattr(self, f)[indices]
        return cast(S, type(self)(**kwargs))

    @classmethod
    def concat(cls, batches: Iterable[Self]) -> Self:
        r"""Concatenate multiple ragged containers into a single contiguous batch."""
        batches_list = [b for b in batches if len(b) > 0]
        if not batches_list:
            return cls.empty()
        if len(batches_list) == 1:
            return batches_list[0]
        all_lengths = np.concatenate([b.lengths for b in batches_list]).astype(np.int32, copy=False)
        kwargs: dict[str, Any] = {
            "offsets": lengths_to_offsets(all_lengths),
            "lengths": all_lengths,
        }
        for f in cls._ragged_fields:
            parts = [getattr(b, f) for b in batches_list if len(getattr(b, f)) > 0]
            kwargs[f] = np.concatenate(parts) if parts else np.empty(0, dtype=cls._field_dtypes[f])
        for f in cls._record_fields:
            kwargs[f] = np.concatenate([getattr(b, f) for b in batches_list])
        return cls(**kwargs)

    @classmethod
    def from_list(cls, items: Sequence[T]) -> Self:
        r"""Pack a sequence of [`ScalarView`][seroverse.core.collections.ScalarView] instances into a ragged container."""
        if not items:
            return cls.empty()
        scalar_cls = cls._scalar_type or type(items[0])
        s_fields = tuple(f.name for f in fields(scalar_cls))
        field_map = dict(zip(cls._data_fields, s_fields, strict=True))
        first_ragged_s = field_map[cls._ragged_fields[0]]
        lengths = np.array([len(getattr(it, first_ragged_s)) for it in items], dtype=np.int32)
        has_elems = bool(np.any(lengths > 0))
        kwargs: dict[str, Any] = {
            "offsets": lengths_to_offsets(lengths),
            "lengths": lengths,
        }
        for b_f in cls._ragged_fields:
            s_f = field_map[b_f]
            dt = cls._field_dtypes[b_f]
            parts = [np.asarray(getattr(it, s_f), dtype=dt) for it in items if len(getattr(it, s_f)) > 0]
            kwargs[b_f] = np.concatenate(parts) if (has_elems and parts) else np.empty(0, dtype=dt)
        for b_f in cls._record_fields:
            s_f = field_map[b_f]
            dt = cls._field_dtypes[b_f]
            kwargs[b_f] = np.array([getattr(it, s_f) for it in items], dtype=dt)
        return cls(**kwargs)

    def _slice_ragged(
        self,
        item: slice | np.ndarray[Any, Any] | list[Any],
        *arrays: np.ndarray[tuple[int], Any],
    ) -> tuple[
        Int32Array,
        Int32Array,
        Int32Array,
        tuple[np.ndarray[tuple[int], Any], ...],
    ]:
        r"""Slice or gather one or more parallel ragged data arrays while maintaining contiguous packing.

        Uses an $O(1)$ zero-copy view slice on `arrays` when `item` is a contiguous `slice(start, stop, 1)`,
        and falls back to the parallel Numba `_gather_ragged_kernel` for stepped slices, boolean masks,
        or arbitrary index arrays.

        Args:
            item: Slice, boolean mask, or integer index list/array.
            *arrays: Parallel 1D flat data arrays indexed by `self.offsets` and `self.lengths`.

        Returns:
            tuple: `(indices, new_offsets, new_lengths, sliced_arrays)`.
        """
        n = len(self)
        if isinstance(item, slice) and item.step in (None, 1):
            start, stop, _ = item.indices(n)
            if stop <= start:
                empty_i32 = np.empty(0, dtype=np.int32)
                return (
                    empty_i32,
                    empty_i32,
                    empty_i32,
                    tuple(np.empty(0, dtype=a.dtype) for a in arrays),
                )
            indices = np.arange(start, stop, dtype=np.int32)
            new_lengths = self.lengths[start:stop]
            base_offset = int(self.offsets[start])
            end_offset = int(self.offsets[stop - 1]) + int(self.lengths[stop - 1])
            new_offsets = self.offsets[start:stop] - base_offset
            return (
                indices,
                new_offsets,
                new_lengths,
                tuple(a[base_offset:end_offset] for a in arrays),
            )

        indices = self._normalize_indices(item)
        if len(indices) == 0:
            empty_i32 = np.empty(0, dtype=np.int32)
            return (
                indices,
                empty_i32,
                empty_i32,
                tuple(np.empty(0, dtype=a.dtype) for a in arrays),
            )

        new_lengths = self.lengths[indices].astype(np.int32, copy=False)
        new_offsets = lengths_to_offsets(new_lengths)
        total_len = int(new_offsets[-1] + new_lengths[-1])
        if total_len == 0:
            return (
                indices,
                new_offsets,
                new_lengths,
                tuple(np.empty(0, dtype=a.dtype) for a in arrays),
            )

        sliced = tuple(
            _gather_ragged_kernel(a, self.offsets, new_offsets, new_lengths, indices, total_len) for a in arrays
        )
        return indices, new_offsets, new_lengths, sliced

    def pad(
        self,
        field_name: str | None = None,
        *,
        max_len: int | None = None,
        fill_value: Any = 0,
    ) -> np.ndarray[tuple[int, int], Any]:
        r"""Pack a ragged 1D field into a dense 2D `(N, max_len)` NumPy matrix.

        Args:
            field_name: Name of the ragged field in `_ragged_fields` to pad. Defaults to the sole
                ragged field when `len(_ragged_fields) == 1`.
            max_len: Target column width $L$ of the output 2D matrix. Defaults to `max(self.lengths)`
                (or `0` when the container is empty); rows longer than `max_len` are truncated.
            fill_value: Padding value used to fill positions `[lengths[i] : max_len]`.

        Returns:
            np.ndarray[tuple[int, int], Any]: Dense 2D array of shape `(len(self), max_len)`.

        Raises:
            ValueError: If `field_name` is omitted on a multi-field ragged container, is not in
                `_ragged_fields`, or `max_len < 0`.
        """
        if field_name is None:
            if len(self._ragged_fields) != 1:
                raise ValueError(
                    f"{type(self).__name__} has {len(self._ragged_fields)} ragged fields {self._ragged_fields}; "
                    "specify field_name explicitly"
                )
            target_field = self._ragged_fields[0]
        else:
            if field_name not in self._ragged_fields_set:
                raise ValueError(
                    f"{field_name!r} is not a ragged field of {type(self).__name__} ({self._ragged_fields})"
                )
            target_field = field_name

        data = getattr(self, target_field)
        n = len(self)
        target_len = (int(np.max(self.lengths)) if n > 0 else 0) if max_len is None else int(max_len)
        if target_len < 0:
            raise ValueError(f"max_len must be non-negative, got {target_len}")
        if n == 0 or target_len == 0:
            return np.empty((n, max(target_len, 0)), dtype=data.dtype)
        if data.dtype.kind in ("S", "U", "O"):
            out = np.full((n, target_len), fill_value, dtype=data.dtype)
            for i in range(n):
                row_len = min(int(self.lengths[i]), target_len)
                if row_len > 0:
                    off = int(self.offsets[i])
                    out[i, :row_len] = data[off : off + row_len]
            return out
        typed_fill = data.dtype.type(fill_value)
        return _pad_ragged_kernel(data, self.offsets, self.lengths, target_len, typed_fill)