Skip to content

Commit 6a4bcb3

Browse files
authored
GH-34359: [Python] Add select method to pyarrow.RecordBatch (#34360)
### Rationale for this change There is a `select` method defined for `pa.Table` and we should add the same for `pa.RecordBatch.` ### What changes are included in this PR? Added method to `RecordBatch` class in `table.pxi`. ### Are these changes tested? Yes, tests are added to _python/pyarrow/tests/test_table.py_. ### Are there any user-facing changes? No. * Closes: #34359 Authored-by: Alenka Frim <frim.alenka@gmail.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
1 parent 52f4fe1 commit 6a4bcb3

3 files changed

Lines changed: 103 additions & 0 deletions

File tree

python/pyarrow/includes/libarrow.pxd

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -815,6 +815,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
815815

816816
const vector[shared_ptr[CArray]]& columns()
817817

818+
CResult[shared_ptr[CRecordBatch]] SelectColumns(const vector[int]&)
819+
818820
int num_columns()
819821
int64_t num_rows()
820822

python/pyarrow/table.pxi

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2254,6 +2254,56 @@ cdef class RecordBatch(_PandasConvertible):
22542254
"""
22552255
return _pc().drop_null(self)
22562256

2257+
def select(self, object columns):
2258+
"""
2259+
Select columns of the RecordBatch.
2260+
2261+
Returns a new RecordBatch with the specified columns, and metadata
2262+
preserved.
2263+
2264+
Parameters
2265+
----------
2266+
columns : list-like
2267+
The column names or integer indices to select.
2268+
2269+
Returns
2270+
-------
2271+
RecordBatch
2272+
2273+
Examples
2274+
--------
2275+
>>> import pyarrow as pa
2276+
>>> n_legs = pa.array([2, 2, 4, 4, 5, 100])
2277+
>>> animals = pa.array(["Flamingo", "Parrot", "Dog", "Horse", "Brittle stars", "Centipede"])
2278+
>>> batch = pa.record_batch([n_legs, animals],
2279+
... names=["n_legs", "animals"])
2280+
2281+
Select columns my indices:
2282+
2283+
>>> batch.select([1])
2284+
pyarrow.RecordBatch
2285+
animals: string
2286+
2287+
Select columns by names:
2288+
2289+
>>> batch.select(["n_legs"])
2290+
pyarrow.RecordBatch
2291+
n_legs: int64
2292+
"""
2293+
cdef:
2294+
shared_ptr[CRecordBatch] c_batch
2295+
vector[int] c_indices
2296+
2297+
for idx in columns:
2298+
idx = self._ensure_integer_index(idx)
2299+
idx = _normalize_index(idx, self.num_columns)
2300+
c_indices.push_back(<int> idx)
2301+
2302+
with nogil:
2303+
c_batch = GetResultValue(self.batch.SelectColumns(move(c_indices)))
2304+
2305+
return pyarrow_wrap_batch(c_batch)
2306+
22572307
def sort_by(self, sorting, **kwargs):
22582308
"""
22592309
Sort the RecordBatch by one or multiple columns.

python/pyarrow/tests/test_table.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -697,6 +697,57 @@ def test_recordbatch_select_column():
697697
batch.column(4)
698698

699699

700+
def test_recordbatch_select():
701+
a1 = pa.array([1, 2, 3, None, 5])
702+
a2 = pa.array(['a', 'b', 'c', 'd', 'e'])
703+
a3 = pa.array([[1, 2], [3, 4], [5, 6], None, [9, 10]])
704+
batch = pa.record_batch([a1, a2, a3], ['f1', 'f2', 'f3'])
705+
706+
# selecting with string names
707+
result = batch.select(['f1'])
708+
expected = pa.record_batch([a1], ['f1'])
709+
assert result.equals(expected)
710+
711+
result = batch.select(['f3', 'f2'])
712+
expected = pa.record_batch([a3, a2], ['f3', 'f2'])
713+
assert result.equals(expected)
714+
715+
# selecting with integer indices
716+
result = batch.select([0])
717+
expected = pa.record_batch([a1], ['f1'])
718+
assert result.equals(expected)
719+
720+
result = batch.select([2, 1])
721+
expected = pa.record_batch([a3, a2], ['f3', 'f2'])
722+
assert result.equals(expected)
723+
724+
# preserve metadata
725+
batch2 = batch.replace_schema_metadata({"a": "test"})
726+
result = batch2.select(["f1", "f2"])
727+
assert b"a" in result.schema.metadata
728+
729+
# selecting non-existing column raises
730+
with pytest.raises(KeyError, match='Field "f5" does not exist'):
731+
batch.select(['f5'])
732+
733+
with pytest.raises(IndexError, match="index out of bounds"):
734+
batch.select([5])
735+
736+
# duplicate selection gives duplicated names in resulting recordbatch
737+
result = batch.select(['f2', 'f2'])
738+
expected = pa.record_batch([a2, a2], ['f2', 'f2'])
739+
assert result.equals(expected)
740+
741+
# selection duplicated column raises
742+
batch = pa.record_batch([a1, a2, a3], ['f1', 'f2', 'f1'])
743+
with pytest.raises(KeyError, match='Field "f1" exists 2 times'):
744+
batch.select(['f1'])
745+
746+
result = batch.select(['f2'])
747+
expected = pa.record_batch([a2], ['f2'])
748+
assert result.equals(expected)
749+
750+
700751
def test_recordbatch_from_struct_array_invalid():
701752
with pytest.raises(TypeError):
702753
pa.RecordBatch.from_struct_array(pa.array(range(5)))

0 commit comments

Comments
 (0)