diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index 0fa23b7c272e..4725ae8f44d4 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -815,6 +815,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: const vector[shared_ptr[CArray]]& columns() + CResult[shared_ptr[CRecordBatch]] SelectColumns(const vector[int]&) + int num_columns() int64_t num_rows() diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi index 8b950c14b136..5b7989e58250 100644 --- a/python/pyarrow/table.pxi +++ b/python/pyarrow/table.pxi @@ -2254,6 +2254,56 @@ cdef class RecordBatch(_PandasConvertible): """ return _pc().drop_null(self) + def select(self, object columns): + """ + Select columns of the RecordBatch. + + Returns a new RecordBatch with the specified columns, and metadata + preserved. + + Parameters + ---------- + columns : list-like + The column names or integer indices to select. + + Returns + ------- + RecordBatch + + Examples + -------- + >>> import pyarrow as pa + >>> n_legs = pa.array([2, 2, 4, 4, 5, 100]) + >>> animals = pa.array(["Flamingo", "Parrot", "Dog", "Horse", "Brittle stars", "Centipede"]) + >>> batch = pa.record_batch([n_legs, animals], + ... names=["n_legs", "animals"]) + + Select columns my indices: + + >>> batch.select([1]) + pyarrow.RecordBatch + animals: string + + Select columns by names: + + >>> batch.select(["n_legs"]) + pyarrow.RecordBatch + n_legs: int64 + """ + cdef: + shared_ptr[CRecordBatch] c_batch + vector[int] c_indices + + for idx in columns: + idx = self._ensure_integer_index(idx) + idx = _normalize_index(idx, self.num_columns) + c_indices.push_back( idx) + + with nogil: + c_batch = GetResultValue(self.batch.SelectColumns(move(c_indices))) + + return pyarrow_wrap_batch(c_batch) + def sort_by(self, sorting, **kwargs): """ Sort the RecordBatch by one or multiple columns. diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index 42419575a8d5..1a0eaa33772e 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -697,6 +697,57 @@ def test_recordbatch_select_column(): batch.column(4) +def test_recordbatch_select(): + a1 = pa.array([1, 2, 3, None, 5]) + a2 = pa.array(['a', 'b', 'c', 'd', 'e']) + a3 = pa.array([[1, 2], [3, 4], [5, 6], None, [9, 10]]) + batch = pa.record_batch([a1, a2, a3], ['f1', 'f2', 'f3']) + + # selecting with string names + result = batch.select(['f1']) + expected = pa.record_batch([a1], ['f1']) + assert result.equals(expected) + + result = batch.select(['f3', 'f2']) + expected = pa.record_batch([a3, a2], ['f3', 'f2']) + assert result.equals(expected) + + # selecting with integer indices + result = batch.select([0]) + expected = pa.record_batch([a1], ['f1']) + assert result.equals(expected) + + result = batch.select([2, 1]) + expected = pa.record_batch([a3, a2], ['f3', 'f2']) + assert result.equals(expected) + + # preserve metadata + batch2 = batch.replace_schema_metadata({"a": "test"}) + result = batch2.select(["f1", "f2"]) + assert b"a" in result.schema.metadata + + # selecting non-existing column raises + with pytest.raises(KeyError, match='Field "f5" does not exist'): + batch.select(['f5']) + + with pytest.raises(IndexError, match="index out of bounds"): + batch.select([5]) + + # duplicate selection gives duplicated names in resulting recordbatch + result = batch.select(['f2', 'f2']) + expected = pa.record_batch([a2, a2], ['f2', 'f2']) + assert result.equals(expected) + + # selection duplicated column raises + batch = pa.record_batch([a1, a2, a3], ['f1', 'f2', 'f1']) + with pytest.raises(KeyError, match='Field "f1" exists 2 times'): + batch.select(['f1']) + + result = batch.select(['f2']) + expected = pa.record_batch([a2], ['f2']) + assert result.equals(expected) + + def test_recordbatch_from_struct_array_invalid(): with pytest.raises(TypeError): pa.RecordBatch.from_struct_array(pa.array(range(5)))