IO

`scan_spatial(path_or_buffer, layer=None, encoding=None, bbox=None, mask=None)`

Scans a data source supported by pyogrio or a geoparquet file to produce a polars LazyFrame.

Note

Although geoparquet is supported, this implementation, in its current state, leaves a lot to be desired.

Parameters:

Name	Type	Description	Default
`path_or_buffer`	`str \| Path \| BytesIO`	A dataset path or URI, raw buffer, or file-like object with a read method.	required
`layer`	`Optional[str \| int]`	If an integer is provided, it corresponds to the index of the layer with the data source. If a string is provided, it must match the name of the layer in the data source. Defaults to first layer in data source.	`None`
`encoding`	`Optional[str]`	If present, will be used as the encoding for reading string values from the data source. By default will automatically try to detect the native encoding and decode to UTF-8.	`None`
`bbox`	`Optional[tuple[float, float, float, float]]`	If present, will be used to filter records whose geometry intersects this box. This must be in the same CRS as the dataset. If GEOS is present and used by GDAL, only geometries that intersect this bbox will be returned; if GEOS is not available or not used by GDAL, all geometries with bounding boxes that intersect this bbox will be returned. Cannot be combined with mask keyword. Tuple should be in the format of (xmin, ymin, xmax, ymax).	`None`
`mask`	`Optional[Polygon]`	If present, will be used to filter records whose geometry intersects this geometry. This must be in the same CRS as the dataset. If GEOS is present and used by GDAL, only geometries that intersect this geometry will be returned; if GEOS is not available or not used by GDAL, all geometries with bounding boxes that intersect the bounding box of this geometry will be returned. Requires Shapely >= 2.0. Cannot be combined with bbox keyword.	`None`

Examples:

Scanning a layer from a geopackage:

>>> my_geopackage = r"c:\data\hiking_club.gpkg"
>>> lf = scan_spatial(my_geopackage, layer="hike")
>>> lf
naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
PYTHON SCAN []
PROJECT */4 COLUMNS

Scanning a shapefile:

>>> from spatial_polars import scan_spatial
>>> my_shapefile = r"c:\data\roads.shp"
>>> lf = scan_spatial(my_shapefile)
>>> lf
naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
PYTHON SCAN []
PROJECT */11 COLUMNS

Scanning a shapefile from within a zipped directory:

>>> zipped_shapefiles = r"C:\data\illinois-latest-free.shp.zip"
>>> lf = scan_spatial(zipped_shapefiles, layer="gis_osm_roads_free_1")
>>> lf
naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
PYTHON SCAN []
PROJECT */11 COLUMNS

Source code in src\spatial_polars\io.py

