Use native factorization for object string categories

news/476.performance.md

@@ -0,0 +1 @@
+Use native factorization for object arrays containing Python strings, preserving category order and handling mixed or missing values safely.

python/xy/channels.py

@@ -342,7 +342,22 @@ def _category_code_dtype(category_count: int) -> type[np.uint8] | type[np.uint32
     return np.uint8 if category_count <= MAX_CATEGORIES else np.uint32
 
 
-def _use_native_fixed_factorizer(arr: np.ndarray) -> bool:
+def _factorize_probe(arr: np.ndarray) -> np.ndarray:
+    """Return the bounded sample used to select the native factorizer."""
+    n = len(arr)
+    if n <= _FACTORIZE_PROBE_ROWS:
+        return arr
+    rows = np.linspace(0, n - 1, _FACTORIZE_PROBE_ROWS, dtype=np.intp)
+    return arr[rows]
+
+
+def _use_native_fixed_factorizer(
+    arr: np.ndarray,
+    *,
+    normalized_itemsize: Optional[int] = None,
+    probe: Optional[np.ndarray] = None,
+    distinct: Optional[int] = None,
+) -> bool:
     """Choose the O(N) hash path unless a bounded global probe says it cannot pay.
 
     The native pass earns its keep by keeping N records out of Python: only the
@@ -356,21 +371,43 @@ def _use_native_fixed_factorizer(arr: np.ndarray) -> bool:
     repeats get scarce while narrow ones hold it until the probe is entirely
     distinct. Sampling across the full array keeps the decision independent of N.
     """
-    n = len(arr)
-    if n <= _FACTORIZE_PROBE_ROWS:
-        probe = arr
-    else:
-        rows = np.linspace(0, n - 1, _FACTORIZE_PROBE_ROWS, dtype=np.intp)
-        probe = arr[rows]
-    distinct = len(np.unique(probe))
+    probe = _factorize_probe(arr) if probe is None else probe
+    distinct = len(np.unique(probe)) if distinct is None else distinct
     if distinct <= _FACTORIZE_NATIVE_MAX_PROBE_CATEGORIES:
         return True
     near_unique = (
-        1.0 if arr.dtype.itemsize <= _FACTORIZE_NARROW_ITEMSIZE else _FACTORIZE_NEAR_UNIQUE_RATIO
+        1.0
+        if (normalized_itemsize or arr.dtype.itemsize) <= _FACTORIZE_NARROW_ITEMSIZE
+        else _FACTORIZE_NEAR_UNIQUE_RATIO
     )
     return distinct < near_unique * len(probe)
 
 
+def _object_string_width(arr: np.ndarray, *, sample: bool = False) -> Optional[int]:
+    """Return the width of safe Python-string values, or ``None``.
+
+    Object arrays can contain values with display-label semantics that cannot
+    be represented by a simple Unicode cast. Restrict this path to exact
+    Python strings so mixed objects, missing values, and custom ``__str__``
+    implementations retain the canonical fallback behavior. NUL codepoints
+    are excluded because NumPy treats them as fixed-width string padding.
+    """
+    if arr.dtype.kind != "O" or arr.size == 0:
+        return None
+    values = _factorize_probe(arr) if sample else arr.reshape(-1)
+    width = 0
+    for value in values:
+        if type(value) is not str or "\x00" in value:
+            return None
+        width = max(width, len(value))
+    return width
+
+
+def _normalize_object_strings(arr: np.ndarray, width: int) -> np.ndarray:
+    """Normalize validated object-backed strings to fixed-width Unicode."""
+    return np.asarray(arr, dtype=f"<U{max(width, 1)}")
+
+
 def _factorize_categories(
     arr: np.ndarray,
 ) -> tuple[
@@ -385,18 +422,46 @@ def _factorize_categories(
     not mutually orderable. Chart labels are strings on the client anyway, so
     canonicalize to display labels first, sort those labels for deterministic
     palettes, and then map each row back to its code. Fixed-width NumPy
-    strings/bytes/bools can identify equal records in Rust without creating N
-    Python objects; only their compact unique set crosses the label-policy path.
+    strings/bytes/bools and safe object-backed Python strings can identify equal
+    records in Rust without creating N Python objects; only their compact unique
+    set crosses the label-policy path. Mixed objects, missing values, custom
+    string semantics, and strings containing NUL codepoints retain the Python
+    fallback so display-label behavior is unchanged.
     """
-    if arr.dtype.kind in ("U", "S", "b") and _use_native_fixed_factorizer(arr):
+    factorizer_arr = arr
+    if arr.dtype.kind == "O":
+        # Probe object strings before the full validation and Unicode copy so
+        # near-unique columns take the existing Python fallback cheaply.
+        probe = distinct = None
+        sample_width = _object_string_width(arr, sample=True)
+        if sample_width is not None:
+            probe = _factorize_probe(arr)
+            distinct = len(np.unique(probe))
+        if sample_width is not None and _use_native_fixed_factorizer(
+            arr,
+            normalized_itemsize=4 * max(sample_width, 1),
+            probe=probe,
+            distinct=distinct,
+        ):
+            width = _object_string_width(arr)
+            if width is not None and _use_native_fixed_factorizer(
+                arr,
+                normalized_itemsize=4 * max(width, 1),
+                probe=probe,
+                distinct=distinct,
+            ):
+                factorizer_arr = _normalize_object_strings(arr, width)
+    if factorizer_arr.dtype.kind in ("U", "S", "b") and _use_native_fixed_factorizer(
+        factorizer_arr
+    ):
         compact = (
-            kernels.factorize_unicode1_u8_counts(arr, MAX_CATEGORIES)
-            if arr.dtype.kind == "U" and arr.dtype.itemsize == 4
-            else kernels.factorize_fixed_u8_counts(arr, MAX_CATEGORIES)
+            kernels.factorize_unicode1_u8_counts(factorizer_arr, MAX_CATEGORIES)
+            if factorizer_arr.dtype.kind == "U" and factorizer_arr.dtype.itemsize == 4
+            else kernels.factorize_fixed_u8_counts(factorizer_arr, MAX_CATEGORIES)
         )
         if compact is not None:
             raw_codes, unique_indices, raw_counts = compact
-            unique_labels = [category_label(value) for value in arr[unique_indices]]
+            unique_labels = [category_label(value) for value in factorizer_arr[unique_indices]]
             categories = sorted(set(unique_labels))
             index = {label: i for i, label in enumerate(categories)}
             remap = np.fromiter(
@@ -412,8 +477,8 @@ def _factorize_categories(
                 counts[index[label]] += count
             return categories, raw_codes, counts
 
-        raw_codes, unique_indices = kernels.factorize_fixed(arr)
-        unique_labels = [category_label(value) for value in arr[unique_indices]]
+        raw_codes, unique_indices = kernels.factorize_fixed(factorizer_arr)
+        unique_labels = [category_label(value) for value in factorizer_arr[unique_indices]]
         categories = sorted(set(unique_labels))
         index = {label: i for i, label in enumerate(categories)}
         dtype = _category_code_dtype(len(categories))

tests/test_custom_ramps_and_palette.py

@@ -375,6 +375,98 @@ def test_the_literal_color_probe_does_not_materialize_category_columns():
     assert channels._literal_color_rgba(column) is None
 
 
+def test_homogeneous_object_categories_use_native_factorization(monkeypatch):
+    """Object strings use the native path without an N-entry label list."""
+    values = np.array(["group-a", "group-b", "group-a", "group-c"] * 250, dtype=object)
+    seen: list[object] = []
+    original = channels.category_label
+
+    def record(value: object) -> str:
+        seen.append(value)
+        return original(value)
+
+    monkeypatch.setattr(channels, "category_label", record)
+    categories, codes, counts = channels._factorize_categories(values)
+
+    assert categories == ["group-a", "group-b", "group-c"]
+    np.testing.assert_array_equal(codes[:4], [0, 1, 0, 2])
+    np.testing.assert_array_equal(counts, [500, 250, 250])
+    assert len(seen) == len(categories)
+
+
+def test_object_factorization_matches_fixed_width_string_semantics():
+    """The object fast path preserves category order, codes, and counts."""
+    values = np.array(["beta", "alpha", "beta", "gamma", "alpha"], dtype=object)
+    object_result = channels._factorize_categories(values)
+    unicode_result = channels._factorize_categories(values.astype("U5"))
+
+    assert object_result[0] == unicode_result[0]
+    np.testing.assert_array_equal(object_result[1], unicode_result[1])
+    np.testing.assert_array_equal(object_result[2], unicode_result[2])
+
+
+def test_object_factorization_preserves_nul_distinct_categories():
+    """NUL-containing strings retain Python equality semantics."""
+    categories, codes, counts = channels._factorize_categories(
+        np.array(["a", "a\x00", "a", "a\x00"], dtype=object)
+    )
+
+    assert categories == ["a", "a\x00"]
+    np.testing.assert_array_equal(codes, [0, 1, 0, 1])
+    assert counts is None
+
+
+def test_near_unique_object_strings_do_not_run_full_width_scan(monkeypatch):
+    """Near-unique object columns are rejected before full normalization."""
+    values = np.array([f"id-{i}" for i in range(5000)], dtype=object)
+    original = channels._object_string_width
+    calls: list[bool] = []
+
+    def record(arr, *, sample=False):
+        calls.append(sample)
+        return original(arr, sample=sample)
+
+    monkeypatch.setattr(channels, "_object_string_width", record)
+    channels._factorize_categories(values)
+
+    assert calls == [True]
+
+
+def test_object_factorization_reuses_the_bounded_probe(monkeypatch):
+    """Object-string eligibility gates share one bounded distinct-count probe."""
+    values = np.array(["group-a", "group-b", "group-a", "group-c"] * 1500, dtype=object)
+    original = channels.np.unique
+    object_probe_calls: list[np.ndarray] = []
+
+    def record(arr, *args, **kwargs):
+        if arr.dtype.kind == "O":
+            object_probe_calls.append(arr)
+        return original(arr, *args, **kwargs)
+
+    monkeypatch.setattr(channels.np, "unique", record)
+    channels._factorize_categories(values)
+
+    assert len(object_probe_calls) == 1
+    assert len(object_probe_calls[0]) == channels._FACTORIZE_PROBE_ROWS
+
+
+@pytest.mark.parametrize(
+    "values",
+    [
+        np.array(["a", None, "a"], dtype=object),
+        np.array(["a", 1, "a"], dtype=object),
+        np.array([b"a", b"b", b"a"], dtype=object),
+    ],
+)
+def test_mixed_object_categories_keep_the_fallback(values):
+    """Values without exact string semantics remain on the safe fallback."""
+    categories, codes, counts = channels._factorize_categories(values)
+
+    assert len(categories) == len(set(categories))
+    assert len(codes) == len(values)
+    assert counts is None
+
+
 def test_the_probe_still_reads_a_column_that_actually_looks_like_paint():
     column = np.array(["#ff0000", "#00ff00", "#0000ff"])
     rgba = channels._literal_color_rgba(column)