pyarrow: Check compatibility of pyarrow.array with string type

pygmt/clib/conversion.py

@@ -9,6 +9,15 @@
 import numpy as np
 from pygmt.exceptions import GMTInvalidInput
 
+StringArrayTypes = Sequence[str]
+
+try:
+    import pyarrow as pa
+
+    StringArrayTypes |= pa.StringArray
+except ImportError:
+    pa = None
+
 
 def dataarray_to_matrix(grid):
     """
@@ -263,14 +272,15 @@ def sequence_to_ctypes_array(
     return (ctype * size)(*sequence)
 
 
-def strings_to_ctypes_array(strings: Sequence[str]) -> ctp.Array:
+def strings_to_ctypes_array(strings: StringArrayTypes) -> ctp.Array:
     """
-    Convert a sequence (e.g., a list) of strings into a ctypes array.
+    Convert a sequence (e.g., a list) of strings or a pyarrow.StringArray into a ctypes
+    array.
 
     Parameters
     ----------
     strings
-        A sequence of strings.
+        A sequence of strings or a pyarrow.StringArray.
 
     Returns
     -------
@@ -286,7 +296,12 @@ def strings_to_ctypes_array(strings: Sequence[str]) -> ctp.Array:
     >>> [s.decode() for s in ctypes_array]
     ['first', 'second', 'third']
     """
-    return (ctp.c_char_p * len(strings))(*[s.encode() for s in strings])
+    try:
+        bytes_string_list = [s.encode() for s in strings]
+    except AttributeError:  # 'pyarrow.StringScalar' object has no attribute 'encode'
+        # Convert pyarrow.StringArray to Python list first
+        bytes_string_list = [s.encode() for s in strings.to_pylist()]
+    return (ctp.c_char_p * len(strings))(*bytes_string_list)
 
 
 def array_to_datetime(array):

pygmt/clib/session.py

@@ -34,6 +34,16 @@
     tempfile_from_image,
 )
 
+StringArrayTypes = Sequence[str]
+
+try:
+    import pyarrow as pa
+
+    StringArrayTypes |= pa.StringArray
+except ImportError:
+    pa = None
+
+
 FAMILIES = [
     "GMT_IS_DATASET",  # Entity is a data table
     "GMT_IS_GRID",  # Entity is a grid
@@ -936,39 +946,43 @@ def put_vector(self, dataset, column, vector):
                 f"in column {column} of dataset."
             )
 
-    def put_strings(self, dataset, family, strings):
+    def put_strings(
+        self,
+        dataset: ctp.c_void_p,
+        family: Literal["GMT_IS_VECTOR", "GMT_IS_MATRIX"],
+        strings: StringArrayTypes,
+    ):
         """
-        Attach a numpy 1-D array of dtype str as a column on a GMT dataset.
+        Attach a 1-D numpy array of dtype str or pyarrow.StringArray as a column on a
+        GMT dataset.
 
-        Use this function to attach string type numpy array data to a GMT
-        dataset and pass it to GMT modules. Wraps ``GMT_Put_Strings``.
+        Use this function to attach string type array data to a GMT dataset and pass it
+        to GMT modules. Wraps ``GMT_Put_Strings``.
 
-        The dataset must be created by :meth:`pygmt.clib.Session.create_data`
-        first.
+        The dataset must be created by :meth:`pygmt.clib.Session.create_data` first.
 
         .. warning::
