@@ -53,6 +53,34 @@ class _AggState:
5353 count : int = 0
5454
5555
56+ @dataclasses .dataclass
57+ class _Utf8KeyChunk :
58+ """A utf8 key-column chunk, factorized to chunk-local integer codes.
59+
60+ ``codes[i]`` indexes ``uniques`` (a ``StringDType`` array sorted
61+ ascending), so null detection, live-row masking, and per-chunk
62+ ``np.unique`` all run on int64 codes; only the (few) distinct strings are
63+ ever decoded. Produced by :meth:`CTableGroupBy._read_key_chunk` via
64+ ``Utf8Array.factorizer``.
65+ """
66+
67+ codes : np .ndarray
68+ uniques : np .ndarray
69+
70+ def __len__ (self ) -> int :
71+ return len (self .codes )
72+
73+ def take (self , mask : np .ndarray ) -> _Utf8KeyChunk :
74+ return _Utf8KeyChunk (self .codes [mask ], self .uniques )
75+
76+ def code_of (self , value : str ) -> int :
77+ """Code of *value* in this chunk, or -1 when absent (uniques are sorted)."""
78+ i = int (np .searchsorted (self .uniques , value ))
79+ if i < len (self .uniques ) and self .uniques [i ] == value :
80+ return i
81+ return - 1
82+
83+
5684def _is_column_like (value : Any ) -> bool :
5785 return isinstance (getattr (value , "_col_name" , None ), str )
5886
@@ -123,6 +151,10 @@ def __init__(
123151 self .dropna = bool (dropna )
124152 self .engine = engine
125153 self .chunk_size = chunk_size
154+ # Per-key incremental Utf8Factorizer instances, shared across the
155+ # chunk loop so the string vocabulary is built once (see
156+ # _read_key_chunk).
157+ self ._utf8_factorizers : dict [str , Any ] = {}
126158
127159 for name in self .keys :
128160 if name in table ._computed_cols :
@@ -503,7 +535,12 @@ def _execute_with_result_target(self, specs: list[_AggSpec]):
503535 if not np .any (live_mask ):
504536 continue
505537
506- keys_live = [np .asarray (values )[live_mask ] for values in raw_keys ]
538+ keys_live = [
539+ values .take (live_mask )
540+ if isinstance (values , _Utf8KeyChunk )
541+ else np .asarray (values )[live_mask ]
542+ for values in raw_keys
543+ ]
507544 n_live = len (keys_live [0 ])
508545 if n_live == 0 :
509546 continue
@@ -1505,49 +1542,138 @@ def _chunk_size(self) -> int:
15051542 if self .chunk_size <= 0 :
15061543 raise ValueError ("chunk_size must be positive" )
15071544 return int (self .chunk_size )
1545+ target = 1 << 20
15081546 chunks = getattr (self .table ._valid_rows , "chunks" , None )
1509- if chunks :
1510- return max (int (chunks [0 ]), 1 )
1511- return 65536
1547+ if not chunks :
1548+ return target
1549+ base = max (int (chunks [0 ]), 1 )
1550+ if base >= target :
1551+ return base
1552+ # Batching the loop at the raw validity chunk shape can be pathological
1553+ # (in-memory tables created small and grown by resize keep their tiny
1554+ # initial chunk shape, e.g. 64 rows), and even 64 Ki-row batches leave
1555+ # the per-batch Python bookkeeping (group display/merge) visible for
1556+ # multi-key groupbys. Scale up to ~1 Mi rows while staying
1557+ # chunk-aligned; per-column batch memory stays modest (8 MB per int64
1558+ # column).
1559+ return base * - (- target // base )
15121560
15131561 def _read_key_chunk (self , name : str , start : int , stop : int ) -> np .ndarray :
15141562 col_info = self .table ._schema .columns_by_name [name ]
15151563 if self .table ._is_dictionary_column (col_info ):
15161564 return np .asarray (self .table ._cols [name ].codes [start :stop ], dtype = np .int32 )
15171565 if self .table ._is_utf8_column (col_info ):
1566+ # Factorize the chunk from its raw offsets/bytes buffers: no row
1567+ # is decoded, only the distinct values (codes flow through the
1568+ # rest of the pipeline). The factorizer is shared across chunks
1569+ # so values seen before are hash-matched instead of re-sorted.
15181570 # Utf8Array is sized to the logical row count, not the physical
1519- # valid_rows capacity, so a chunk boundary can run past its end.
1520- # Rows beyond it are never live (the row can't have been written
1521- # without this column), so the padding value is never read live.
1571+ # valid_rows capacity, so a chunk boundary can run past its end;
1572+ # rows beyond it are never live (the row can't have been written
1573+ # without this column), so the padding code is never read live.
15221574 col = self .table ._cols [name ]
1575+ fact = self ._utf8_factorizers .get (name )
1576+ if fact is None :
1577+ fact = self ._utf8_factorizers [name ] = col .factorizer ()
15231578 n = len (col )
1524- values = np .asarray (col [start : min (stop , n )])
1579+ codes = fact .codes_for_span (start , min (stop , n )) if start < n else np .empty (0 , dtype = np .int64 )
1580+ uniques = fact .uniques ()
1581+ # Rank-normalize so this chunk keeps the np.unique contract the
1582+ # pipeline expects (uniques ascending, codes = string ranks).
1583+ order = np .argsort (uniques , kind = "stable" )
1584+ rank = np .empty (len (order ), dtype = np .int64 )
1585+ rank [order ] = np .arange (len (order ))
1586+ codes = rank [codes ] if len (order ) else codes
15251587 if stop > n :
1526- values = np .concatenate ([values , np .full (stop - max (start , n ), "" , dtype = values . dtype )])
1527- return values
1588+ codes = np .concatenate ([codes , np .zeros (stop - max (start , n ), dtype = np . int64 )])
1589+ return _Utf8KeyChunk ( codes , uniques [ order ])
15281590 return np .asarray (self .table ._cols [name ][start :stop ])
15291591
15301592 def _factorize_keys (
15311593 self , keys_live : list [np .ndarray ]
15321594 ) -> tuple [np .ndarray | list [np .ndarray ], np .ndarray ]:
15331595 if len (keys_live ) == 1 :
15341596 arr = keys_live [0 ]
1597+ if isinstance (arr , _Utf8KeyChunk ):
1598+ # The chunk is already factorized to dense string-rank codes;
1599+ # dedupe them with an O(n) bincount instead of a sort. The
1600+ # np.unique contract — uniques ascending by string — holds
1601+ # because the codes are ranks into the sorted uniques.
1602+ present = np .bincount (arr .codes , minlength = len (arr .uniques )) > 0
1603+ remap = np .cumsum (present ) - 1
1604+ return arr .uniques [present ], remap [arr .codes ]
15351605 if arr .dtype .kind in ("U" , "S" ) and arr .dtype .itemsize % 4 == 0 and arr .dtype .itemsize :
15361606 return _factorize_fixed_width_str (arr )
15371607 unique , inverse = np .unique (arr , return_inverse = True )
15381608 return unique , inverse
15391609
1540- # StringDType (utf8 columns) cannot be a structured-array field, so pack
1541- # those columns as plain Python-object arrays instead; np.unique still
1542- # dedupes them correctly by str equality.
1543- pack_arrs = [arr .astype (object ) if arr .dtype .kind == "T" else arr for arr in keys_live ]
1544- dtype = [(f"k{ i } , arr .dtype ) for i , arr in enumerate (pack_arrs )]
1545- packed = np .empty (len (pack_arrs [0 ]), dtype = dtype )
1546- for i , arr in enumerate (pack_arrs ):
1547- packed [f"k{ i } ] = arr
1548- unique , inverse = np .unique (packed , return_inverse = True )
1610+ # utf8 key chunks pack as their int64 codes (codes are string-rank
1611+ # within the chunk, so the packed sort order matches the string sort
1612+ # order); the codes in the deduped rows are mapped back to strings
1613+ # below, into object fields (StringDType cannot be a structured-array
1614+ # field).
1615+ pack_arrs = [arr .codes if isinstance (arr , _Utf8KeyChunk ) else arr for arr in keys_live ]
1616+ composite = self ._composite_int_factorize (pack_arrs )
1617+ if composite is not None :
1618+ unique_fields , inverse = composite
1619+ else :
1620+ dtype = [(f"k{ i } , arr .dtype ) for i , arr in enumerate (pack_arrs )]
1621+ packed = np .empty (len (pack_arrs [0 ]), dtype = dtype )
1622+ for i , arr in enumerate (pack_arrs ):
1623+ packed [f"k{ i } ] = arr
1624+ packed_unique , inverse = np .unique (packed , return_inverse = True )
1625+ unique_fields = [packed_unique [f"k{ i } ] for i in range (len (pack_arrs ))]
1626+ out_dtype = [
1627+ (f"k{ i } , object if isinstance (arr , _Utf8KeyChunk ) else pack_arrs [i ].dtype )
1628+ for i , arr in enumerate (keys_live )
1629+ ]
1630+ unique = np .empty (len (unique_fields [0 ]), dtype = out_dtype )
1631+ for i , arr in enumerate (keys_live ):
1632+ field = unique_fields [i ]
1633+ unique [f"k{ i } ] = arr .uniques [field ] if isinstance (arr , _Utf8KeyChunk ) else field
15491634 return unique , inverse
15501635
1636+ @staticmethod
1637+ def _composite_int_factorize (
1638+ pack_arrs : list [np .ndarray ],
1639+ ) -> tuple [list [np .ndarray ], np .ndarray ] | None :
1640+ """Dedup rows of all-integer key columns via one combined int64 key.
1641+
1642+ ``np.unique`` over a structured dtype compares rows field-by-field
1643+ through void comparisons and is ~an order of magnitude slower than
1644+ over a plain int64 array, so when every key column is integral and
1645+ the product of the per-column value ranges fits int64, combine the
1646+ columns into a single integer (Horner over the zero-based fields —
1647+ the combined sort order equals the structured lexicographic order)
1648+ and dedup that. Returns ``(per-field unique values, inverse)``, or
1649+ ``None`` when the keys don't fit this scheme.
1650+ """
1651+ if not all (arr .dtype .kind in "biu" for arr in pack_arrs ):
1652+ return None
1653+ if len (pack_arrs [0 ]) == 0 :
1654+ return None
1655+ mins = [int (arr .min ()) for arr in pack_arrs ]
1656+ spans = [int (arr .max ()) - mn + 1 for arr , mn in zip (pack_arrs , mins , strict = True )]
1657+ if math .prod (spans ) >= 1 << 62 :
1658+ return None
1659+ combined = np .zeros (len (pack_arrs [0 ]), dtype = np .int64 )
1660+ for arr , mn , span in zip (pack_arrs , mins , spans , strict = True ):
1661+ if arr .dtype .kind == "u" :
1662+ # Zero-base within the unsigned dtype first: the raw values may
1663+ # not fit int64, but value - min always fits (span is checked).
1664+ zero_based = (arr - arr .dtype .type (mn )).astype (np .int64 , copy = False )
1665+ else :
1666+ zero_based = arr .astype (np .int64 , copy = False ) - mn
1667+ combined *= span
1668+ combined += zero_based
1669+ unique_c , inverse = np .unique (combined , return_inverse = True )
1670+ unique_fields : list [np .ndarray ] = [None ] * len (pack_arrs ) # type: ignore[list-item]
1671+ rem = unique_c
1672+ for i in range (len (pack_arrs ) - 1 , - 1 , - 1 ):
1673+ unique_fields [i ] = (rem % spans [i ] + mins [i ]).astype (pack_arrs [i ].dtype )
1674+ rem = rem // spans [i ]
1675+ return unique_fields , inverse
1676+
15511677 def _resolve_sort (self , * , cheap : bool ) -> bool :
15521678 """Resolve the tri-state ``self.sort`` request for the current path.
15531679
@@ -1990,6 +2116,11 @@ def _result_spec_for_agg(self, spec: _AggSpec) -> SchemaSpec:
19902116 def _null_mask (self , name : str , values : np .ndarray , * , is_key : bool ) -> np .ndarray :
19912117 col_info = self .table ._schema .columns_by_name [name ]
19922118 spec = col_info .spec
2119+ if isinstance (values , _Utf8KeyChunk ):
2120+ null_value = getattr (spec , "null_value" , None )
2121+ if null_value is None :
2122+ return np .zeros (len (values ), dtype = bool )
2123+ return values .codes == values .code_of (null_value )
19932124 if isinstance (spec , DictionarySpec ):
19942125 mask = values == np .int32 (spec .null_code )
19952126 return mask if is_key or getattr (spec , "nullable" , False ) else np .zeros (len (values ), dtype = bool )
0 commit comments