GH-34359: [Python] Add select method to pyarrow.RecordBatch

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()
 

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(<int> 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.

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)))