feat(geoarrow-pyarrow): Add make_point() helper (#40)

geoarrow-pyarrow/src/geoarrow/pyarrow/__init__.py

@@ -48,6 +48,7 @@
     infer_type_common,
     as_geoarrow,
     format_wkt,
+    make_point,
     unique_geometry_types,
     box,
     box_agg,
@@ -98,6 +99,7 @@
     "infer_type_common",
     "as_geoarrow",
     "format_wkt",
+    "make_point",
     "unique_geometry_types",
     "box",
     "box_agg",

geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py

@@ -309,6 +309,44 @@ def format_wkt(obj, precision=None, max_element_size_bytes=None):
     )
 
 
+def make_point(x, y, z=None, m=None, crs=None, crs_type=None):
+    """Create a geoarrow-encoded point array from two or more arrays
+    representing x, y, and/or z, and/or m values. In many cases, this
+    is a zero-copy operation if the input arrays are already in a
+    column-based format (e.g., numpy array, pandas series, or pyarrow
+    Array/ChunkedArray).
+
+    >>> import geoarrow.pyarrow as ga
+    >>> ga.make_point([1, 2, 3], [4, 5, 6])
+    PointArray:PointType(geoarrow.point)[3]
+    <POINT (1 4)>
+    <POINT (2 5)>
+    <POINT (3 6)>
+    """
+    import pyarrow.compute as pc
+
+    if z is not None and m is not None:
+        dimensions = Dimensions.XYZM
+        field_names = ["x", "y", "z", "m"]
+    elif m is not None:
+        dimensions = Dimensions.XYM
+        field_names = ["x", "y", "m"]
+    elif z is not None:
+        dimensions = Dimensions.XYZ
+        field_names = ["x", "y", "z"]
+    else:
+        dimensions = Dimensions.XY
+        field_names = ["x", "y"]
+
+    type = _type.extension_type(
+        GeometryType.POINT, dimensions, crs=crs, crs_type=crs_type
+    )
+    args = [x, y] + [el for el in [z, m] if el is not None]
+    args = [pa.array(el, pa.float64()) for el in args]
+    storage = pc.make_struct(*args, field_names=field_names)
+    return type.wrap_array(storage)
+
+
 def _box_point_struct(storage):
     arrays = storage.flatten()
     return pa.StructArray.from_arrays(

geoarrow-pyarrow/tests/test_compute.py

@@ -234,6 +234,48 @@ def test_as_geoarrow():
     assert array.type.geoarrow_id == ga.wkb().geoarrow_id
 
 
+def test_make_point():
+    xs = [1, 2, 3]
+    ys = [4, 5, 6]
+    zs = [7, 8, 9]
+    ms = [10, 11, 12]
+
+    xy = _compute.make_point(xs, ys)
+    assert xy.type.dimensions == ga.Dimensions.XY
+    assert _compute.format_wkt(xy).to_pylist() == [
+        "POINT (1 4)",
+        "POINT (2 5)",
+        "POINT (3 6)",
+    ]
+
+    xyz = _compute.make_point(xs, ys, zs)
+    assert xyz.type.dimensions == ga.Dimensions.XYZ
+    assert _compute.format_wkt(xyz).to_pylist() == [
+        "POINT Z (1 4 7)",
+        "POINT Z (2 5 8)",
+        "POINT Z (3 6 9)",
+    ]
+
+    xym = _compute.make_point(xs, ys, m=ms)
+    assert xym.type.dimensions == ga.Dimensions.XYM
+    assert _compute.format_wkt(xym).to_pylist() == [
+        "POINT M (1 4 10)",
+        "POINT M (2 5 11)",
+        "POINT M (3 6 12)",
+    ]
+
+    xyzm = _compute.make_point(xs, ys, zs, ms)
+    assert xyzm.type.dimensions == ga.Dimensions.XYZM
+    assert _compute.format_wkt(xyzm).to_pylist() == [
+        "POINT ZM (1 4 7 10)",
+        "POINT ZM (2 5 8 11)",
+        "POINT ZM (3 6 9 12)",
+    ]
+
+    xy_crs = _compute.make_point(xs, ys, crs="EPSG:1234")
+    assert xy_crs.type.crs == "EPSG:1234"
+
+
 def test_box():
     wkt_array = ga.array(["POINT (0 1)", "POINT (2 3)"])
     box = _compute.box(wkt_array)