GH-34686: [Python] Add RunEndEncodedScalar class (#34924)

jorisvandenbossche · web-flow · commit 84e54308f139 · 2023-04-06T17:09:52.000+02:00
### Rationale for this change Follow-up on #34570 (exposing the new RunEndEncoded array and type in pyarrow) to also add the scalar class. ### Are there any user-facing changes? The example from the issue now works: ``` In [15]: run_ends = [3, 5, 10, 19] ...: values = [1, 2, 1, 3] ...: ree_array = pa.RunEndEncodedArray.from_arrays(run_ends, values) In [16]: ree_array[0] Out[16]: <pyarrow.RunEndEncodedScalar: 1> ``` * Closes: #34686 Authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
diff --git a/docs/source/python/api/arrays.rst b/docs/source/python/api/arrays.rst
@@ -123,6 +123,7 @@ classes may expose data type-specific methods or properties.
    MonthDayNanoIntervalScalar
    Decimal128Scalar
    DictionaryScalar
+   RunEndEncodedScalar
    ListScalar
    LargeListScalar
    MapScalar
diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
@@ -225,7 +225,7 @@ def print_entry(label, value):
                          StringScalar, LargeStringScalar,
                          FixedSizeBinaryScalar, DictionaryScalar,
                          MapScalar, StructScalar, UnionScalar,
-                         ExtensionScalar)
+                         RunEndEncodedScalar, ExtensionScalar)
 
 # Buffers, allocation
 from pyarrow.lib import (Buffer, ResizableBuffer, foreign_buffer, py_buffer,
diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd
@@ -1155,6 +1155,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         vector[shared_ptr[CScalar]] value
         int child_id
 
+    cdef cppclass CRunEndEncodedScalar" arrow::RunEndEncodedScalar"(CScalar):
+        shared_ptr[CScalar] value
+
     cdef cppclass CExtensionScalar" arrow::ExtensionScalar"(CScalar):
         CExtensionScalar(shared_ptr[CScalar] storage,
                          shared_ptr[CDataType], c_bool is_valid)
diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi
@@ -880,6 +880,25 @@ cdef class DictionaryScalar(Scalar):
         return self.value.as_py() if self.is_valid else None
 
 
+cdef class RunEndEncodedScalar(Scalar):
+    """
+    Concrete class for RunEndEncoded scalars.
+    """
+    @property
+    def value(self):
+        """
+        Return underlying value as a scalar.
+        """
+        cdef CRunEndEncodedScalar* sp = <CRunEndEncodedScalar*> self.wrapped.get()
+        return Scalar.wrap(sp.value)
+
+    def as_py(self):
+        """
+        Return underlying value as a Python object.
+        """
+        return self.value.as_py()
+
+
 cdef class UnionScalar(Scalar):
     """
     Concrete class for Union scalars.
@@ -1010,6 +1029,7 @@ cdef dict _scalar_classes = {
     _Type_STRUCT: StructScalar,
     _Type_MAP: MapScalar,
     _Type_DICTIONARY: DictionaryScalar,
+    _Type_RUN_END_ENCODED: RunEndEncodedScalar,
     _Type_SPARSE_UNION: UnionScalar,
     _Type_DENSE_UNION: UnionScalar,
     _Type_INTERVAL_MONTH_DAY_NANO: MonthDayNanoIntervalScalar,
diff --git a/python/pyarrow/tests/test_misc.py b/python/pyarrow/tests/test_misc.py
@@ -203,6 +203,7 @@ def test_runtime_info():
     pa.UnionScalar,
     pa.StructScalar,
     pa.DictionaryScalar,
+    pa.RunEndEncodedScalar,
     pa.ipc.Message,
     pa.ipc.MessageReader,
     pa.MemoryPool,
diff --git a/python/pyarrow/tests/test_scalars.py b/python/pyarrow/tests/test_scalars.py
@@ -681,6 +681,27 @@ def test_dictionary():
         assert restored.equals(s)
 
 
+def test_run_end_encoded():
+    run_ends = [3, 5, 10, 12, 19]
+    values = [1, 2, 1, None, 3]
+    arr = pa.RunEndEncodedArray.from_arrays(run_ends, values)
+
+    scalar = arr[0]
+    assert isinstance(scalar, pa.RunEndEncodedScalar)
+    assert isinstance(scalar.value, pa.Int64Scalar)
+    assert scalar.value == pa.array(values)[0]
+    assert scalar.as_py() == 1
+
+    # null -> .value is still a scalar, as_py returns None
+    scalar = arr[10]
+    assert isinstance(scalar.value, pa.Int64Scalar)
+    assert scalar.as_py() is None
+
+    # constructing a scalar directly doesn't work yet
+    with pytest.raises(NotImplementedError):
+        pa.scalar(1, pa.run_end_encoded(pa.int64(), pa.int64()))
+
+
 def test_union():
     # sparse
     arr = pa.UnionArray.from_sparse(
