Address the Copilot review on PR #684

Three findings, all real:

- @blosc2.jit forwarded every decorator kwarg to blosc2.asarray() when the
  traced function returned a NumPy array, so combining a storage kwarg with
  an execution-tuning one -- @blosc2.jit(jit=False, cparams=...) -- raised
  instead of returning an NDArray.  Only storage kwargs go there now; the
  function has already run, so there is nothing left to tune.  The sibling
  compute() call keeps taking all of them on purpose: compute() names
  fp_accuracy while lazyudf() does not, so stripping them would drop it.

- The mandelbrot benchmark's header still said any non-None jit() kwarg
  flips the return container.  Only storage kwargs do.

- _fill_me_udata leaked np_data/np_typesizes on every allocation-failure
  exit after the one that allocates them.  All the exits now go through one
  _free_me_udata_tables() helper, so a table added later cannot be missed by
  some of them again.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: FrancescAlted

RELEASE_NOTES.md

@@ -166,6 +166,12 @@ XXX version-specific blurb XXX
 
 ### Bug fixes
 
+- **`@blosc2.jit` raised when a storage kwarg and an execution-tuning kwarg
+  were combined** and the decorated function returned a NumPy array —
+  `@blosc2.jit(jit=False, cparams=...)` ended in
+  `blosc2.asarray(retval, jit=False, ...)`, which rejects the tuning kwargs.
+  Only storage kwargs reach `asarray()` now; the function has already run, so
+  there is nothing left to tune.
 - **A DSL kernel over a utf8 column registered as a computed column, then
   broke the table.** `add_computed_column(name, kernel, inputs=["utf8_col"])`
   was accepted, after which every read of that column *and* `str(table)`

bench/ndarray/jit-dsl-mandelbrot.py

@@ -12,9 +12,11 @@
 # doc/guides/optimization_tips.md ("Let @blosc2.jit compile control flow
 # instead of tracing it").
 #
-# Return paths are equalized (both calls end in a plain NumPy array): any
-# non-None jit() kwarg flips the return from `retval[()]` to `.compute()`,
-# which would otherwise skew the comparison.
+# Return paths are equalized (both calls end in a plain NumPy array): a storage
+# kwarg (cparams/chunks/urlpath/...) flips the return from `retval[()]` to
+# `.compute()`, i.e. to an NDArray, which would otherwise skew the comparison.
+# Execution-tuning kwargs (jit/jit_backend/fp_accuracy) do not, so `@blosc2.jit`
+# is used bare here.
 
 from __future__ import annotations
 

src/blosc2/blosc2_ext.pyx

@@ -1027,6 +1027,22 @@ def me_output_dtype(expression, operands):
         free(variables)
 
 
+cdef inline void _free_me_udata_tables(me_udata* udata, b2nd_array_t** inputs_,
+                                       uint8_t** np_data, int32_t* np_typesizes):
+    """Release the per-input tables and the udata block itself.
+
+    ``_fill_me_udata`` has half a dozen allocation-failure exits; routing them
+    all through here is what keeps the tables from drifting apart, which they
+    already had once -- ``np_data``/``np_typesizes`` arrived later and every
+    exit but the first kept freeing only ``inputs``.  ``free(NULL)`` is a no-op,
+    so no exit needs to know which tables it got as far as allocating.
+    """
+    free(inputs_)
+    free(np_data)
+    free(np_typesizes)
+    free(udata)
+
+
 cdef inline str _me_compile_status_name(int rc):
     if rc == ME_COMPILE_SUCCESS:
         return "ME_COMPILE_SUCCESS"
@@ -4155,10 +4171,7 @@ cdef class NDArray:
             np_data = <uint8_t**> calloc(ninputs, sizeof(uint8_t*))
             np_typesizes = <int32_t*> calloc(ninputs, sizeof(int32_t))
             if np_data == NULL or np_typesizes == NULL:
-                free(inputs_)
-                free(np_data)
-                free(np_typesizes)
-                free(udata)
+                _free_me_udata_tables(udata, inputs_, np_data, np_typesizes)
                 raise MemoryError("Cannot allocate miniexpr raw-input tables")
         for i, operand in enumerate(operands):
             if isinstance(operand, np.ndarray):
@@ -4176,8 +4189,7 @@ cdef class NDArray:
         if ninputs > 0:
             input_chunk_caches = <me_input_cache_s*> calloc(ninputs, sizeof(me_input_cache_s))
             if input_chunk_caches == NULL:
