def write_geospatial()

in data/geospatial/geospatial-gen.py [0:0]

• 20 lines of code
• 2 McCabe index (conditional complexity)

def write_geospatial():
    with open(HERE / "geospatial.yaml") as f:
        examples = yaml.safe_load(f)

    schema = pa.schema({"group": pa.utf8(), "wkt": pa.utf8(), "geometry": ga.wkb()})

    with parquet.ParquetWriter(
        HERE / "geospatial.parquet",
        schema,
        store_schema=False,
        compression="none",
    ) as writer:
        for group_name, geometries_wkt in examples.items():
            # Unfortunately we can't use Shapely to generate the test WKB
            # because of https://github.com/libgeos/geos/issues/888, so we use
            # geoarrow.pyarrow.as_wkb() instead.
            # geometries = shapely.from_wkt(geometries_wkt)
            # wkbs = shapely.to_wkb(geometries, flavor="iso")
            # wkb_array = ga.wkb().wrap_array(pa.array(wkbs, pa.binary()))
            wkt_array = pa.array(geometries_wkt, pa.utf8())

            batch = pa.record_batch(
                {
                    "group": [group_name] * len(geometries_wkt),
                    "wkt": wkt_array,
                    "geometry": ga.as_wkb(wkt_array),
                }
            )

            writer.write_batch(batch)