Merge branch 'main' into copilot/use-stats-to-create-filters

Author: joocer

.gitignore

@@ -219,3 +219,4 @@ tests/draken/performance/.perf_baseline.json
 opteryx/draken/vectors/_hash_api.c*
 third_party/mabel/draken/vectors/interval_vector.c*
 third_party/mabel/draken/morsels/align.c*
+third_party/mabel/draken/compiled/maskops.cpp

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1907
+__build__ = 1926
 __author__ = "@joocer"
-__version__ = "0.26.2-beta.1907"
+__version__ = "0.26.2-beta.1926"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/compiled/aggregations/group_by_draken.pyx

@@ -51,6 +51,8 @@ cpdef Morsel group_by_morsel(
         object iter_indices
         int64_t group_val_int
         double group_val_double
+        uint64_t h
+        set seen_set
 
     if morsel is None or morsel.num_rows == 0:
         # return empty table similar to what pyarrow would return
@@ -85,7 +87,6 @@ cpdef Morsel group_by_morsel(
 
         # Build map of hash -> vector of row indices
         # Store seen keys in seen_keys (Python list) to iterate later
-        cdef uint64_t h
         for i in range(num_rows):
             h = <uint64_t> row_hashes[i]
             fmap.insert(h, i)
@@ -98,7 +99,7 @@ cpdef Morsel group_by_morsel(
 
         # Build unique keys dict by iterating again and filling a Python dict
         seen_keys = []
-        cdef set seen_set = set()
+        seen_set = set()
         for i in range(num_rows):
             h = <uint64_t> row_hashes[i]
             if h not in seen_set:

opteryx/compiled/functions/strings.pyx

@@ -7,6 +7,7 @@
 # cython: boundscheck=False
 
 import platform
+from libcpp.vector cimport vector
 
 cdef extern from "simd_search.h":
     size_t neon_count(const char* data, size_t length, char target)
@@ -15,21 +16,30 @@ cdef extern from "simd_search.h":
     int avx_search(const char* data, size_t length, char target)
     int neon_find_delimiter(const char* data, size_t length)
     int avx_find_delimiter(const char* data, size_t length)
+    vector[size_t] neon_find_all(const char* data, size_t length, char target)
+    vector[size_t] avx_find_all(const char* data, size_t length, char target)
+
+cdef extern from "simd_string_ops.h":
+    void simd_to_upper(char* data, size_t length)
+    void simd_to_lower(char* data, size_t length)
 
 # Architecture detection
 cdef size_t (*simd_count)(const char*, size_t, char)
 cdef int (*simd_search)(const char*, size_t, char)
 cdef int (*simd_find_delimiter)(const char*, size_t)
+cdef vector[size_t] (*simd_find_all)(const char*, size_t, char)
 
 _arch = platform.machine().lower()
 if _arch in ('arm64', 'aarch64'):
     simd_count = neon_count
     simd_search = neon_search
     simd_find_delimiter = neon_find_delimiter
+    simd_find_all = neon_find_all
 else:
     simd_count = avx_count
     simd_search = avx_search
     simd_find_delimiter = avx_find_delimiter
+    simd_find_all = avx_find_all
 
 
 # Expose simd_count to Python so higher-level code can reuse the optimized newline/count routine.
@@ -43,3 +53,70 @@ cpdef size_t count_instances(const unsigned char[::1] mv, char target=10):
     cdef const char* data = <const char*>udata
     cdef size_t data_len = mv.shape[0]
     return simd_count(data, data_len, target)
+
+
+# String case conversion functions
+def to_upper(bytes data):
+    """
+    Convert ASCII characters in bytes to uppercase using SIMD.
+    Non-ASCII bytes are left unchanged.
+
+    Returns a new bytes object with uppercase characters.
+    """
+    cdef size_t length = len(data)
+    cdef bytearray result = bytearray(data)
+    cdef char* ptr = <char*> result
+    simd_to_upper(ptr, length)
+    return bytes(result)
+
+
+def to_lower(bytes data):
+    """
+    Convert ASCII characters in bytes to lowercase using SIMD.
+    Non-ASCII bytes are left unchanged.
+
+    Returns a new bytes object with lowercase characters.
+    """
+    cdef size_t length = len(data)
+    cdef bytearray result = bytearray(data)
+    cdef char* ptr = <char*> result
+    simd_to_lower(ptr, length)
+    return bytes(result)
+
+
+# Character search functions
+def find_char(bytes data, int target):
+    """
+    Find the first occurrence of a character in bytes using SIMD.
+
+    Returns the index of the first occurrence, or -1 if not found.
+    """
+    cdef const char* ptr = <const char*> (<char*> data)
+    cdef size_t length = len(data)
+    cdef char char_target = <char> target
+    return simd_search(ptr, length, char_target)
+
+
+def count_char(bytes data, int target):
+    """
+    Count occurrences of a character in bytes using SIMD.
+
+    Returns the count of occurrences.
+    """
+    cdef const char* ptr = <const char*> (<char*> data)
+    cdef size_t length = len(data)
+    cdef char char_target = <char> target
+    return simd_count(ptr, length, char_target)
+
+
+def find_all_char(bytes data, int target):
+    """
+    Find all occurrences of a character in bytes using SIMD.
+
+    Returns a list of indices where the character occurs.
+    """
+    cdef const char* ptr = <const char*> (<char*> data)
+    cdef size_t length = len(data)
+    cdef char char_target = <char> target
+    cdef vector[size_t] positions = simd_find_all(ptr, length, char_target)
+    return list(positions)

opteryx/compiled/list_ops/list_initcap.pyx

@@ -22,8 +22,12 @@ cdef inline numpy.ndarray[object, ndim=1] _ensure_object_array(object data):
             return data
         return data.astype(object)
 
-    if hasattr(data, "combine_chunks"):
-        data = data.combine_chunks()
+    if hasattr(data, "chunks"):
+        # Handle ChunkedArray by processing chunks individually to avoid massive copy
+        return numpy.concatenate([_ensure_object_array(chunk) for chunk in data.chunks])
+
+    # if hasattr(data, "combine_chunks"):
+    #     data = data.combine_chunks()
 
     if hasattr(data, "to_numpy"):
         arr = data.to_numpy(zero_copy_only=False)

opteryx/compiled/list_ops/list_length.pyx

@@ -26,7 +26,7 @@ cpdef numpy.ndarray[numpy.uint32_t, ndim=1] list_length(object array):
     # 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()
+            return numpy.concatenate([list_length(chunk) for chunk in array.chunks])
 
         n = len(array)
 

opteryx/compiled/structures/buffers.pxd

@@ -9,7 +9,7 @@
 import numpy
 cimport numpy
 
-from libc.stdint cimport int64_t
+from libc.stdint cimport int64_t, int32_t
 from libc.stddef cimport size_t
 from libcpp.vector cimport vector
 
@@ -20,14 +20,29 @@ cdef extern from "intbuffer.h" namespace "" nogil:
         void extend(const vector[int64_t]& values)
         void extend(const int64_t* values, size_t count)
         void reserve(size_t additional_capacity)
+        void resize(size_t new_size)
         const int64_t* data() const
+        int64_t* mutable_data()
         size_t size() const
         void append_repeated(int64_t value, size_t count)
 
+    cdef cppclass CInt32Buffer:
+        CInt32Buffer(size_t size_hint) except +
+        void append(int32_t value)
+        void extend(const vector[int32_t]& values)
+        void extend(const int32_t* values, size_t count)
+        void reserve(size_t additional_capacity)
+        void resize(size_t new_size)
+        const int32_t* data() const
+        int32_t* mutable_data()
+        size_t size() const
+
 
 cdef class IntBuffer:
 
     cdef CIntBuffer* c_buffer
+    cdef Py_ssize_t _shape[1]
+    cdef Py_ssize_t _strides[1]
 
     cpdef void append(self, int64_t value)
     cpdef void extend(self, iterable)
@@ -37,3 +52,16 @@ cdef class IntBuffer:
     cpdef void extend_numpy(self, numpy.ndarray[int64_t, ndim=1] arr)
     cpdef void reserve(self, size_t capacity)
     cpdef void append_batch(self, int64_t[::1] values)
+
+cdef class Int32Buffer:
+
+    cdef CInt32Buffer* c_buffer
+    cdef Py_ssize_t _shape[1]
+    cdef Py_ssize_t _strides[1]
+
+    cpdef void append(self, int32_t value)
+    cpdef void extend(self, iterable)
+    cpdef numpy.ndarray[int32_t, ndim=1] to_numpy(self)
+    cpdef size_t size(self)
+    cpdef void reserve(self, size_t capacity)
+    cpdef void extend_numpy(self, numpy.ndarray[int32_t, ndim=1] arr)