-                free(inputs_)
-                free(udata)
+                _free_me_udata_tables(udata, inputs_, np_data, np_typesizes)
                 raise MemoryError("Cannot allocate miniexpr chunk caches")
             for i in range(ninputs):
                 input_chunk_caches[i].nchunk = -1
@@ -4191,8 +4203,7 @@ cdef class NDArray:
                         if input_chunk_caches[i].ready_lock != NULL:
                             PyThread_free_lock(input_chunk_caches[i].ready_lock)
                     free(input_chunk_caches)
-                    free(inputs_)
-                    free(udata)
+                    _free_me_udata_tables(udata, inputs_, np_data, np_typesizes)
                     raise MemoryError("Cannot allocate miniexpr chunk cache state lock")
                 input_chunk_caches[i].ready_lock = PyThread_allocate_lock()
                 if input_chunk_caches[i].ready_lock == NULL:
@@ -4205,8 +4216,7 @@ cdef class NDArray:
                         if input_chunk_caches[i].ready_lock != NULL:
                             PyThread_free_lock(input_chunk_caches[i].ready_lock)
                     free(input_chunk_caches)
-                    free(inputs_)
-                    free(udata)
+                    _free_me_udata_tables(udata, inputs_, np_data, np_typesizes)
                     raise MemoryError("Cannot allocate miniexpr chunk cache ready lock")
         udata.input_chunk_caches = input_chunk_caches
         eval_params = <me_eval_params*> malloc(sizeof(me_eval_params))
@@ -4217,8 +4227,7 @@ cdef class NDArray:
                 if input_chunk_caches[i].ready_lock != NULL:
                     PyThread_free_lock(input_chunk_caches[i].ready_lock)
             free(input_chunk_caches)
-            free(inputs_)
-            free(udata)
+            _free_me_udata_tables(udata, inputs_, np_data, np_typesizes)
             raise MemoryError("Cannot allocate miniexpr eval params")
         eval_params.disable_simd = False
         eval_params.simd_ulp_mode = ME_SIMD_ULP_3_5 if fp_accuracy == blosc2.FPAccuracy.MEDIUM else ME_SIMD_ULP_1

src/blosc2/proxy.py

@@ -1073,6 +1073,9 @@ def dsl_wrapper(*args, **func_kwargs):
             return out
 
         if storage_kwargs and any(v is not None for v in storage_kwargs.values()):
+            # Execution-tuning kwargs go along too: compute() names all three,
+            # while lazyudf() above only names jit/jit_backend, so fp_accuracy
+            # would otherwise be dropped on this path.
             return lexpr.compute(**decorator_kwargs)
         return lexpr[()]
 
@@ -1270,8 +1273,11 @@ def wrapper(*args, **func_kwargs):
             # If it is a numpy array, return it as is
             if isinstance(retval, np.ndarray):
                 if storage_kwargs and any(v is not None for v in storage_kwargs.values()):
-                    # But if storage kwargs are provided, return a NDArray instead
-                    return blosc2.asarray(retval, **kwargs)
+                    # But if storage kwargs are provided, return a NDArray instead.
+                    # Only storage kwargs: asarray() rejects the execution-tuning
+                    # ones, and there is nothing left to tune -- the function has
+                    # already run.
+                    return blosc2.asarray(retval, **storage_kwargs)
                 return retval
 
             # In some instances, the return value is not a LazyExpr

tests/ndarray/test_jit.py

@@ -205,3 +205,19 @@ def f(a, b):
     assert isinstance(res, blosc2.NDArray)
     assert res.schunk.cparams.clevel == 2
     np.testing.assert_allclose(res[:], a * 2.0 + b)
+
+
+def test_jit_numpy_return_with_storage_and_tuning_kwargs():
+    # A traced function whose return is already a NumPy array takes the
+    # asarray() branch, which accepts storage kwargs only -- forwarding the
+    # execution-tuning ones there raised instead of returning an NDArray.
+    @blosc2.jit(jit=False, cparams=blosc2.CParams(clevel=2))
+    def f(a, b):
+        return np.sum(a * 2.0 + b, axis=0)
+
+    a = np.arange(1000, dtype=np.float64).reshape(10, 100)
+    b = np.arange(1000, dtype=np.float64).reshape(10, 100) * 0.5
+    res = f(a, b)
+    assert isinstance(res, blosc2.NDArray)
+    assert res.schunk.cparams.clevel == 2
+    np.testing.assert_allclose(res[:], np.sum(a * 2.0 + b, axis=0))