Optimization for the case float keys are integral and fit a compact non-negative range

Author: FrancescAlted

src/blosc2/groupby.py

@@ -264,10 +264,13 @@ def _execute_with_result_target(self, specs: list[_AggSpec]):
             fast = self._try_execute_cython_float_integral_key_f64_sum(specs)
             if fast is not None:
                 return fast
-            fast = self._try_execute_cython_float_hash(specs)
+            # Dense single-key path also covers integral-valued float keys with a
+            # compact non-negative range (e.g. float32 second/id columns), so try
+            # it before the generic float hash path, which is markedly slower.
+            fast = self._try_execute_dense_single_int_key(specs)
             if fast is not None:
                 return fast
-            fast = self._try_execute_dense_single_int_key(specs)
+            fast = self._try_execute_cython_float_hash(specs)
             if fast is not None:
                 return fast
 
@@ -1033,8 +1036,12 @@ def _try_execute_dense_single_int_key(self, specs: list[_AggSpec]):  # noqa: C90
         key_info = self.table._schema.columns_by_name[key_name]
         key_is_dict = self.table._is_dictionary_column(key_info)
         key_dtype = np.dtype(np.int32) if key_is_dict else getattr(key_info.spec, "dtype", None)
-        if key_dtype is None or key_dtype.kind not in "biu":
+        if key_dtype is None or key_dtype.kind not in "biuf":
             return None
+        # Float keys are accepted only when their values are integral and fit a
+        # compact non-negative range; per-chunk casting verifies integrality and
+        # bails (falling back to the hash path) on the first fractional value.
+        key_is_integral_float = (not key_is_dict) and key_dtype.kind == "f"
         if any(spec.op in {"min", "max"} and spec.input_col is not None for spec in specs):
             for spec in specs:
                 if spec.op in {"min", "max"} and spec.input_col is not None:
@@ -1107,6 +1114,16 @@ def ensure_size(size: int) -> bool:
             keys = np.asarray(raw_keys[live_mask])
             if keys.dtype.kind == "b":
                 keys = keys.astype(np.int8, copy=False)
+            elif key_is_integral_float:
+                # Non-finite keys (NaN/inf, e.g. a retained NaN group when
+                # dropna=False) and fractional keys both make the dense integer
+                # mapping invalid; defer to the generic/hash fallback.
+                if not np.all(np.isfinite(keys)):
+                    return None
+                keys_int = keys.astype(np.int64)
+                if not np.array_equal(keys_int, keys):
+                    return None
+                keys = keys_int
             if len(keys) == 0:
                 continue
             min_key = int(np.min(keys))
@@ -1159,7 +1176,12 @@ def ensure_size(size: int) -> bool:
         group_codes = np.nonzero(present)[0]
         rows = []
         for code in group_codes:
-            key_value = self.table._cols[key_name].decode(int(code)) if key_is_dict else _python_scalar(code)
+            if key_is_dict:
+                key_value = self.table._cols[key_name].decode(int(code))
+            elif key_is_integral_float:
+                key_value = float(code)
+            else:
+                key_value = _python_scalar(code)
             row = {key_name: key_value}
             for spec in specs:
                 state = states[spec.output_col]

tests/ctable/test_groupby.py

@@ -265,6 +265,31 @@ def test_groupby_float_integral_fast_path_falls_back_for_nan_group_when_kept():
     assert got[1][1] == 2.0
 
 
+@pytest.mark.parametrize("row_type", [Float64KeyRow, Float32KeyRow])
+def test_groupby_integral_float_key_dense_min_max(row_type):
+    # Non-negative integral float keys take the dense single-key path (the same
+    # one used for integer keys), which also covers min/max -- not just sum.
+    t = CTable(row_type, new_data=[(2.0, 5.0), (1.0, 9.0), (2.0, 3.0), (1.0, 4.0)])
+
+    out_max = t.group_by("key", sort=True).max("value")
+    out_min = t.group_by("key", sort=True).min("value")
+
+    assert rows(out_max) == [(1.0, 9.0), (2.0, 5.0)]
+    assert rows(out_min) == [(1.0, 4.0), (2.0, 3.0)]
+    # The key column keeps its original float dtype in the result.
+    assert out_max._cols["key"][:].dtype == t._cols["key"][:].dtype
+
+
+def test_groupby_integral_float_key_falls_back_for_negative_keys():
+    # Negative keys cannot use the dense (non-negative) mapping; the generic
+    # path must still produce correct max results.
+    t = CTable(Float64KeyRow, new_data=[(-1.0, 5.0), (-1.0, 8.0), (2.0, 3.0)])
+
+    out = t.group_by("key", sort=True).max("value")
+
+    assert rows(out) == [(-1.0, 8.0), (2.0, 3.0)]
+
+
 def test_group_reduce_object_keys_sort_with_none():
     groups, sizes = blosc2.group_reduce(
         np.array([None, "b", "a", "b"], dtype=object), sort=True, dropna=False