GH-35415: [Python] RecordBatch string reprsentation includes column preview (#35416)

danepitkin · web-flow · commit b211c1896916 · 2023-05-08T20:41:58.000+02:00
### Rationale for this change Table and RecordBatch now share a common parent class and common APIs should behave the same. ### What changes are included in this PR? Remove override of RecordBatch string representation. ### Are these changes tested? Pytests and doctests updated. ### Are there any user-facing changes? Yes, the string representation of `RecordBatch` includes additional info. * Closes: #35415 Authored-by: Dane Pitkin <dane@voltrondata.com> Signed-off-by: Alenka Frim <frim.alenka@gmail.com>
diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
@@ -811,6 +811,9 @@ cdef class _PandasConvertible(_Weakrefable):
         pyarrow.RecordBatch
         n_legs: int64
         animals: string
+        ----
+        n_legs: [2,4,5,100]
+        animals: ["Flamingo","Horse","Brittle stars","Centipede"]
         >>> batch.to_pandas()
            n_legs        animals
         0       2       Flamingo
diff --git a/python/pyarrow/ipc.pxi b/python/pyarrow/ipc.pxi
@@ -645,8 +645,12 @@ cdef class RecordBatchReader(_Weakrefable):
     ...     print(batch)
     pyarrow.RecordBatch
     x: int64
+    ----
+    x: [1,2,3]
     pyarrow.RecordBatch
     x: int64
+    ----
+    x: [1,2,3]
     """
 
     # cdef block is in lib.pxd
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
@@ -1594,6 +1594,9 @@ cdef class RecordBatch(_Tabular):
     pyarrow.RecordBatch
     n_legs: int64
     animals: string
+    ----
+    n_legs: [2,2,4,4,5,100]
+    animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
     >>> pa.RecordBatch.from_arrays([n_legs, animals], names=names).to_pandas()
        n_legs        animals
     0       2       Flamingo
@@ -1618,6 +1621,12 @@ cdef class RecordBatch(_Tabular):
     day: int64
     n_legs: int64
     animals: string
+    ----
+    year: [2020,2022,2021,2022]
+    month: [3,5,7,9]
+    day: [1,5,9,13]
+    n_legs: [2,4,5,100]
+    animals: ["Flamingo","Horse","Brittle stars","Centipede"]
     >>> pa.RecordBatch.from_pandas(df).to_pandas()
        year  month  day  n_legs        animals
     0  2020      3    1       2       Flamingo
@@ -1652,6 +1661,12 @@ cdef class RecordBatch(_Tabular):
     day: int64
     n_legs: int64
     animals: string
+    ----
+    year: [2020,2022,2021,2022]
+    month: [3,5,7,9]
+    day: [1,5,9,13]
+    n_legs: [2,4,5,100]
+    animals: ["Flamingo","Horse","Brittle stars","Centipede"]
     """
 
     def __cinit__(self):
@@ -1733,6 +1748,9 @@ cdef class RecordBatch(_Tabular):
         pyarrow.RecordBatch
         n_legs: int64
         animals: string
+        ----
+        n_legs: [2,2,4,4,5,100]
+        animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
         >>> pa.RecordBatch.from_pydict(pydict).to_pandas()
            n_legs        animals
         0       2       Flamingo
@@ -1788,6 +1806,10 @@ cdef class RecordBatch(_Tabular):
         pyarrow.RecordBatch
         n_legs: int64
         animals: string
+        ----
+        n_legs: [2,4]
+        animals: ["Flamingo","Dog"]
+
         >>> pa.RecordBatch.from_pylist(pylist).to_pandas()
            n_legs   animals
         0       2  Flamingo
@@ -1820,12 +1842,6 @@ cdef class RecordBatch(_Tabular):
         except TypeError:
             return NotImplemented
 
-    def __repr__(self):
-        # TODO remove this and update pytests/doctests for
-        # RecordBatch.to_string(preview_cols=10) usage in
-        # parent class
-        return self.to_string()
-
     def validate(self, *, full=False):
         """
         Perform validation checks.  An exception is raised if validation fails.
@@ -2402,12 +2418,16 @@ cdef class RecordBatch(_Tabular):
         >>> batch.select([1])
         pyarrow.RecordBatch
         animals: string
+        ----
+        animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
 
         Select columns by names:
 
         >>> batch.select(["n_legs"])
         pyarrow.RecordBatch
         n_legs: int64
+        ----
+        n_legs: [2,2,4,4,5,100]
         """
         cdef:
             shared_ptr[CRecordBatch] c_batch
@@ -2558,6 +2578,12 @@ cdef class RecordBatch(_Tabular):
         day: int64
         n_legs: int64
         animals: string
+        ----
+        year: [2020,2022,2021,2022]
+        month: [3,5,7,9]
+        day: [1,5,9,13]
+        n_legs: [2,4,5,100]
+        animals: ["Flamingo","Horse","Brittle stars","Centipede"]
 
         Convert pandas DataFrame to RecordBatch using schema:
 
@@ -2569,12 +2595,17 @@ cdef class RecordBatch(_Tabular):
         pyarrow.RecordBatch
         n_legs: int64
         animals: string
+        ----
+        n_legs: [2,4,5,100]
+        animals: ["Flamingo","Horse","Brittle stars","Centipede"]
 
         Convert pandas DataFrame to RecordBatch specifying columns:
 
         >>> pa.RecordBatch.from_pandas(df, columns=["n_legs"])
         pyarrow.RecordBatch
         n_legs: int64
+        ----
+        n_legs: [2,4,5,100]
         """
         from pyarrow.pandas_compat import dataframe_to_arrays
         arrays, schema, n_rows = dataframe_to_arrays(
@@ -2622,6 +2653,9 @@ cdef class RecordBatch(_Tabular):
         pyarrow.RecordBatch
         n_legs: int64
         animals: string
+        ----
+        n_legs: [2,2,4,4,5,100]
+        animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
         >>> pa.RecordBatch.from_arrays([n_legs, animals], names=names).to_pandas()
            n_legs        animals
         0       2       Flamingo
@@ -5063,6 +5097,9 @@ def record_batch(data, names=None, schema=None, metadata=None):
     pyarrow.RecordBatch
     n_legs: int64
     animals: string
+    ----
+    n_legs: [2,2,4,4,5,100]
+    animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
     >>> pa.record_batch([n_legs, animals], names=["n_legs", "animals"]).to_pandas()
        n_legs        animals
     0       2       Flamingo
@@ -5081,6 +5118,9 @@ def record_batch(data, names=None, schema=None, metadata=None):
     pyarrow.RecordBatch
     n_legs: int64
     animals: string
+    ----
+    n_legs: [2,2,4,4,5,100]
+    animals: ["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]
     >>> pa.record_batch([n_legs, animals],
     ...                  names=names,
     ...                  metadata = my_metadata).schema
@@ -5104,6 +5144,13 @@ def record_batch(data, names=None, schema=None, metadata=None):
     day: int64
     n_legs: int64
     animals: string
+    ----
+    year: [2020,2022,2021,2022]
+    month: [3,5,7,9]
+    day: [1,5,9,13]
+    n_legs: [2,4,5,100]
+    animals: ["Flamingo","Horse","Brittle stars","Centipede"]
+
     >>> pa.record_batch(df).to_pandas()
        year  month  day  n_legs        animals
     0  2020      3    1       2       Flamingo
diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py
@@ -529,7 +529,10 @@ def test_recordbatch_basics():
     assert batch.schema == schema
     assert str(batch) == """pyarrow.RecordBatch
 c0: int16
-c1: int32"""
+c1: int32
+----
+c0: [0,1,2,3,4]
+c1: [-10,-5,0,null,10]"""
 
     assert batch.to_string(show_metadata=True) == """\
 pyarrow.RecordBatch
