pyarrow: Check compatibility of pyarrow.array with string type

doc/conf.py

@@ -85,6 +85,7 @@
     "contextily": ("https://contextily.readthedocs.io/en/stable/", None),
     "geopandas": ("https://geopandas.org/en/stable/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
+    "pyarrow": ("https://arrow.apache.org/docs/", None),
     "python": ("https://docs.python.org/3/", None),
     "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
     "rasterio": ("https://rasterio.readthedocs.io/en/stable/", None),

doc/ecosystem.md

@@ -94,9 +94,9 @@ Python objects. They are based on the C++ implementation of Arrow.
 ```{note}
 If you have [PyArrow][] installed, PyGMT does have some initial support for
 `pandas.Series` and `pandas.DataFrame` objects with Apache Arrow-backed arrays.
-Specifically, only uint/int/float and date32/date64 are supported for now.
-Support for string Array dtypes, Duration types and GeoArrow geometry types is still a work in progress.
-For more details, see
+Specifically, only uint/int/float, date32/date64 and string types are supported for now.
+Support for Duration types and GeoArrow geometry types is still a work in progress. For
+more details, see
 [issue #2800](https://github.com/GenericMappingTools/pygmt/issues/2800).
 ```
 

pygmt/_typing.py

@@ -2,7 +2,17 @@
 Type aliases for type hints.
 """
 
+import contextlib
+import importlib
+from collections.abc import Sequence
 from typing import Literal
 
+import numpy as np
+
 # Anchor codes
 AnchorCode = Literal["TL", "TC", "TR", "ML", "MC", "MR", "BL", "BC", "BR"]
+
+# String array types
+StringArrayTypes = Sequence[str] | np.ndarray
+with contextlib.suppress(ImportError):
+    StringArrayTypes |= importlib.import_module(name="pyarrow").StringArray

pygmt/clib/conversion.py

@@ -11,6 +11,7 @@
 import pandas as pd
 import xarray as xr
 from packaging.version import Version
+from pygmt._typing import StringArrayTypes
 from pygmt.exceptions import GMTInvalidInput
 
 
@@ -273,14 +274,15 @@ def sequence_to_ctypes_array(
     return (ctype * size)(*sequence)
 
 
-def strings_to_ctypes_array(strings: Sequence[str] | np.ndarray) -> 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) or numpy.ndarray of strings or a
+    pyarrow.StringArray into a ctypes array.
 
     Parameters
     ----------
     strings
-        A sequence of strings.
+        A sequence of strings, a numpy.ndarray of str dtype, or a pyarrow.StringArray.
 
     Returns
     -------
@@ -296,7 +298,7 @@ def strings_to_ctypes_array(strings: Sequence[str] | np.ndarray) -> 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])
+    return (ctp.c_char_p * len(strings))(*[s.encode() for s in np.asarray(strings)])
 
 
 def array_to_datetime(array: Sequence[Any] | np.ndarray) -> np.ndarray:

pygmt/clib/session.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
+from pygmt._typing import StringArrayTypes
 from pygmt.clib.conversion import (
     array_to_datetime,
     dataarray_to_matrix,
@@ -999,9 +1000,12 @@ def put_vector(self, dataset: ctp.c_void_p, column: int, vector: np.ndarray):
             )
             raise GMTCLibError(msg)
 
-    def put_strings(self, dataset: ctp.c_void_p, family: str, strings: np.ndarray):
+    def put_strings(
+        self, dataset: ctp.c_void_p, family: str, strings: StringArrayTypes
+    ):
         """
-        Attach a 1-D numpy 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``.
@@ -1050,7 +1054,8 @@ def put_strings(self, dataset: ctp.c_void_p, family: str, strings: np.ndarray):
             self.session_pointer, family_int, dataset, strings_pointer
         )
         if status != 0:
-            msg = f"Failed to put strings of type {strings.dtype} into dataset."
+            dtype = strings.dtype if hasattr(strings, "dtype") else type(strings)
+            msg = f"Failed to put strings of type {dtype} into dataset."
             raise GMTCLibError(msg)
 
     def put_matrix(self, dataset: ctp.c_void_p, matrix: np.ndarray, pad: int = 0):

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": np.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=np.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

@@ -10,8 +10,14 @@
 from pygmt.clib.session import DTYPES_NUMERIC
 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"
 
 
@@ -137,3 +143,37 @@ def test_open_virtual_file():
         bounds = "\t".join([f"<{col.min():.0f}/{col.max():.0f}>" for col in data.T])
         expected = f"<matrix memory>: N = {shape[0]}\t{bounds}\n"
         assert output == expected
+
+
+@pytest.mark.benchmark
+@pytest.mark.parametrize(
+    ("array_func", "dtype"),
+    [
+        pytest.param(np.array, {"dtype": np.str_}, id="str"),
+        pytest.param(np.array, {"dtype": np.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 (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 = 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:
+                lib.call_module("convert", [vfile, f"->{outfile.name}"])
+                output = outfile.read(keep_tabs=True)
+        expected = "".join(
+            f"{i}\t{j}\t{k}\n" for i, j, k in zip(x, y, strings, strict=True)
+        )
+        assert output == expected

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