numbarrow 0.0.0a0__py3-none-any.whl

numbarrow/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.0.0a0'

numbarrow/core/adapters.py

@@ -0,0 +1,108 @@
+import numpy as np
+import pyarrow as pa
+
+from functools import singledispatch
+from typing import Union
+
+from numbarrow.core.is_null import is_null
+from numbarrow.utils.arrow_array_utils import (
+    create_bitmap, create_str_array, structured_array_adapter,
+    structured_list_array_adapter, uniform_arrow_array_adapter
+)
+from numbarrow.utils.utils import arrays_viewers
+
+
+def cast_64bit_date_arrow_to_numpy_array(pa_array: pa.Array, np_dtype: np.dtype):
+    """ Can be used to cast PyArrow arrays of date types that are represented by
+    64-bit integers to numpy arrays of various date types (np.datetime64[...],
+    which are always represented  by 64-bit integers whose meaning is determined
+    by the precision, such as, 's', 'ms', 'us').
+
+    Since underlying data layout of both arrays in int64, a copy is avoided,
+
+    The associated bitmap (if any) is also returned.
+    """
+    int64_array = pa_array.cast(pa.int64())
+    assert int64_array.buffers()[1].address == pa_array.buffers()[1].address, "got copied"
+    bitmap, int64_data = uniform_arrow_array_adapter(int64_array)
+    data = int64_data.view(np_dtype)
+    assert data.ctypes.data == int64_data.ctypes.data, "got copied"
+    return bitmap, data
+
+
+@singledispatch
+def arrow_array_adapter(pa_array: pa.Array):
+    """ Dispatcher for PyArrow array adapters of various types. """
+    raise NotImplementedError(f"Not implemented for {pa_array} of type {type(pa_array)} and elements {pa_array.type}")
+
+
+@arrow_array_adapter.register(pa.BooleanArray)
+def _(pa_array: pa.BooleanArray):
+    """ PyArrow stores boolean arrays bit-wise,
+     following the same kind of layout it uses
+     for bitmaps. This requires creating a copy
+     when casting to numpy arrays of booleans.
+    """
+    bitmap_buf, data_buf = pa_array.buffers()
+    data_buf_p = data_buf.address
+    num_of_bool_elements = len(pa_array)
+    num_of_bytes = num_of_bool_elements // 8
+    if num_of_bool_elements % 8 > 0:
+        num_of_bytes += 1
+    packed_boolean_data_viewer = arrays_viewers[np.uint8]
+    packed_boolean_data = packed_boolean_data_viewer(data_buf_p, num_of_bytes)
+    data_lst = [not is_null(i, packed_boolean_data) for i in range(num_of_bool_elements)]
+    data = np.array(data_lst, dtype=np.bool_)
+    bitmap = create_bitmap(bitmap_buf)
+    return bitmap, data
+
+
+@arrow_array_adapter.register(pa.Date32Array)
+def _(pa_array: pa.Date32Array):
+    """
+    PyArrow Date32 dates are represented by 32bit integers.
+    Since all numpy dates are represented by 64bit integers, this
+    creates a copy when it re-interprets numpy array of int32 integers
+     (number of days since 1970-01-01) as datetime64[D] (int64)"""
+    int32_array = pa_array.cast(pa.int32())
+    assert int32_array.buffers()[1].address == pa_array.buffers()[1].address, "got copied"
+    bitmap, int32_data = uniform_arrow_array_adapter(int32_array)
+    data = int32_data.astype(np.dtype("datetime64[D]"))
+    assert int32_data.ctypes.data != data.ctypes.data
+    return bitmap, data
+
+
+@arrow_array_adapter.register(pa.Date64Array)
+def _(pa_array: pa.Date64Array):
+    return cast_64bit_date_arrow_to_numpy_array(pa_array, np.dtype("datetime64[ms]"))
+
+
+@arrow_array_adapter.register(pa.lib.DoubleArray)
+@arrow_array_adapter.register(pa.Int32Array)
+@arrow_array_adapter.register(pa.Int64Array)
+def _(pa_array: Union[
+    pa.lib.DoubleArray, pa.Int32Array, pa.Int64Array
+]):
+    return uniform_arrow_array_adapter(pa_array)
+
+
+@arrow_array_adapter.register(pa.ListArray)
+def _(pa_array: pa.ListArray):
+    return structured_list_array_adapter(pa_array)
+
+
+@arrow_array_adapter.register(pa.StructArray)
+def _(pa_array: pa.StructArray):
+    return structured_array_adapter(pa_array)
+
+
+@arrow_array_adapter.register(pa.StringArray)
+def _(pa_array: pa.StringArray):
+    return None, create_str_array(pa_array)
+
+
+@arrow_array_adapter.register(pa.TimestampArray)
+def _(pa_array: pa.TimestampArray):
+    timestamp_type: pa.TimestampType = pa_array.type
+    timestamp_unit = timestamp_type.unit
+    return cast_64bit_date_arrow_to_numpy_array(pa_array, np.dtype(f"datetime64[{timestamp_unit}]"))

