PERF: build read_sas string columns as pyarrow arrays directly (GH#47339) (pandas-dev#66464)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Author: jbrockmendel

doc/source/whatsnew/v3.1.0.rst

@@ -228,6 +228,7 @@ Performance improvements
 - Performance improvement in :func:`read_sas` for RLE- and RDC-compressed SAS7BDAT files (:issue:`47339`)
 - Performance improvement in :func:`read_sas` for SAS7BDAT files by pre-computing date/datetime column classification once during metadata parsing instead of per chunk (:issue:`47339`)
 - Performance improvement in :func:`read_sas` for SAS7BDAT files with full-precision (8-byte) numeric columns, with up to ~2x speedup on bulk reads (:issue:`47339`)
+- Performance improvement in :func:`read_sas` for SAS7BDAT files with string columns when an ``encoding`` is given, with up to ~7x speedup on bulk reads (:issue:`47339`)
 - Performance improvement in :func:`read_sas` for compressed SAS7BDAT files by reusing the decompression buffer instead of allocating per row (:issue:`47339`)
 - Performance improvement in :func:`read_sas` for compressed SAS7BDAT files with many rows per page, up to ~1.7x faster (:issue:`47339`)
 - Performance improvement in :func:`read_sas` when decoding strings (:issue:`47339`)

pandas/_libs/sas.pyi

@@ -1,7 +1,10 @@
+import numpy as np
+
 from pandas.io.sas.sas7bdat import SAS7BDATReader
 
 class Parser:
     def __init__(self, parser: SAS7BDATReader) -> None: ...
     def read(self, nrows: int) -> None: ...
+    def string_values(self) -> list[np.ndarray]: ...
 
 def collect_page_subheaders(parser: SAS7BDATReader) -> int: ...

pandas/_libs/sas.pyx

@@ -13,6 +13,7 @@ from libc.stdint cimport (
 from libc.stdlib cimport (
     calloc,
     free,
+    realloc,
 )
 from libc.string cimport (
     memcpy,
@@ -91,6 +92,41 @@ cdef buf_free(Buffer buf):
     if buf.data != NULL:
         free(buf.data)
 
+
+cdef struct StrColumn:
+    # Growable utf-8 value buffer backing one string column.
+    #
+    # Grown by doubling rather than preallocated to the declared column width
+    # times nrows: SAS pads character fields out to a fixed width, so the
+    # worst case badly overstates how many bytes the data actually needs.
+    uint8_t *data
+    int64_t size
+    int64_t capacity
+
+
+cdef bint strcol_reserve(StrColumn *col, int64_t extra) except 0:
+    """Ensure `col` has room for `extra` more bytes, growing by doubling."""
+    cdef:
+        int64_t needed = col.size + extra
+        int64_t capacity = col.capacity
+        uint8_t *data
+
+    # A zero capacity means data is still NULL, so allocate even when nothing is
+    # needed: an empty first cell would otherwise leave callers memcpy-ing into
+    # NULL, which is undefined behavior even for a length of zero.
+    if needed <= capacity and capacity > 0:
+        return True
+    if capacity == 0:
+        capacity = 4096
+    while capacity < needed:
+        capacity *= 2
+    data = <uint8_t *>realloc(col.data, capacity)
+    if data is NULL:
+        raise MemoryError(f"Failed to allocate {capacity} bytes")
+    col.data = data
+    col.capacity = capacity
+    return True
+
 # rle_decompress decompresses data using a Run Length Encoding
 # algorithm.  It is partially documented here:
 #
@@ -269,6 +305,12 @@ cdef:
     int page_data_type = const.page_data_type
     int subheader_pointers_offset = const.subheader_pointers_offset
 
+    # See sas_constants.py for what each string mode means. string_mode_table
+    # needs no alias here: it is whatever the other two are not.
+    int string_mode_object = const.string_mode_object
+    int string_mode_utf8 = const.string_mode_utf8
+    uint8_t undefined_byte = const.undefined_byte
+
     # Copy of subheader_signature_to_index that allows for much faster lookups.
     # Lookups are done in lookup_subheader_index_32/64. The C structures are
     # initialized in _init_subheader_signatures().
@@ -482,6 +524,24 @@ cdef class Parser:
         int64_t[:] column_types
         uint8_t[:, :] byte_chunk
         object[:, :] string_chunk
+        # Output for the pyarrow string fast path (string_mode_utf8/_table):
+        # str_cols[js] accumulates column js's utf-8 bytes, str_offsets[js, r+1]
+        # records the end of row r within it, and str_valid[js, r] is 0 for a
+        # null cell. Unused when str_mode is string_mode_object.
+        int str_mode
+        StrColumn *str_cols
+        int n_str_cols
+        int64_t[:, ::1] str_offsets
+        uint8_t[:, ::1] str_valid
+        # str_table[b * str_table_width : ...] holds the utf-8 for source byte
+        # b, str_table_len[b] its length (undefined_byte if b is not a valid
+        # character in str_encoding). Only set for string_mode_table. const
+        # because these are shared, read-only arrays cached across readers.
+        const uint8_t[::1] str_table
+        const uint8_t[::1] str_table_len
+        int str_table_width
+        bint str_ascii_identity
+        object str_encoding
         # (offset, length) of each compressed data subheader on the current
         # page, filled in by collect_page_subheaders and indexed by
         # current_row_on_page_index in readline().
@@ -510,6 +570,8 @@ cdef class Parser:
     def __cinit__(self, object parser):
         self.decompress = NULL
         self.decompressed_buf = Buffer(NULL, 0)
+        self.str_cols = NULL
+        self.n_str_cols = 0
 
     @cython.wraparound(False)
     @cython.boundscheck(False)
@@ -526,6 +588,22 @@ cdef class Parser:
         self.offsets = parser.column_data_offsets()
         self.byte_chunk = parser._byte_chunk
         self.string_chunk = parser._string_chunk
+        # sas7bdat.py always supplies these, using empty arrays for whichever
+        # of the two string outputs is inactive, so the memoryview assignments
+        # below never see None.
+        self.str_mode = parser._str_mode
+        self.str_offsets = parser._str_offsets
+        self.str_valid = parser._str_valid
+        self.str_table = parser._str_table
+        self.str_table_len = parser._str_table_len
+        self.str_table_width = parser._str_table_width
+        self.str_ascii_identity = parser._str_ascii_identity
+        self.str_encoding = parser._str_encoding
+        if self.str_mode != string_mode_object:
+            self.n_str_cols = self.str_offsets.shape[0]
+            self.str_cols = <StrColumn *>calloc(self.n_str_cols, sizeof(StrColumn))
+            if self.str_cols is NULL:
+                raise MemoryError("Failed to allocate string column buffers")
         self.data_subheader_offsets = parser._data_subheader_offsets
         self.data_subheader_lengths = parser._data_subheader_lengths
         self.row_length = parser.row_length
@@ -566,7 +644,34 @@ cdef class Parser:
         self.current_row_on_page_index = parser._current_row_on_page_index
 
     def __dealloc__(self):
+        cdef Py_ssize_t js
         buf_free(self.decompressed_buf)
+        if self.str_cols is not NULL:
+            for js in range(self.n_str_cols):
+                free(self.str_cols[js].data)
+            free(self.str_cols)
+
+    def string_values(self):
+        """
+        Return one exact-size uint8 array of utf-8 bytes per string column.
+
+        Copies out of the growable parse buffers so the returned arrays do not
+        pin their (up to 2x) spare capacity for the lifetime of the result.
+        """
+        cdef:
+            Py_ssize_t js
+            StrColumn *col
+            uint8_t[::1] out
+
+        result = []
+        for js in range(self.n_str_cols):
+            col = &self.str_cols[js]
+            values = np.empty(col.size, dtype=np.uint8)
+            if col.size > 0:
+                out = values
+                memcpy(&out[0], col.data, col.size)
+            result.append(values)
+        return result
 
     def read(self, int nrows):
         cdef:
@@ -719,13 +824,22 @@ cdef class Parser:
         cdef:
             Py_ssize_t j
             int s, k, m, jb, js, current_row, rpos
+            int bi, blen, table_width
             int64_t lngt, start, ct
+            uint8_t bval
+            bint ascii_identity, mode_object, mode_utf8
             Buffer source
+            StrColumn *col
             int64_t[:] column_types
             int64_t[:] lengths
             int64_t[:] offsets
             uint8_t[:, :] byte_chunk
             object[:, :] string_chunk
+            StrColumn *str_cols
+            int64_t[:, ::1] str_offsets
+            uint8_t[:, ::1] str_valid
+            const uint8_t[::1] str_table
+            const uint8_t[::1] str_table_len
             bint compressed
 
         assert offset + length <= self.cached_page_len, "Out of bounds read"
@@ -747,6 +861,17 @@ cdef class Parser:
         offsets = self.offsets
         byte_chunk = self.byte_chunk
         string_chunk = self.string_chunk
+        # Resolved once per row so that the per-cell branches below test a
+        # local rather than reloading the module-level mode constants.
+        mode_object = self.str_mode == string_mode_object
+        mode_utf8 = self.str_mode == string_mode_utf8
+        str_cols = self.str_cols
+        str_offsets = self.str_offsets
+        str_valid = self.str_valid
+        str_table = self.str_table
+        str_table_len = self.str_table_len
+        table_width = self.str_table_width
+        ascii_identity = self.str_ascii_identity
         s = 8 * self.current_row_in_chunk_index
         js = 0
         jb = 0
@@ -777,10 +902,57 @@ cdef class Parser:
                 # .rstrip(b"\x00 ") but without Python call overhead.
                 while lngt > 0 and buf_get(source, start + lngt - 1) in b"\x00 ":
                     lngt -= 1
-                if lngt == 0 and self.blank_missing:
-                    string_chunk[js, current_row] = np_nan
+                if mode_object:
+                    if lngt == 0 and self.blank_missing:
+                        string_chunk[js, current_row] = np_nan
+                    else:
+                        string_chunk[js, current_row] = buf_as_bytes(
+                            source, start, lngt
+                        )
                 else:
-                    string_chunk[js, current_row] = buf_as_bytes(source, start, lngt)
+                    col = &str_cols[js]
+                    if lngt > 0 or not self.blank_missing:
+                        # Spell out the bound the raw-pointer reads below rely
+                        # on; the rstrip loop's buf_get established it.
+                        assert start + lngt <= <int64_t>source.length, \
+                            "Out of bounds read"
+                        # Cells start null (str_valid is zero-initialized), so
+                        # only non-null ones need marking.
+                        str_valid[js, current_row] = 1
+                        strcol_reserve(col, lngt * table_width)
+                        if mode_utf8:
+                            memcpy(col.data + col.size, &source.data[start], lngt)
+                            col.size += lngt
+                        else:
+                            k = 0
+                            if ascii_identity:
+                                # Bytes below 0x80 are their own utf-8 in this
+                                # encoding, so an all-ascii cell (the common
+                                # case) costs one memcpy instead of a byte loop.
+                                while k < lngt and source.data[start + k] < 0x80:
+                                    k += 1
+                                memcpy(col.data + col.size, &source.data[start], k)
+                                col.size += k
+                            while k < lngt:
+                                bval = source.data[start + k]
+                                blen = str_table_len[bval]
+                                if blen == undefined_byte:
+                                    # Re-decode the cell so the error matches
+                                    # the UnicodeDecodeError that the per-cell
+                                    # strict decode used to raise.
+                                    buf_as_bytes(source, start, lngt).decode(
+                                        self.str_encoding
+                                    )
+                                    raise AssertionError(
+                                        "decoding an undefined byte should have raised"
+                                    )
+                                for bi in range(blen):
+                                    col.data[col.size + bi] = (
+                                        str_table[bval * table_width + bi]
+                                    )
+                                col.size += blen
+                                k += 1
+                    str_offsets[js, current_row + 1] = col.size
                 js += 1
 
         self.current_row_on_page_index += 1