Skip to content

Commit e8a6014

Browse files
dangotbannedzen-xu
andauthored
fix: Match runtime behavior of (Table|RecordBatch).select (#221)
* fix: Match runtime behavior of `(Table|RecordBatch).select` ## Resolves - https://github.com/MarcoGorelli/narwhals/blob/5b02b592183b8d39e2d32e0aedd6c234bb22d405/narwhals/_arrow/dataframe.py#L305-L307 - https://github.com/MarcoGorelli/narwhals/blob/5b02b592183b8d39e2d32e0aedd6c234bb22d405/narwhals/_arrow/dataframe.py#L285-L294 ##Description Following up on what I thought was a simple stub issue, but we're both *too strict* and *too permissive* in different ways ##Examples {placeholder} ##Related - https://github.com/apache/arrow/blob/d2ddee62329eb711572b4d71d6380673d7f7edd1/python/pyarrow/table.pxi#L4367-L4374 - https://github.com/apache/arrow/blob/d2ddee62329eb711572b4d71d6380673d7f7edd1/python/pyarrow/table.pxi#L1721-L1739 * update select * update select --------- Co-authored-by: ZhengYu, Xu <zen-xu@outlook.com>
1 parent dcadf97 commit e8a6014

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

pyarrow-stubs/__lib_pxi/table.pyi

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ from typing import (
2828
import numpy as np
2929
import pandas as pd
3030

31+
from numpy.typing import NDArray
3132
from pyarrow._compute import (
3233
CastOptions,
3334
CountOptions,
@@ -525,7 +526,7 @@ class RecordBatch(_Tabular[Array]):
525526
def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ...
526527
def slice(self, offset: int = 0, length: int | None = None) -> Self: ...
527528
def equals(self, other: Self, check_metadata: bool = False) -> bool: ...
528-
def select(self, columns: list[str] | Indices) -> Self: ...
529+
def select(self, columns: Iterable[str] | Iterable[int] | NDArray[np.str_]) -> Self: ...
529530
def cast(
530531
self, target_schema: Schema, safe: bool | None = None, options: CastOptions | None = None
531532
) -> Self: ...
@@ -592,7 +593,7 @@ JoinType: TypeAlias = Literal[
592593
class Table(_Tabular[ChunkedArray[Any]]):
593594
def validate(self, *, full=False) -> None: ...
594595
def slice(self, offset=0, length=None) -> Self: ...
595-
def select(self, columns: list[str] | Indices) -> Self: ...
596+
def select(self, columns: Iterable[str] | Iterable[int] | NDArray[np.str_]) -> Self: ...
596597
def replace_schema_metadata(self, metadata: dict | None = None) -> Self: ...
597598
def flatten(self, memory_pool: MemoryPool | None = None) -> Self: ...
598599
def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Self: ...

0 commit comments

Comments
 (0)