diff --git a/geoarrow-pyarrow/src/geoarrow/pyarrow/__init__.py b/geoarrow-pyarrow/src/geoarrow/pyarrow/__init__.py index 8680888..59a24c3 100644 --- a/geoarrow-pyarrow/src/geoarrow/pyarrow/__init__.py +++ b/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", diff --git a/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py b/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py index 76cf6d0..50a5b4f 100644 --- a/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py +++ b/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] + + + + """ + 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( diff --git a/geoarrow-pyarrow/tests/test_compute.py b/geoarrow-pyarrow/tests/test_compute.py index dac66a2..08ebdca 100644 --- a/geoarrow-pyarrow/tests/test_compute.py +++ b/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)