Fast path for UDFs and new benchmarks

Author: FrancescAlted

bench/ndarray/array-constructor-memray.py

@@ -0,0 +1,58 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+from time import time
+import os
+
+import numpy as np
+import memray
+
+import blosc2
+
+N = 100_000_000
+
+
+def info(a, t1):
+    size = a.schunk.nbytes
+    csize = a.schunk.cbytes
+    print(
+        f"Time: {t1:.3f} s - size: {size / 2 ** 30:.2f} GB ({size / t1 / 2 ** 30:.2f} GB/s)"
+        f"\tStorage required: {csize / 2 ** 20:.2f} MB (cratio: {size / csize:.1f}x)"
+    )
+
+
+def run_benchmark():
+    shape = (N,)
+    shape = (100, 1000, 1000)
+    print(f"*** Creating a blosc2 array with {N:_} elements (shape: {shape}) ***")
+    t0 = time()
+    #a = blosc2.arange(N, shape=shape, dtype=np.int32, urlpath="a.b2nd", mode="w")
+    a = blosc2.linspace(0, 1, N, shape=shape, dtype=np.float64, urlpath="a.b2nd", mode="w")
+    elapsed = time() - t0
+    info(a, elapsed)
+    return a
+
+
+# Check if we're being tracked by memray
+if not os.environ.get("MEMRAY_TRACKING", False):
+    # Run the benchmark with memray tracking
+    output_file = "array_constructor_memray.bin"
+    print(f"Starting memray profiling. Results will be saved to {output_file}")
+
+    with memray.Tracker(output_file):
+        array = run_benchmark()
+
+    print(f"\nMemray profiling completed. To view results, run:")
+    print(f"memray flamegraph {output_file}")
+    print(f"# or")
+    print(f"memray summary {output_file}")
+    print(f"# or")
+    print(f"memray tree {output_file}")
+else:
+    # We're already being tracked by memray
+    run_benchmark()

bench/ndarray/array-constructor.py

@@ -0,0 +1,32 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# This source code is licensed under a BSD-style license (found in the
+# LICENSE file in the root directory of this source tree)
+#######################################################################
+
+from time import time
+
+import numpy as np
+
+import blosc2
+
+N = 100_000_000
+
+def info(a, t1):
+    size = a.schunk.nbytes
+    csize = a.schunk.cbytes
+    print(
+        f"Time: {t1:.3f} s - size: {size / 2 ** 30:.2f} GB ({size / t1 / 2 ** 30:.2f} GB/s)"
+        f"\tStorage required: {csize / 2 ** 20:.2f} MB (cratio: {size / csize:.1f}x)"
+    )
+
+
+shape = (N,)
+shape = (100, 1000, 1000)
+print(f"*** Creating a blosc2 array with {N:_} elements (shape: {shape}) ***")
+t0 = time()
+# a = blosc2.arange(N, shape=shape, dtype=np.int32, urlpath="a.b2nd", mode="w")
+a = blosc2.linspace(0, 1, N, shape=shape, dtype=np.float64, urlpath="a.b2nd", mode="w")
+info(a, time() - t0)

src/blosc2/ndarray.py

@@ -334,7 +334,14 @@ def reshape(
             slice_size = src_slice.stop - src_slice.start
             dst_buf_slice = slice(dst_buf_len, dst_buf_len + slice_size)
             dst_buf_len += slice_size
-            dst_buf[dst_buf_slice] = src[src_slice]
+            if hasattr(src, "res_getitem"):
+                # Fast path for lazy UDFs (important for e.g. arange or linspace)
+                # This essentially avoids the need to create a new,
+                # potentially large NumPy array in memory.
+                # This is not critical for Linux, but it is for Windows/Mac.
+                dst_buf[dst_buf_slice] = src.res_getitem[src_slice]
+            else:
+                dst_buf[dst_buf_slice] = src[src_slice]
         # Compute the shape of dst_slice
         dst_slice_shape = tuple(s.stop - s.start for s in dst_slice)
         # ... and assign the buffer to the destination array