PERF: Optimize CSV categorical parsing when categories are known

When read_csv receives a CategoricalDtype with pre-specified categories,
map parsed values directly to category codes in a single pass using a
pre-built hash table, avoiding the factorize-then-recode steps.

For non-string category types (datetime, float edge cases, bool), the
optimization is attempted first and falls back to the existing path if
str() representations don't match the raw CSV tokens.

closes pandas-dev#17743

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: jbrockmendel

asv_bench/benchmarks/io/csv.py

@@ -9,6 +9,7 @@
 
 from pandas import (
     Categorical,
+    CategoricalDtype,
     DataFrame,
     Index,
     concat,
@@ -448,6 +449,12 @@ def time_convert_post(self, engine):
     def time_convert_direct(self, engine):
         read_csv(self.fname, engine=engine, dtype="category")
 
+    def time_convert_known_categories(self, engine):
+        # GH#17743
+        cats = ["aaaaaaaa", "bbbbbbb", "cccccccc", "dddddddd", "eeeeeeee"]
+        dtype = {col: CategoricalDtype(cats) for col in ["a", "b", "c"]}
+        read_csv(self.fname, engine=engine, dtype=dtype)
+
 
 class ReadCSVParseDates(StringIORewind):
     params = ["c", "python"]

doc/source/whatsnew/v3.1.0.rst

@@ -132,6 +132,7 @@ Performance improvements
 - Performance improvement in :func:`merge` with ``how="cross"`` (:issue:`38082`)
 - Performance improvement in :func:`merge` with ``how="left"`` (:issue:`64370`)
 - Performance improvement in :func:`merge` with ``sort=False`` for single-key ``how="left"``/``how="right"`` joins when the opposite join key is sorted, unique, and range-like (:issue:`64146`)
+- Performance improvement in :func:`read_csv` with ``engine="c"`` when parsing columns as :class:`CategoricalDtype` with known categories, by mapping parsed values directly to category codes in a single pass instead of factorizing and then recoding (:issue:`17743`)
 - Performance improvement in :func:`read_csv` with ``engine="c"`` when reading from binary file-like objects (e.g. PyArrow S3 file handles) by avoiding unnecessary ``TextIOWrapper`` wrapping (:issue:`46823`)
 - Performance improvement in :func:`read_html` and the Python CSV parser when ``thousands`` is set, fixing catastrophic regex backtracking on cells with many comma-separated digit groups followed by non-numeric text (:issue:`52619`)
 - Performance improvement in :func:`read_sas` by reading page header fields directly in Cython instead of falling back to Python (:issue:`47339`)

pandas/_libs/parsers.pyx

@@ -95,6 +95,7 @@ from pandas._libs.khash cimport (
 
 from pandas.errors import (
     EmptyDataError,
+    Pandas4Warning,
     ParserError,
     ParserWarning,
 )
@@ -1113,8 +1114,36 @@ cdef class TextReader:
                              kh_str_starts_t *na_hashset,
                              set na_fset, bint raise_on_invalid):
         if isinstance(dtype, CategoricalDtype):
-            # TODO: I suspect that _categorical_convert could be
-            # optimized when dtype is an instance of CategoricalDtype
+            if dtype.categories is not None:
+                # GH#17743: try optimized path for known categories.
+                # Convert categories to str and match against raw CSV tokens.
+                # For string/integer categories str() matches CSV text exactly.
+                # For other types (datetime, float, bool) it may not match,
+                # in which case we fall through to the general path.
+                codes, na_count, not_found_count = _categorical_convert_known(
+                    self.parser, i, start, end, na_filter, na_hashset,
+                    dtype.categories)
+                if not_found_count == 0 or (
+                    dtype.categories.dtype == object
+                ):
+                    # Either all values matched, or categories are strings
+                    # (so unmatched values are genuinely unexpected).
+                    if not_found_count > 0:
+                        warnings.warn(
+                            "Constructing a Categorical with a dtype and "
+                            "values containing non-null entries not in that "
+                            "dtype's categories is deprecated and will raise "
+                            "in a future version.",
+                            Pandas4Warning,
+                            stacklevel=find_stack_level(),
+                        )
+                    array_type = dtype.construct_array_type()
+                    cat = array_type._simple_new(codes, dtype=dtype)
+                    return cat, na_count
+
+                # Fall through: non-string type where str() didn't match
+                # all CSV tokens (e.g. datetime, float edge cases).
+
             codes, cats, na_count = _categorical_convert(
                 self.parser, i, start, end, na_filter, na_hashset)
 
@@ -1574,6 +1603,71 @@ cdef _categorical_convert(parser_t *parser, int64_t col,
     return np.asarray(codes), result, na_count
 
 
+@cython.wraparound(False)
+@cython.boundscheck(False)
+cdef _categorical_convert_known(parser_t *parser, int64_t col,
+                                int64_t line_start, int64_t line_end,
+                                bint na_filter,
+                                kh_str_starts_t *na_hashset,
+                                categories):
+    """Convert column data into codes using pre-specified categories.
+
+    When categories are known ahead of time, we build the hash table from
+    the categories directly and map parsed values to category codes in a
+    single pass, avoiding the factorize-then-recode steps.
+    """
+    cdef:
+        int na_count = 0
+        int not_found_count = 0
+        Py_ssize_t i, num_cats, lines
+        coliter_t it
+        const char *word = NULL
+
+        int64_t NA = -1
+        int64_t[::1] codes
+
+        int ret = 0
+        kh_str_t *table
+        khiter_t k
+
+    lines = line_end - line_start
+    codes = np.empty(lines, dtype=np.int64)
+
+    num_cats = len(categories)
+    # Keep references to encoded category strings to prevent GC
+    # while the hash table holds pointers to their buffers.
+    cat_bytes_list = [str(cat).encode("utf-8") for cat in categories]
+
+    # Build hash table: encoded category string -> category index
+    table = kh_init_str()
+    for i in range(num_cats):
+        k = kh_put_str(table, PyBytes_AsString(cat_bytes_list[i]), &ret)
+        table.vals[k] = <int64_t>i
+
+    with nogil:
+        coliter_setup(&it, parser, col, line_start)
+
+        for i in range(lines):
+            COLITER_NEXT(it, word)
+
+            if na_filter:
+                if kh_get_str_starts_item(na_hashset, word):
+                    na_count += 1
+                    codes[i] = NA
+                    continue
+
+            k = kh_get_str(table, word)
+            if k != table.n_buckets:
+                codes[i] = table.vals[k]
+            else:
+                # Value not in known categories
+                not_found_count += 1
+                codes[i] = NA
+
+    kh_destroy_str(table)
+    return np.asarray(codes), na_count, not_found_count
+
+
 # -> ndarray[f'|S{width}']
 cdef _to_fw_string(parser_t *parser, int64_t col, int64_t line_start,
                    int64_t line_end, int64_t width):