numbarrow/core/configurations.py

@@ -0,0 +1,3 @@
+default_jit_options = {
+    "cache": True
+}

numbarrow/core/is_null.py

@@ -0,0 +1,12 @@
+import numpy as np
+from numba import njit
+from numba.core.types import boolean, int64, Array, uint8
+
+from numbarrow.core.configurations import default_jit_options
+
+
+@njit(boolean(int64, Array(uint8, 1, "C")), **default_jit_options)
+def is_null(index_: int, bitmap: np.ndarray) -> bool:
+    byte_for_index = bitmap[index_ // 8]
+    bit_position_in_byte = index_ % 8
+    return not (byte_for_index >> bit_position_in_byte) % 2

numbarrow/core/mapinarrow_factory.py

@@ -0,0 +1,44 @@
+import numpy as np
+import pyarrow as pa
+
+from pyspark.broadcast import Broadcast
+from typing import Callable, Dict, List, Optional
+
+from numbarrow.core.adapters import arrow_array_adapter
+
+
+def make_mapinarrow_func(
+    main_func: Callable,
+    input_columns: Optional[List[str]] = None,
+    broadcasts: Optional[Dict] = None
+):
+    """
+    Creates a function that can be given as an argument to `mapInArrow`
+
+    :param main_func: should have the following signature:
+        - `data_dict: Dict[str, np.ndarray]`, values are arrays of data of various supported types
+        - `bitmap_dict: Dict[str, np.ndarray]`, values are uint8 aligned arrays of bitmap data
+        - `broadcasts` Optional[Dict[str, Any]]
+        returns: Dict[str, np.ndarray] that will be used to create PyArrow `RecordBatch`
+    :param input_columns: optional list of column names that will be expected to be needed for in
+    `data_dict` for the calculation done by `main_func`. When not given, all columns in the iterated
+     over PySpark DataFrame will be used.
+    :param broadcasts: optional dictionary of broadcast values
+    """
+    broadcasts = broadcasts if broadcasts is not None else {}
+
+    def _(iterator):
+        for batch in iterator:
+            data_dict: Dict[str, np.ndarray] = {}
+            bitmap_dict: Dict[str, np.ndarray] = {}
+            input_columns_ = input_columns if input_columns is not None else batch.schema.names
+            for col in input_columns_:
+                col_pa: pa.Array = batch.column(col)
+                col_bitmap, col_data = arrow_array_adapter(col_pa)
+                col_bitmap = col_bitmap if isinstance(col_bitmap, dict) else {} if col_bitmap is None else {col: col_bitmap}  # noqa: E501
+                col_data = col_data if isinstance(col_data, dict) else {col: col_data}
+                bitmap_dict = {**bitmap_dict, **col_bitmap}
+                data_dict = {**data_dict, **col_data}
+            outputs = main_func(data_dict, bitmap_dict, broadcasts)
+            yield pa.RecordBatch.from_pydict({col: pa.array(output) for col, output in outputs.items()})
+    return _

numbarrow/utils/arrow_array_utils.py

@@ -0,0 +1,102 @@
+import ctypes
+import numpy as np
+import pyarrow as pa
+
+from typing import Dict, Optional, Tuple
+
+from numbarrow.utils.utils import arrays_viewers
+
+
+def create_bitmap(bitmap_buf: Optional[pa.Buffer]):
+    """ Create numpy array of uint8 type containing
+    bit-map of valid array entries """
+    if bitmap_buf is None:
+        return None
+    bitmap_p = bitmap_buf.address
+    bitmap_len = bitmap_buf.size
+    bitmap_viewer = arrays_viewers[np.uint8]
+    bitmap = bitmap_viewer(bitmap_p, bitmap_len)
+    return bitmap
+
+
+def create_str_array(pa_str_array: pa.StringArray) -> np.ndarray:
+    """ Copy data from densely packed `pa.StringArray` into
+     padded numpy array of the character sequence type determined
+     by the length of the longest string. """
+    bitmap_buf, offsets_buf, data_buf = pa_str_array.buffers()
+    data_p = data_buf.address
+    offsets_p = offsets_buf.address
+    offsets_len = offsets_buf.size // np.dtype(np.int32).itemsize
+    offsets_array = arrays_viewers[np.int32](offsets_p, offsets_len)
+    diffs = np.diff(offsets_array)
+    str_array_len = len(diffs)
+    item_sz = diffs.max()
+    str_array = np.empty((str_array_len,), dtype=f"|U{item_sz}")
+    for i in range(len(offsets_array) - 1):
+        start = offsets_array[i]
+        end = offsets_array[i + 1]
+        s = (ctypes.c_char * int(end - start)).from_address(data_p + int(start)).value
+        str_array[i] = s
+    return str_array
+
+
+def structured_array_adapter(struct_array: pa.StructArray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+    """
+    NumPy adapter of PyArrow `StructArray`.
+
+    Returns tuple of two dictionaries, the first dictionary maps names of
+    the structure fields to the contiguous bitmap arrays, the second maps
+    these names to the contiguous value arrays.
+    """
+    assert isinstance(struct_array, pa.StructArray)
+    data_type: pa.StructType = struct_array.type
+    assert isinstance(data_type, pa.StructType)
+    bitmaps = {}
+    datas = {}
+    for field_ind in range(len(data_type)):
+        field: pa.Field = data_type[field_ind]
+        field_name = field.name
+        pa_array = struct_array.field(field_name)
+        bitmap, data = uniform_arrow_array_adapter(pa_array)
+        bitmaps[field_name] = bitmap
+        datas[field_name] = data
+    return bitmaps, datas
+
+
+def structured_list_array_adapter(list_array: pa.ListArray) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]:
+    """
+    NumPy adapter of PyArrow array of same-length lists of structures.
+
+    :param list_array: PyArrow array with elements being of `pa.ListType`.
+    Each list is in turn of the same length, and each element of the list
+    is of `pa.StructType`.
+
+    Returns tuple of two dictionaries, the first dictionary maps names of
+    the structure fields to the contiguous bitmap array, the second maps
+    these names to the contiguous values arrays.
+
+    Data is not copied as it is uniformly stored in a columnar format,
+    that is, the underlying values are stored contiguously in a
+    `pa.StructArray`.
+    """
+    assert isinstance(list_array, pa.ListArray)
+    data_values: pa.StructArray = list_array.values
+    return structured_array_adapter(data_values)
+
+
+def uniform_arrow_array_adapter(pa_array: pa.Array) -> Tuple[Optional[np.ndarray], np.ndarray]:
+    """ NumPy adapter for PyArrow arrays with uniformly sized elements.
+     Returns views over bitmap and data contiguous memory regions as numpy arrays. """
+    bitmap_buf, data_buf = pa_array.buffers()
+    data_arrow_ty = pa_array.type
+    data_np_ty = data_arrow_ty.to_pandas_dtype()
+    data_viewer = arrays_viewers.get(data_np_ty, None)
+    if data_viewer is None:
+        raise ValueError(f"There is no {data_np_ty} in `utils.arrays_viewers`. Add it?")
+    data_p = data_buf.address
+    data_buf_byte_size = data_buf.size
+    data_item_byte_size = np.dtype(data_np_ty).itemsize
+    data_len = data_buf_byte_size // data_item_byte_size
+    data = data_viewer(data_p, data_len)
+    bitmap = create_bitmap(bitmap_buf)
+    return bitmap, data

numbarrow/utils/utils.py

@@ -0,0 +1,31 @@
+import numpy as np
+from numba import carray, from_dtype, int64, intp, njit
+from numba.core.types import Array, voidptr
+from numba.extending import intrinsic
+
+from numbarrow.core.configurations import default_jit_options
+
+
+@intrinsic
+def _ptr_as_int_to_voidptr(typingctx, arg_type):
+    def codegen(context, builder, signature, args):
+        return builder.inttoptr(args[0], context.get_value_type(voidptr))
+    return voidptr(arg_type), codegen
+
+
+def numpy_array_from_ptr_factory(dtype_):
+    @njit(Array(from_dtype(dtype_), 1, "C")(intp, int64), **default_jit_options)
+    def _(ptr_as_int: int, sz: int):
+        return carray(_ptr_as_int_to_voidptr(ptr_as_int), shape=(sz,), dtype=dtype_)
+    return _
+
+
+arrays_viewers = {
+    np_type: numpy_array_from_ptr_factory(np_type) for np_type in [
+        np.bool_,
+        np.float64,
+        np.int32,
+        np.int64,
+        np.uint8
+    ]
+}

numbarrow-0.0.0a0.dist-info/LICENSE

@@ -0,0 +1,27 @@
+MIT License (with Citation Clause)
+
+Copyright (c) 2025 Mikhail Goykhman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+- The above copyright notice and this permission notice shall be included in all
+  copies or substantial portions of the Software.
+
+- If this software or any derivative work is used, modified, or distributed,
+  you must provide proper credit to the original author. This includes:
+  - Mentioning the original author in documentation, README files, or any
+    publication describing work based on this software.
+  - Including a link to the original repository, when applicable.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

numbarrow-0.0.0a0.dist-info/METADATA

@@ -0,0 +1,49 @@
+Metadata-Version: 2.1
+Name: numbarrow
+Version: 0.0.0a0
+Author: Mikhail Goykhman
+License: MIT License (with Citation Clause)
+        
+        Copyright (c) 2025 Mikhail Goykhman
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        - The above copyright notice and this permission notice shall be included in all
+          copies or substantial portions of the Software.
+        
+        - If this software or any derivative work is used, modified, or distributed,
+          you must provide proper credit to the original author. This includes:
+          - Mentioning the original author in documentation, README files, or any
+            publication describing work based on this software.
+          - Including a link to the original repository, when applicable.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Keywords: pyarrow,pyspark,numba,numpy
+Requires-Python: ~=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: llvmlite<0.45.0,>=0.43.0
+Requires-Dist: numba<0.62.0,>=0.60.0
+Requires-Dist: numpy<2.0.0,>=1.22.4
+Requires-Dist: pandas==1.5.0
+Requires-Dist: py4j==0.10.9.5
+Requires-Dist: pyarrow<=15.0.0
+Requires-Dist: pyspark~=3.3.1
+Requires-Dist: python-dateutil==2.9.0.post0
+Requires-Dist: pytz==2025.2
+Requires-Dist: six==1.17.0
+
+# numbarrow
+Numba adapters for PyArrow and PySpark

numbarrow-0.0.0a0.dist-info/RECORD

@@ -0,0 +1,14 @@
+numbarrow/__init__.py,sha256=ZOPI6V3g5yKt5gCJVpzpQrEJVh0AfpD9NTTDyJ1J2CU,24
+numbarrow/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+numbarrow/core/adapters.py,sha256=nY5_upQUrNOq-M_4glSjkVp5PNc833-Xckqc8jlWW9Y,4100
+numbarrow/core/configurations.py,sha256=dQe7TN0jJ4HUGYCxYBAMdyMm8dC5rhRVZ1PHRfWD1OE,44
+numbarrow/core/is_null.py,sha256=LcPky2pRIKMtlu2WMteU9bMsKEzZv5Q1HxFGje45GNQ,425
+numbarrow/core/mapinarrow_factory.py,sha256=JQIEnLLlNY50jxerFc6VNeWF2RWkVk52rVn_0NV64cQ,2110
+numbarrow/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+numbarrow/utils/arrow_array_utils.py,sha256=I8yW4wOPlVXEA7lc4_aHXuAfZ-NYz_EnEAsAKdfdFZw,4038
+numbarrow/utils/utils.py,sha256=jRPZjFAwr1iSm9TtYtO4ioF_gDJzYCZ1eb0BEkwMzWg,894
+numbarrow-0.0.0a0.dist-info/LICENSE,sha256=YYgNvjH_p6-1NsdrIqGJnr1GUbZzA_8DxsP6vVfM6nY,1446
+numbarrow-0.0.0a0.dist-info/METADATA,sha256=6DASatw_x9Gq5qm30ph2W9JjVF7osqiOIAxp0tR2q_s,2260
+numbarrow-0.0.0a0.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+numbarrow-0.0.0a0.dist-info/top_level.txt,sha256=MOoPFxa928DKw4fAl7Q9KnV7bNx8LKN3cbzHSpauZSM,10
+numbarrow-0.0.0a0.dist-info/RECORD,,

numbarrow-0.0.0a0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.3.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

numbarrow-0.0.0a0.dist-info/top_level.txt

@@ -0,0 +1 @@
+numbarrow