Merge pull request #2702 from mabel-dev/cythonize-LENGTH-function

Cythonize LENGTH function

Author: joocer

opteryx/__version__.py

@@ -1,4 +1,4 @@
-__build__ = 1394
+__build__ = 1396
 
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

opteryx/compiled/list_ops/list_length.pyx

@@ -0,0 +1,53 @@
+# cython: language_level=3
+# cython: nonecheck=False
+# cython: cdivision=True
+# cython: initializedcheck=False
+# cython: infer_types=True
+# cython: wraparound=False
+# cython: boundscheck=False
+
+import numpy
+cimport numpy
+numpy.import_array()
+
+import pyarrow
+
+from libc.stdint cimport int32_t, uint32_t
+
+cpdef numpy.ndarray[numpy.uint32_t, ndim=1] list_length(object array):
+
+    cdef Py_ssize_t n
+    cdef numpy.ndarray[numpy.uint32_t, ndim=1] result
+    cdef uint32_t[::1] result_view
+    cdef object val
+    cdef uint32_t i
+    cdef numpy.ndarray[numpy.int32_t, ndim=1] offsets
+    cdef int32_t[::1] offsets_view
+
+    # PyArrow fast path (uses offsets buffer)
+    if isinstance(array, (pyarrow.Array, pyarrow.ChunkedArray, pyarrow.lib.StringArray)):
+        if isinstance(array, pyarrow.ChunkedArray):
+            array = array.combine_chunks()
+
+        n = len(array)
+
+        try:
+            offsets_buffer = array.buffers()[1]
+            offsets = numpy.frombuffer(offsets_buffer, dtype=numpy.int32, count=n + 1)
+            offsets_view = offsets
+            return (offsets[1:] - offsets[:-1]).astype(numpy.uint32)
+        except Exception:
+            pass  # fallback if offsets unavailable
+
+    n = array.shape[0]
+    result = numpy.empty(n, dtype=numpy.uint32)
+    result_view = result
+
+    for i in range(n):
+        val = array[i]
+        if isinstance(val, (str, bytes, list, numpy.ndarray)):
+            result_view[i] = len(val)
+        else:
+            result_view[i] = 0
+
+    return result

opteryx/functions/__init__.py

@@ -22,6 +22,7 @@
 import opteryx
 from opteryx.compiled.list_ops.list_contains_any import list_contains_any
 from opteryx.compiled.list_ops.list_encode_utf8 import list_encode_utf8 as to_blob
+from opteryx.compiled.list_ops.list_length import list_length
 from opteryx.exceptions import FunctionExecutionError
 from opteryx.exceptions import IncorrectTypeError
 from opteryx.functions import date_functions
@@ -260,19 +261,6 @@ def _inner(array, literal):
     return _inner
 
 
-def get_len(obj):
-    """len, but nullsafe"""
-    if hasattr(obj, "__len__"):
-        return len(obj)
-    if hasattr(obj, "length"):  # Some Arrow scalar types have .length property
-        return obj.length
-    if hasattr(obj, "nbytes"):  # NumPy scalar types have .nbytes
-        return obj.nbytes
-    if hasattr(obj, "as_py") and isinstance(obj.as_py(), (bytes, str)):  # PyArrow string scalar
-        return len(obj.as_py())
-    return None
-
-
 def _coalesce(*arrays):
     """
     Element-wise coalesce function for multiple numpy arrays.
@@ -394,7 +382,7 @@ def sleep(x):
     "ASCII": (string_functions.to_ascii, "INTEGER", 1.0),
 
     # STRINGS
-    "LENGTH": (_iterate_single_parameter(get_len), "INTEGER", 1.0),  # LENGTH(str) -> int
+    "LENGTH": (list_length, "INTEGER", 1.0),  # LENGTH(str) -> int
     "UPPER": (compute.utf8_upper, "VARCHAR", 1.0),  # UPPER(str) -> str
     "LOWER": (compute.utf8_lower, "VARCHAR", 1.0),  # LOWER(str) -> str
     "LEFT": (string_functions.string_slicer_left, "VARCHAR", 1.0),

opteryx/managers/expression/ops.py

@@ -5,6 +5,7 @@
 """
 
 import re
+from contextlib import suppress
 
 import numpy
 import pyarrow
@@ -147,7 +148,7 @@ def _inner_filter_operations(arr, operator, value):
     """
     if not operator.startswith(("AnyOp", "AllOp")) and len(value) == 1:
         value = value[0]
-        if hasattr(value, "item"):
+        with suppress(AttributeError):
             value = value.item()
         if isinstance(value, (tuple, list)):
             value = pyarrow.array(value)
@@ -165,24 +166,24 @@ def _inner_filter_operations(arr, operator, value):
     if operator == "GtEq":
         return compute.greater_equal(arr, value).to_numpy(False).astype(dtype=bool)
     if operator == "InList":
-        if hasattr(value, "to_pylist"):
+        with suppress(AttributeError):
             value = value.to_pylist()
-        if hasattr(value, "to_numpy"):
+        with suppress(AttributeError):
             value = value.to_numpy(zero_copy_only=False)
         values = set(value)
-        if hasattr(arr, "to_numpy"):
+        with suppress(AttributeError):
             arr = arr.to_numpy(zero_copy_only=False)
         if arr.dtype == numpy.int64:
             return list_ops.list_in_list.list_in_list_int64(memoryview(arr), values, len(arr))
         else:
             return list_ops.list_in_list.list_in_list(arr.astype(object), values)
     if operator == "NotInList":
-        if hasattr(value, "to_pylist"):
+        with suppress(AttributeError):
             value = value.to_pylist()
-        if hasattr(value, "to_numpy"):
+        with suppress(AttributeError):
             value = value.to_numpy(zero_copy_only=False)
         values = set(value)
-        if hasattr(arr, "to_numpy"):
+        with suppress(AttributeError):
             arr = arr.to_numpy(zero_copy_only=False)
         if arr.dtype == numpy.int64:
             matches = list_ops.list_in_list.list_in_list_int64(memoryview(arr), values, len(arr))