11from pathlib import Path
22import json
33
4- # Currently requires a special branch of pyarrow with extra GeoArrow features
5- # https://github.com/apache/arrow/compare/main...paleolimbot:arrow:parquet-geo-write-files-from-geoarrow
4+ # Currently requires nightly pyarrow
65import pyarrow as pa
76from pyarrow import parquet
87import geoarrow .pyarrow as ga
9- from geoarrow .pyarrow import io
108
119here = Path (__file__ ).parent
1210
@@ -26,45 +24,22 @@ def list_wkb_files():
2624 return wkb_files
2725
2826
29- def convert_arrow_wkb_to_parquet (
30- src , dst , compression , write_geoparquet_metadata = False
31- ):
32- # Calculate the basic GeoParquet metadata to add to the file for readers that
33- # don't support this type (but can return the storage of an unknown logical type)
34- with pa .ipc .open_stream (src ) as reader :
35- schema = reader .schema
36-
37- if write_geoparquet_metadata :
38- columns = io ._geoparquet_columns_from_schema (
39- schema ,
40- geometry_columns = ["geometry" ],
41- primary_geometry_column = ["geometry" ],
42- add_geometry_types = False ,
43- )
44- geo_metadata = {
45- "version" : "1.0.0" ,
46- "primary_column" : "geometry" ,
47- "columns" : columns ,
48- }
49-
27+ def convert_arrow_wkb_to_parquet (src , dst , compression ):
5028 # Maintain chunking from IPC into Parquet so that the statistics
5129 # are theoretically the same.
5230 with (
5331 pa .ipc .open_stream (src ) as reader ,
5432 parquet .ParquetWriter (
5533 dst ,
5634 reader .schema ,
57- store_schema = write_geoparquet_metadata ,
35+ store_schema = False ,
5836 compression = compression ,
5937 ) as writer ,
6038 ):
6139 print (f"Reading { src } )
6240 for batch in reader :
6341 writer .write_batch (batch )
6442
65- if write_geoparquet_metadata :
66- writer .add_key_value_metadata ({"geo" : json .dumps (geo_metadata )})
67-
6843 print (f"Wrote { dst } )
6944
7045
@@ -87,48 +62,18 @@ def check_parquet_file(src, dst):
8762 return True
8863
8964
90- def generate_parquet_testing_files (wkb_files , parquet_testing_path ):
91- successful_checks = 0
92- written_files = 0
93- for path in wkb_files :
94- # Skip big files + one CRS example that includes a non-PROJJSON value
95- # on purpose (allowed in GeoArrow), which is rightly rejected
96- # by Parquet
97- name = path .name .replace ("_wkb.arrows" , "" )
98- if (
99- "microsoft-buildings" in name
100- or ("ns-water" in name and name != "ns-water_water-point" )
101- or "wkt2" in name
102- ):
103- print (f"Skipping { name } )
104- continue
105-
106- dst = parquet_testing_path / f"{ name }
107- convert_arrow_wkb_to_parquet (path , dst , compression = "none" )
108- written_files += 1
109- successful_checks += check_parquet_file (path , dst )
110-
111- if successful_checks != written_files :
112- raise ValueError ("Some checks failed when generating testing files" )
113-
114-
11565def generate_geoarrow_data_parquet_files (wkb_files ):
11666 successful_checks = 0
11767 written_files = 0
11868 for path in wkb_files :
11969 name = path .name .replace ("_wkb.arrows" , "" )
120- if "wkt2" in name :
121- print (f"Skipping { name } )
122- continue
12370 if name .startswith ("ns-water" ) or name .startswith ("microsoft" ):
12471 compression = "zstd"
12572 else :
12673 compression = "none"
12774
12875 dst = path .parent / f"{ name }
129- convert_arrow_wkb_to_parquet (
130- path , dst , compression = compression , write_geoparquet_metadata = True
131- )
76+ convert_arrow_wkb_to_parquet (path , dst , compression = compression )
13277 written_files += 1
13378 successful_checks += check_parquet_file (path , dst )
13479
@@ -139,5 +84,4 @@ def generate_geoarrow_data_parquet_files(wkb_files):
13984if __name__ == "__main__" :
14085 parquet_testing_path = here .parent / "parquet-testing" / "data" / "geospatial"
14186 wkb_files = list_wkb_files ()
142- generate_parquet_testing_files (wkb_files , parquet_testing_path )
14387 generate_geoarrow_data_parquet_files (wkb_files )
0 commit comments