Skip to content

Commit

Permalink
support pyarrow compute funcs (#61)
Browse files Browse the repository at this point in the history
* update compute.pyi

* impl Aggregation funcs

* impl arithmetic

* imit bit-wise functions

* imit rounding functions

* optimize annotation

* impl logarithmic functions

* update

* impl comparisons funcs

* impl logical funcs

* impl string predicates and transforms

* impl string padding

* impl string trimming

* impl string splitting and component extraction

* impl string joining and slicing

* impl Containment tests

* impl Categorizations

* impl Structural transforms

* impl Conversions

* impl Temporal component extraction

* impl random, Timezone handling

* impl Array-wise functions

* fix timestamp scalar
  • Loading branch information
zen-xu authored Sep 4, 2024
1 parent 43b8bc0 commit 343f221
Show file tree
Hide file tree
Showing 3 changed files with 1,849 additions and 39 deletions.
13 changes: 12 additions & 1 deletion pyarrow-stubs/__lib_pxi/scalar.pyi
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# mypy: disable-error-code="overload-overlap,misc"
# mypy: disable-error-code="overload-overlap,misc,type-arg"
import collections.abc
import datetime as dt

Expand Down Expand Up @@ -281,6 +281,10 @@ def scalar(
value: Decimal, *, from_pandas: bool | None = None, memory_pool: MemoryPool | None = None
) -> Decimal128Scalar: ...
@overload
def scalar(
value: dt.datetime, *, from_pandas: bool | None = None, memory_pool: MemoryPool | None = None
) -> TimestampScalar: ...
@overload
def scalar(
value: dt.date, *, from_pandas: bool | None = None, memory_pool: MemoryPool | None = None
) -> Date32Scalar: ...
Expand Down Expand Up @@ -349,6 +353,13 @@ def scalar(
memory_pool: MemoryPool | None = None,
) -> ListScalar[types.ListType[types.Decimal128Type]]: ...
@overload
def scalar(
value: CollectionValue[dt.datetime],
*,
from_pandas: bool | None = None,
memory_pool: MemoryPool | None = None,
) -> ListScalar[types.ListType[types.TimestampType]]: ...
@overload
def scalar(
value: CollectionValue[dt.date],
*,
Expand Down
7 changes: 6 additions & 1 deletion pyarrow-stubs/__lib_pxi/table.pyi
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# mypy: disable-error-code="overload-overlap"
# mypy: disable-error-code="overload-overlap,type-arg,misc"

import datetime as dt

Expand Down Expand Up @@ -162,6 +162,11 @@ def chunked_array(
type: None = None,
) -> ChunkedArray[scalar.StructScalar]: ...
@overload
def chunked_array(
values: NullableIterable[dt.datetime],
type: None = None,
) -> ChunkedArray[scalar.TimestampScalar]: ...
@overload
def chunked_array(
values: NullableIterable[dt.date],
type: None = None,
Expand Down
Loading

0 comments on commit 343f221

Please sign in to comment.