diff --git a/lonboard/_geoarrow/ops/coord_layout.py b/lonboard/_geoarrow/ops/coord_layout.py index b23b8ff7..728dd9bf 100644 --- a/lonboard/_geoarrow/ops/coord_layout.py +++ b/lonboard/_geoarrow/ops/coord_layout.py @@ -12,6 +12,7 @@ ) from lonboard._constants import EXTENSION_NAME +from lonboard._geoarrow.extension_types import CoordinateDimension, coord_storage_type from lonboard._geoarrow.ops.reproject import ( _map_coords_nest_0, _map_coords_nest_1, @@ -24,7 +25,7 @@ def make_geometry_interleaved( table: Table, ) -> Table: - """Convert geometry columns in table to interleaved coordinate layout.""" + """Convert geometry column in table from struct to interleaved coordinate layout.""" geom_col_idx = get_geometry_column_index(table.schema) # No geometry column in table if geom_col_idx is None: @@ -38,28 +39,29 @@ def make_geometry_interleaved( if geom_field.metadata.get(b"ARROW:extension:name") == EXTENSION_NAME.BOX: return table - new_field, new_column = transpose_column(field=geom_field, column=geom_column) + new_field, new_column = convert_struct_column_to_interleaved( + field=geom_field, + column=geom_column, + ) return table.set_column(geom_col_idx, new_field, new_column) -def transpose_column( +def convert_struct_column_to_interleaved( *, field: Field, column: ChunkedArray, ) -> tuple[Field, ChunkedArray]: + """Convert a GeoArrow column from struct to interleaved coordinate layout.""" extension_type_name = field.metadata[b"ARROW:extension:name"] - new_chunked_array = _transpose_column( - column, - extension_type_name=extension_type_name, # type: ignore - ) + new_chunked_array = _convert_column(column, extension_type_name=extension_type_name) return field.with_type(new_chunked_array.type), new_chunked_array -def _transpose_column( +def _convert_column( column: ChunkedArray, *, - extension_type_name: EXTENSION_NAME, + extension_type_name: bytes, ) -> ChunkedArray: if extension_type_name == EXTENSION_NAME.POINT: func = _transpose_chunk_nest_0 @@ -74,9 +76,13 @@ def _transpose_column( elif extension_type_name == EXTENSION_NAME.MULTIPOLYGON: func = _transpose_chunk_nest_3 else: - raise ValueError(f"Unexpected extension type name {extension_type_name}") + raise ValueError(f"Unexpected extension type name {extension_type_name!r}") - return ChunkedArray([func(chunk) for chunk in column.chunks]) + arrays = [func(chunk) for chunk in column.chunks] + return ChunkedArray( + arrays, + type=arrays[0].field.with_metadata(column.field.metadata), + ) def _transpose_coords(arr: Array) -> Array: @@ -87,14 +93,22 @@ def _transpose_coords(arr: Array) -> Array: x = struct_field(arr, [0]).to_numpy() y = struct_field(arr, [1]).to_numpy() coords = np.column_stack([x, y]).ravel("C") - return fixed_size_list_array(coords, 2) + return fixed_size_list_array( + coords, + 2, + type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY), + ) if arr.type.num_fields == 3: x = struct_field(arr, [0]).to_numpy() y = struct_field(arr, [1]).to_numpy() z = struct_field(arr, [2]).to_numpy() coords = np.column_stack([x, y, z]).ravel("C") - return fixed_size_list_array(coords, 3) + return fixed_size_list_array( + coords, + 3, + type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XYZ), + ) raise ValueError(f"Expected struct with 2 or 3 fields, got {arr.type.num_fields}") diff --git a/lonboard/experimental/_layer.py b/lonboard/experimental/_layer.py index 9cf6a793..7c9a9c0b 100644 --- a/lonboard/experimental/_layer.py +++ b/lonboard/experimental/_layer.py @@ -25,6 +25,7 @@ PointAccessor, TextAccessor, ) +from lonboard.types.layer import ArcLayerKwargs, PointAccessorInput if TYPE_CHECKING: import sys @@ -59,9 +60,6 @@ class ArcLayer(BaseArrowLayer): This is the fastest way to plot data from an existing GeoArrow source, such as [geoarrow-rust](https://geoarrow.github.io/geoarrow-rs/python/latest) or [geoarrow-pyarrow](https://geoarrow.github.io/geoarrow-python/main/index.html). - - If you have a GeoPandas `GeoDataFrame`, use - [`from_geopandas`][lonboard.ScatterplotLayer.from_geopandas] instead. """ great_circle = t.Bool(None, allow_none=True).tag(sync=True) @@ -152,6 +150,31 @@ class ArcLayer(BaseArrowLayer): - Default: `0`. """ + def __init__( + self, + *, + table: ArrowStreamExportable, + get_source_position: PointAccessorInput, + get_target_position: PointAccessorInput, + _rows_per_chunk: int | None = None, + **kwargs: Unpack[ArcLayerKwargs], + ) -> None: + """Construct an ArcLayer from existing Arrow data. + + Keyword Args: + table: An Arrow table from any Arrow-compatible library. This does not need to have a geometry column as geometries are passed separately in `get_source_position` and `get_target_position`. + get_source_position: Source position of each object. + get_target_position: Target position of each object. + + """ + super().__init__( + table=table, + _rows_per_chunk=_rows_per_chunk, + get_source_position=get_source_position, # type: ignore + get_target_position=get_target_position, # type: ignore + **kwargs, + ) + class TextLayer(BaseArrowLayer): """Render text labels at given coordinates.""" diff --git a/lonboard/traits.py b/lonboard/traits.py index f607c955..a0a61862 100644 --- a/lonboard/traits.py +++ b/lonboard/traits.py @@ -30,6 +30,7 @@ from lonboard._constants import EXTENSION_NAME from lonboard._environment import DEFAULT_HEIGHT from lonboard._geoarrow.box_to_polygon import parse_box_encoded_table +from lonboard._geoarrow.ops.coord_layout import convert_struct_column_to_interleaved from lonboard._serialization import ( ACCESSOR_SERIALIZATION, TABLE_SERIALIZATION, @@ -531,17 +532,19 @@ class PointAccessor(FixedErrorTraitType): Various input is allowed: - - A numpy `ndarray` with two dimensions and data type [`np.uint8`][numpy.uint8]. The - size of the second dimension must be `2` or `3`, and will correspond to either XY - or XYZ positions. - - A pyarrow [`FixedSizeListArray`][pyarrow.FixedSizeListArray] or - [`ChunkedArray`][pyarrow.ChunkedArray] containing `FixedSizeListArray`s. The inner - size of the fixed size list must be `2` or `3` and its child must be of floating - point type. + - A numpy `ndarray` with two dimensions and data type [`np.float64`][numpy.float64]. + The size of the second dimension must be `2` or `3`, and will correspond to either + XY or XYZ positions. + - A pyarrow [`FixedSizeListArray`][pyarrow.FixedSizeListArray], + [`StructArray`][pyarrow.StructArray] or [`ChunkedArray`][pyarrow.ChunkedArray] + containing `FixedSizeListArray`s or `StructArray`s. The inner size of the fixed + size list must be `2` or `3` and its child must be of floating point type. + - Any Arrow fixed size list or struct array from a library that implements the [Arrow PyCapsule + Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html). """ default_value = (0, 0, 0) - info_text = "a numpy ndarray or arrow FixedSizeList representing a point array" + info_text = "a numpy ndarray, arrow FixedSizeList, or arrow Struct representing a point array" def __init__( self: TraitType, @@ -582,6 +585,7 @@ def validate( self.error(obj, value) assert isinstance(value, ChunkedArray) + _, value = convert_struct_column_to_interleaved(field=value.field, column=value) if not DataType.is_fixed_size_list(value.type): self.error(obj, value, info="Point arrow array to be a FixedSizeList") diff --git a/lonboard/types/layer.py b/lonboard/types/layer.py index 8b73db8d..17e4d503 100644 --- a/lonboard/types/layer.py +++ b/lonboard/types/layer.py @@ -59,6 +59,13 @@ ArrowArrayExportable, ArrowStreamExportable, ] +PointAccessorInput = Union[ + NDArray[np.float64], + "pa.FixedSizeListArray", + "pa.ChunkedArray", + ArrowArrayExportable, + ArrowStreamExportable, +] TextAccessorInput = Union[ str, NDArray[np.str_], @@ -79,6 +86,20 @@ class BaseLayerKwargs(TypedDict, total=False): auto_highlight: bool +class ArcLayerKwargs(BaseLayerKwargs, total=False): + great_circle: bool + num_segments: int + width_units: Units + width_scale: IntFloat + width_min_pixels: IntFloat + width_max_pixels: IntFloat + get_source_color: ColorAccessorInput + get_target_color: ColorAccessorInput + get_width: FloatAccessorInput + get_height: FloatAccessorInput + get_tilt: FloatAccessorInput + + class BitmapLayerKwargs(BaseLayerKwargs, total=False): image: str bounds: ( diff --git a/tests/layers/__init__.py b/tests/layers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/layers/test_arc_layer.py b/tests/layers/test_arc_layer.py new file mode 100644 index 00000000..45260b61 --- /dev/null +++ b/tests/layers/test_arc_layer.py @@ -0,0 +1,43 @@ +import numpy as np +import pyarrow as pa +from arro3.core import Table +from geoarrow.rust.core import point, points + +from lonboard import Map +from lonboard.experimental import ArcLayer + + +def test_arc_layer_geoarrow_interleaved(): + start_coords = np.array([[1, 4], [2, 5], [3, 6]], dtype=np.float64) + end_coords = np.array([[7, 4], [8, 5], [9, 6]], dtype=np.float64) + start_array = points(start_coords).cast(point("XY", coord_type="interleaved")) + end_array = points(end_coords).cast(point("XY", coord_type="interleaved")) + + string_col = pa.array(["a", "b", "c"], type=pa.string()) + table = Table.from_arrays([string_col], names=["string"]) + + layer = ArcLayer( + table=table, + get_source_position=start_array, + get_target_position=end_array, + ) + m = Map(layer) + assert isinstance(m.layers[0], ArcLayer) + + +def test_arc_layer_geoarrow_separated(): + start_coords = np.array([[1, 4], [2, 5], [3, 6]], dtype=np.float64) + end_coords = np.array([[7, 4], [8, 5], [9, 6]], dtype=np.float64) + start_array = points(start_coords).cast(point("XY", coord_type="separated")) + end_array = points(end_coords).cast(point("XY", coord_type="separated")) + + string_col = pa.array(["a", "b", "c"], type=pa.string()) + table = Table.from_arrays([string_col], names=["string"]) + + layer = ArcLayer( + table=table, + get_source_position=start_array, + get_target_position=end_array, + ) + m = Map(layer) + assert isinstance(m.layers[0], ArcLayer)