def scan_spatial(
    path_or_buffer: str | Path | BytesIO,
    layer: Optional[str | int] = None,
    encoding: Optional[str] = None,
    bbox: Optional[tuple[float, float, float, float]] = None,
    mask: Optional[shapely.Polygon] = None,
) -> pl.LazyFrame:
    r"""
    Scans a data source [supported by pyogrio](https://pyogrio.readthedocs.io/en/stable/supported_formats.html) or a geoparquet file to produce a polars LazyFrame.

    Note
    ----
    Although geoparquet is supported, this implementation, in its current state, leaves a lot to be desired.

    Parameters
    ----------

    path_or_buffer
        A dataset path or URI, raw buffer, or file-like object with a read method.

    layer
        If an integer is provided, it corresponds to the index of the layer with the
        data source. If a string is provided, it must match the name of the layer in
        the data source. Defaults to first layer in data source.

    encoding
        If present, will be used as the encoding for reading string values from the
        data source. By default will automatically try to detect the native encoding
        and decode to UTF-8.

    bbox
        If present, will be used to filter records whose geometry intersects this
        box. This must be in the same CRS as the dataset. If GEOS is present and
        used by GDAL, only geometries that intersect this bbox will be returned;
        if GEOS is not available or not used by GDAL, all geometries with bounding
        boxes that intersect this bbox will be returned. Cannot be combined with mask
        keyword.  Tuple should be in the format of (xmin, ymin, xmax, ymax).

    mask
        If present, will be used to filter records whose geometry intersects this
        geometry. This must be in the same CRS as the dataset. If GEOS is present
        and used by GDAL, only geometries that intersect this geometry will be
        returned; if GEOS is not available or not used by GDAL, all geometries with
        bounding boxes that intersect the bounding box of this geometry will be
        returned. Requires Shapely >= 2.0. Cannot be combined with bbox keyword.

    Examples
    --------
    **Scanning a layer from a geopackage:**

    >>> my_geopackage = r"c:\data\hiking_club.gpkg"
    >>> lf = scan_spatial(my_geopackage, layer="hike")
    >>> lf
    naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
    PYTHON SCAN []
    PROJECT */4 COLUMNS

    **Scanning a shapefile:**

    >>> from spatial_polars import scan_spatial
    >>> my_shapefile = r"c:\data\roads.shp"
    >>> lf = scan_spatial(my_shapefile)
    >>> lf
    naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
    PYTHON SCAN []
    PROJECT */11 COLUMNS

    **Scanning a shapefile from within a zipped directory:**

    >>> zipped_shapefiles = r"C:\data\illinois-latest-free.shp.zip"
    >>> lf = scan_spatial(zipped_shapefiles, layer="gis_osm_roads_free_1")
    >>> lf
    naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
    PYTHON SCAN []
    PROJECT */11 COLUMNS
    """
    if isinstance(path_or_buffer, (str, Path)) and str(path_or_buffer).endswith(
        ".parquet"
    ):
        # TODO look into libgdal-arrow-parquet from conda forge
        # https://pyogrio.readthedocs.io/en/latest/install.html#conda-forge
        schema = pl.scan_parquet(path_or_buffer).collect_schema()
        if schema.get("geometry") is not None:
            schema["geometry"] = spatial_series_dtype
        if bbox is not None:
            mask = shapely.Polygon(shapely.box(*bbox))

        def source_generator(
            with_columns: list[str] | None,
            predicate: pl.Expr | None,
            n_rows: int | None,
            batch_size: int | None,
        ) -> Iterator[pl.DataFrame]:
            """
            Generator function that creates the source.
            This function will be registered as IO source.
            """
            if mask is not None:
                if "geometry" not in with_columns:
                    with_columns.append("geometry")

            tab = _pq.read_table(path_or_buffer)
            tab_metadata = tab.schema.metadata if tab.schema.metadata else {}
            if b"geo" in tab_metadata:
                geo_meta = json.loads(tab_metadata[b"geo"])
            else:
                geo_meta = {}
            geom_col = geo_meta["primary_column"]
            crs_wkt = pyproj.CRS(geo_meta["columns"][geom_col]["crs"]).to_wkt()

            if batch_size is None:
                batch_size = 10000

            if with_columns is None:
                read_geometry = True
            elif "geometry" in with_columns:
                read_geometry = True
            else:
                read_geometry = False

            lf = pl.scan_parquet(path_or_buffer)

            if with_columns is not None:
                lf = lf.select(with_columns)

            if predicate is not None:
                lf = lf.filter(predicate)

            previous_max = 0
            while n_rows is None or n_rows > 0:
                batch = lf.slice(previous_max, previous_max + batch_size).collect()
                if batch.height is None or batch.height == 0:
                    break
                if n_rows is not None and n_rows <= 0:
                    break

                if read_geometry:
                    # get the geometries from the batch
                    geometries = batch[0:n_rows][geom_col]
                    shapely_goms = shapely.from_wkb(geometries)
                    geometries = shapely.to_wkb(shapely_goms)

                    # create the dataframe with the non geometry columns
                    # then add struct column with the WKB geometries/CRS
                    df = pl.DataFrame(batch[0:n_rows].drop(geom_col)).with_columns(
                        pl.struct(
                            pl.Series("wkb_geometry", geometries, dtype=pl.Binary),
                            pl.lit(crs_wkt, dtype=pl.Categorical).alias("crs"),
                        ).alias("geometry")
                    )
                else:
                    df = pl.DataFrame(batch[0:n_rows])
                previous_max += df.height

                if n_rows is not None:
                    n_rows -= df.height

                if predicate is not None:
                    df = df.filter(predicate)

                if mask is not None:
                    df = df.filter(pl.col("geometry").spatial.intersects(mask))
                if mask is not None:
                    if "geometry" not in with_columns:
                        df = df.drop("geometry")

                yield df

    else:
        # not geoparquet
        layer_info = pyogrio.read_info(path_or_buffer, layer=layer, encoding=encoding)
        schema = dict(
            zip(
                layer_info["fields"],
                [PYOGRIO_POLARS_DTYPES[dt] for dt in layer_info["dtypes"]],
            )
        )
        if layer_info.get("geometry_type"):
            schema["geometry"] = spatial_series_dtype

        def source_generator(
            with_columns: list[str] | None,
            predicate: pl.Expr | None,
            n_rows: int | None,
            batch_size: int | None,
        ) -> Iterator[pl.DataFrame]:
            """
            Generator function that creates the source.
            This function will be registered as IO source.
            """

            if batch_size is None:
                batch_size = 100

            if with_columns is None:
                read_geometry = True
            elif "geometry" in with_columns:
                read_geometry = True
                with_columns.remove("geometry")
            else:
                read_geometry = False

            with pyogrio.open_arrow(
                path_or_buffer,
                layer=layer,
                encoding=encoding,
                columns=with_columns,
                read_geometry=read_geometry,
                force_2d=False,
                bbox=bbox,
                mask=mask,
                batch_size=batch_size,
                use_pyarrow=True,
            ) as source:
                meta, reader = source

                # extract the crs from the metadata
                crs_wkt = pyproj.CRS(meta["crs"]).to_wkt()

                geom_col = meta["geometry_name"] or "wkb_geometry"

                while n_rows is None or n_rows > 0:
                    for batch in reader:
                        if n_rows is not None and n_rows <= 0:
                            break

                        if read_geometry:
                            # get the geometries from the batch
                            geometries = batch[geom_col][0:n_rows]
                            shapely_goms = shapely.from_wkb(geometries)
                            geometries = shapely.to_wkb(shapely_goms)
                            # create the dataframe with the non geometry columns
                            # then add struct column with the WKB geometries/CRS
                            df = pl.DataFrame(
                                batch[0:n_rows].drop_columns(geom_col)
                            ).with_columns(
                                pl.struct(
                                    pl.Series(
                                        "wkb_geometry", geometries, dtype=pl.Binary
                                    ),
                                    pl.lit(crs_wkt, dtype=pl.Categorical).alias("crs"),
                                ).alias("geometry")
                            )
                        else:
                            df = pl.DataFrame(batch[0:n_rows])

                        if n_rows is not None:
                            n_rows -= df.height

                        if predicate is not None:
                            df = df.filter(predicate)

                        yield df
                    if n_rows is None or n_rows <= 0:
                        break

    return register_io_source(io_source=source_generator, schema=schema)

`read_spatial(path_or_buffer, layer=None, encoding=None, bbox=None, mask=None)`

Reads a spatial data source supported by pyogrio to produce a polars DataFrame.