-            The numpy array must be C contiguous in memory. If it comes from a
-            column slice of a 2-D array, for example, you will have to make a
-            copy. Use :func:`numpy.ascontiguousarray` to make sure your vector
-            is contiguous (it won't copy if it already is).
+            The array must be C contiguous in memory. If it comes from a column slice of
+            a 2-D array, for example, you will have to make a copy. Use
+            :func:`numpy.ascontiguousarray` to make sure your vector is contiguous (it
+            won't copy if it already is).
 
         Parameters
         ----------
-        dataset : :class:`ctypes.c_void_p`
+        dataset
             The ctypes void pointer to a ``GMT_Dataset``. Create it with
             :meth:`pygmt.clib.Session.create_data`.
-        family : str
+        family
             The family type of the dataset. Can be either ``GMT_IS_VECTOR`` or
             ``GMT_IS_MATRIX``.
-        strings : numpy 1-D array
-            The array that will be attached to the dataset. Must be a 1-D C
-            contiguous array.
+        strings
+            The array that will be attached to the dataset. Must be a 1-D C contiguous
+            array.
 
         Raises
         ------
         GMTCLibError
-            If given invalid input or ``GMT_Put_Strings`` exits with
-            status != 0.
+            If given invalid input or ``GMT_Put_Strings`` exits with status != 0.
         """
         c_put_strings = self.get_libgmt_func(
             "GMT_Put_Strings",

pygmt/tests/test_clib_put_strings.py

@@ -8,12 +8,31 @@
 from pygmt import clib
 from pygmt.exceptions import GMTCLibError
 from pygmt.helpers import GMTTempFile
+from pygmt.helpers.testing import skip_if_no
+
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
 
 
 @pytest.mark.benchmark
-def test_put_strings():
+@pytest.mark.parametrize(
+    ("array_func", "dtype"),
+    [
+        pytest.param(np.array, {"dtype": str}, id="str"),
+        pytest.param(
+            getattr(pa, "array", None),
+            {"type": "string"},  # pa.string()
+            marks=skip_if_no(package="pyarrow"),
+            id="pyarrow",
+        ),
+    ],
+)
+def test_put_strings(array_func, dtype):
     """
-    Check that assigning a numpy array of dtype str to a dataset works.
+    Check that assigning a numpy array of dtype str, or a pyarrow.StringArray to a
+    dataset works.
     """
     with clib.Session() as lib:
         dataset = lib.create_data(
@@ -24,7 +43,7 @@ def test_put_strings():
         )
         x = np.array([1, 2, 3, 4, 5], dtype=np.int32)
         y = np.array([6, 7, 8, 9, 10], dtype=np.int32)
-        strings = np.array(["a", "bc", "defg", "hijklmn", "opqrst"], dtype=str)
+        strings = array_func(["a", "bc", "defg", "hijklmn", "opqrst"], **dtype)
         lib.put_vector(dataset, column=lib["GMT_X"], vector=x)
         lib.put_vector(dataset, column=lib["GMT_Y"], vector=y)
         lib.put_strings(

pygmt/tests/test_clib_virtualfiles.py

@@ -14,8 +14,14 @@
 from pygmt import clib
 from pygmt.exceptions import GMTCLibError, GMTInvalidInput
 from pygmt.helpers import GMTTempFile
+from pygmt.helpers.testing import skip_if_no
 from pygmt.tests.test_clib import mock
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
 POINTS_DATA = Path(__file__).parent / "data" / "points.txt"
 
 
@@ -225,16 +231,28 @@ def test_virtualfile_from_vectors(dtypes):
             assert output == expected
 
 
-@pytest.mark.benchmark
-@pytest.mark.parametrize("dtype", [str, object])
-def test_virtualfile_from_vectors_one_string_or_object_column(dtype):
+@pytest.mark.parametrize(
+    ("array_func", "dtype"),
+    [
+        pytest.param(np.array, {"dtype": str}, id="str"),
+        pytest.param(np.array, {"dtype": object}, id="object"),
+        pytest.param(
+            getattr(pa, "array", None),
+            {"type": "string"},  # pa.string()
+            marks=skip_if_no(package="pyarrow"),
+            id="pyarrow",
+        ),
+    ],
+)
+def test_virtualfile_from_vectors_one_string_or_object_column(array_func, dtype):
     """
-    Test passing in one column with string or object dtype into virtual file dataset.
+    Test passing in one column with string (numpy/pyarrow) or object (numpy)
+    dtype into virtual file dataset.
     """
     size = 5
     x = np.arange(size, dtype=np.int32)
     y = np.arange(size, size * 2, 1, dtype=np.int32)
-    strings = np.array(["a", "bc", "defg", "hijklmn", "opqrst"], dtype=dtype)
+    strings = array_func(["a", "bc", "defg", "hijklmn", "opqrst"], **dtype)
     with clib.Session() as lib:
         with lib.virtualfile_from_vectors(x, y, strings) as vfile:
             with GMTTempFile() as outfile:

pygmt/tests/test_text.py

@@ -9,6 +9,12 @@
 from pygmt import Figure
 from pygmt.exceptions import GMTCLibError, GMTInvalidInput
 from pygmt.helpers import GMTTempFile
+from pygmt.helpers.testing import skip_if_no
+
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
 
 TEST_DATA_DIR = Path(__file__).parent / "data"
 POINTS_DATA = TEST_DATA_DIR / "points.txt"
@@ -48,8 +54,20 @@ def test_text_single_line_of_text(region, projection):
 
 
 @pytest.mark.benchmark
-@pytest.mark.mpl_image_compare
-def test_text_multiple_lines_of_text(region, projection):
+@pytest.mark.mpl_image_compare(filename="test_text_multiple_lines_of_text.png")
+@pytest.mark.parametrize(
+    "array_func",
+    [
+        list,
+        pytest.param(np.array, id="numpy"),
+        pytest.param(
+            getattr(pa, "array", None),
+            marks=skip_if_no(package="pyarrow"),
+            id="pyarrow",
+        ),
+    ],
+)
+def test_text_multiple_lines_of_text(region, projection, array_func):
     """
     Place multiple lines of text at their respective x, y locations.
     """
@@ -59,7 +77,7 @@ def test_text_multiple_lines_of_text(region, projection):
         projection=projection,
         x=[1.2, 1.6],
         y=[0.6, 0.3],
-        text=["This is a line of text", "This is another line of text"],
+        text=array_func(["This is a line of text", "This is another line of text"]),
     )
     return fig