Skip to content

Commit

Permalink
Zero-copy buffer protocol data import (#204)
Browse files Browse the repository at this point in the history
* Zero-copy buffer protocol

* Update README

* Automatically convert buffer protocol to PyArray

* reorder doc

* Fix type hints to allow buffer protocol objects

* Add test for operations

* Update reamde

* better docstring

* Test that datetime array fails
  • Loading branch information
kylebarron authored Oct 1, 2024
1 parent e78cd54 commit 9a01c80
Show file tree
Hide file tree
Showing 17 changed files with 515 additions and 364 deletions.
8 changes: 4 additions & 4 deletions arro3-compute/python/arro3/compute/_aggregate.pyi
Original file line number Diff line number Diff line change
@@ -1,17 +1,17 @@
from arro3.core import Scalar
from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
from arro3.core.types import ArrayInput, ArrowStreamExportable

def max(input: ArrowArrayExportable | ArrowStreamExportable) -> Scalar:
def max(input: ArrayInput | ArrowStreamExportable) -> Scalar:
"""
Returns the max of values in the array.
"""

def min(input: ArrowArrayExportable | ArrowStreamExportable) -> Scalar:
def min(input: ArrayInput | ArrowStreamExportable) -> Scalar:
"""
Returns the min of values in the array.
"""

def sum(input: ArrowArrayExportable | ArrowStreamExportable) -> Scalar:
def sum(input: ArrayInput | ArrowStreamExportable) -> Scalar:
"""
Returns the sum of values in the array.
"""
22 changes: 11 additions & 11 deletions arro3-compute/python/arro3/compute/_arith.pyi
Original file line number Diff line number Diff line change
@@ -1,32 +1,32 @@
from arro3.core import Array
from arro3.core.types import ArrowArrayExportable
from arro3.core.types import ArrayInput

def add(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def add(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs + rhs`, returning an error on overflow"""

def add_wrapping(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def add_wrapping(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs + rhs`, wrapping on overflow for integer data types."""

def div(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def div(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs / rhs`"""

def mul(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def mul(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs * rhs`, returning an error on overflow"""

def mul_wrapping(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def mul_wrapping(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs * rhs`, wrapping on overflow for integer data types."""

def neg(array: ArrowArrayExportable) -> Array:
def neg(array: ArrayInput) -> Array:
"""Negates each element of array, returning an error on overflow"""

def neg_wrapping(array: ArrowArrayExportable) -> Array:
def neg_wrapping(array: ArrayInput) -> Array:
"""Negates each element of array, wrapping on overflow for integer data types."""

def rem(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def rem(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs % rhs`"""

def sub(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def sub(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs - rhs`, returning an error on overflow"""

def sub_wrapping(lhs: ArrowArrayExportable, rhs: ArrowArrayExportable) -> Array:
def sub_wrapping(lhs: ArrayInput, rhs: ArrayInput) -> Array:
"""Perform `lhs - rhs`, wrapping on overflow for integer data types."""
10 changes: 5 additions & 5 deletions arro3-compute/python/arro3/compute/_boolean.pyi
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
from typing import overload

from arro3.core import Array, ArrayReader
from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
from arro3.core.types import ArrayInput, ArrowStreamExportable

@overload
def is_null(input: ArrowArrayExportable) -> Array: ...
def is_null(input: ArrayInput) -> Array: ...
@overload
def is_null(input: ArrowStreamExportable) -> ArrayReader: ...
def is_null(
input: ArrowArrayExportable | ArrowStreamExportable,
input: ArrayInput | ArrowStreamExportable,
) -> Array | ArrayReader:
"""
Returns a non-null boolean-typed array with whether each value of the array is null.
Expand All @@ -23,11 +23,11 @@ def is_null(
"""

@overload
def is_not_null(input: ArrowArrayExportable) -> Array: ...
def is_not_null(input: ArrayInput) -> Array: ...
@overload
def is_not_null(input: ArrowStreamExportable) -> ArrayReader: ...
def is_not_null(
input: ArrowArrayExportable | ArrowStreamExportable,
input: ArrayInput | ArrowStreamExportable,
) -> Array | ArrayReader:
"""
Returns a non-null boolean-typed array with whether each value of the array is not null.
Expand Down
6 changes: 3 additions & 3 deletions arro3-compute/python/arro3/compute/_cast.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,14 @@ from typing import overload

from arro3.core import Array, ArrayReader
from arro3.core.types import (
ArrowArrayExportable,
ArrayInput,
ArrowSchemaExportable,
ArrowStreamExportable,
)

@overload
def cast(
input: ArrowArrayExportable,
input: ArrayInput,
to_type: ArrowSchemaExportable,
) -> Array: ...
@overload
Expand All @@ -18,7 +18,7 @@ def cast(
to_type: ArrowSchemaExportable,
) -> ArrayReader: ...
def cast(
input: ArrowArrayExportable | ArrowStreamExportable,
input: ArrayInput | ArrowStreamExportable,
to_type: ArrowSchemaExportable,
) -> Array | ArrayReader:
"""
Expand Down
6 changes: 3 additions & 3 deletions arro3-compute/python/arro3/compute/_dictionary.pyi
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
from typing import overload

from arro3.core import Array, ArrayReader
from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
from arro3.core.types import ArrayInput, ArrowStreamExportable

@overload
def dictionary_encode(array: ArrowArrayExportable) -> Array: ...
def dictionary_encode(array: ArrayInput) -> Array: ...
@overload
def dictionary_encode(array: ArrowStreamExportable) -> ArrayReader: ...
def dictionary_encode(
array: ArrowArrayExportable | ArrowStreamExportable,
array: ArrayInput | ArrowStreamExportable,
) -> Array | ArrayReader:
"""
Dictionary-encode array.
Expand Down
10 changes: 5 additions & 5 deletions arro3-compute/python/arro3/compute/_filter.pyi
Original file line number Diff line number Diff line change
@@ -1,21 +1,21 @@
from typing import overload

from arro3.core import Array, ArrayReader
from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
from arro3.core.types import ArrayInput, ArrowStreamExportable

@overload
def filter(
values: ArrowArrayExportable,
predicate: ArrowArrayExportable,
values: ArrayInput,
predicate: ArrayInput,
) -> Array: ...
@overload
def filter(
values: ArrowStreamExportable,
predicate: ArrowStreamExportable,
) -> ArrayReader: ...
def filter(
values: ArrowArrayExportable | ArrowStreamExportable,
predicate: ArrowArrayExportable | ArrowStreamExportable,
values: ArrayInput | ArrowStreamExportable,
predicate: ArrayInput | ArrowStreamExportable,
) -> Array | ArrayReader:
"""
Returns a filtered `values` array where the corresponding elements of
Expand Down
4 changes: 2 additions & 2 deletions arro3-compute/python/arro3/compute/_take.pyi
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from arro3.core import Array
from arro3.core.types import ArrowArrayExportable
from arro3.core.types import ArrayInput

def take(values: ArrowArrayExportable, indices: ArrowArrayExportable) -> Array:
def take(values: ArrayInput, indices: ArrayInput) -> Array:
"""Take elements by index from Array, creating a new Array from those indexes.
```
Expand Down
Loading

0 comments on commit 9a01c80

Please sign in to comment.