Pass storage options, add some typing, minor clean up (#78)

* build sindex takes args, returns self, cleanup

* Pass storage and engine options, add some typing

* Pass kwargs to _load_parquet_pandas_metadata

* Remove unused import in tests

* Fix argument order for pack_partitions_to_parquet

* Add overwrite option, minor clean up

* Add __all__, minor cleanup

Author: brl0

spatialpandas/__init__.py

@@ -1,12 +1,13 @@
 import param as _param
 
-from . import geometry, spatialindex, tools  # noqa
-from .geodataframe import GeoDataFrame  # noqa
-from .geoseries import GeoSeries  # noqa
-from .tools.sjoin import sjoin  # noqa
+from . import geometry, spatialindex, tools
+from .geodataframe import GeoDataFrame
+from .geoseries import GeoSeries
+from .tools.sjoin import sjoin
 
 try:
     import dask.dataframe  # noqa
+
     # Import to trigger registration of types with Dask
     import spatialpandas.dask  # noqa
 except ImportError:
@@ -19,3 +20,12 @@
         archive_commit="$Format:%h$",
         reponame="spatialpandas",
     ))
+
+__all__ = [
+    "GeoDataFrame",
+    "GeoSeries",
+    "geometry",
+    "sjoin",
+    "spatialindex",
+    "tools",
+]

spatialpandas/dask.py

@@ -83,11 +83,11 @@ def cx(self):
     def cx_partitions(self):
         return _DaskPartitionCoordinateIndexer(self, self.partition_sindex)
 
-    def build_sindex(self):
-        def build_sindex(series):
-            series.build_sindex()
+    def build_sindex(self, **kwargs):
+        def build_sindex(series, **kwargs):
+            series.build_sindex(**kwargs)
             return series
-        return self.map_partitions(build_sindex, meta=self._meta)
+        return self.map_partitions(build_sindex, **kwargs, meta=self._meta)
 
     def intersects_bounds(self, bounds):
         return self.map_partitions(lambda s: s.intersects_bounds(bounds))
@@ -200,8 +200,17 @@ def pack_partitions(self, npartitions=None, p=15, shuffle='tasks'):
         return ddf
 
     def pack_partitions_to_parquet(
-            self, path, filesystem=None, npartitions=None, p=15, compression="snappy",
-            tempdir_format=None, _retry_args=None
+        self,
+        path,
+        filesystem=None,
+        npartitions=None,
+        p=15,
+        compression="snappy",
+        tempdir_format=None,
+        _retry_args=None,
+        storage_options=None,
+        engine_kwargs=None,
+        overwrite=False,
     ):
         """
         Repartition and reorder dataframe spatially along a Hilbert space filling curve
@@ -235,14 +244,14 @@ def pack_partitions_to_parquet(
         from .io.utils import validate_coerce_filesystem
 
         # Get fsspec filesystem object
-        filesystem = validate_coerce_filesystem(path, filesystem)
+        filesystem = validate_coerce_filesystem(path, filesystem, storage_options)
 
         # Decorator for operations that should be retried
         if _retry_args is None:
             _retry_args = dict(
                 wait_exponential_multiplier=100,
                 wait_exponential_max=120000,
-                stop_max_attempt_number=24
+                stop_max_attempt_number=24,
             )
         retryit = retry(**_retry_args)
 
@@ -320,7 +329,8 @@ def move_retry(p1, p2):
 
         # Initialize output partition directory structure
         filesystem.invalidate_cache()
-        rm_retry(path)
+        if overwrite:
+            rm_retry(path)
 
         for out_partition in out_partitions:
             part_dir = os.path.join(path, "part.%d.parquet" % out_partition)
@@ -332,17 +342,24 @@ def move_retry(p1, p2):
         @retryit
         def write_partition(df_part, part_path):
             with filesystem.open(part_path, "wb") as f:
-                df_part.to_parquet(f, compression=compression, index=True)
+                df_part.to_parquet(
+                    f,
+                    compression=compression,
+                    index=True,
+                    **(engine_kwargs or {}),
+                )
 
         def process_partition(df, i):
             subpart_paths = {}
             for out_partition, df_part in df.groupby('_partition'):
                 part_path = os.path.join(
                     tempdir_format.format(partition=out_partition, uuid=dataset_uuid),
-                    'part.%d.parquet' % i
+                    'part.%d.parquet' % i,
                 )
-                df_part = df_part.drop('_partition', axis=1).set_index(
-                    'hilbert_distance', drop=True
+                df_part = (
+                    df_part
+                    .drop('_partition', axis=1)
+                    .set_index('hilbert_distance', drop=True)
                 )
                 write_partition(df_part, part_path)
                 subpart_paths[out_partition] = part_path
@@ -528,11 +545,11 @@ def _propagate_props_to_series(self, new_series):
             new_series._partition_sindex = self._partition_sindex[new_series.name]
         return new_series
 
-    def build_sindex(self):
-        def build_sindex(df):
-            df.build_sindex()
+    def build_sindex(self, **kwargs):
+        def build_sindex(df, **kwargs):
+            df.build_sindex(**kwargs)
             return df
-        return self.map_partitions(build_sindex, meta=self._meta)
+        return self.map_partitions(build_sindex, **kwargs, meta=self._meta)
 
     def persist(self, **kwargs):
         return self._propagate_props_to_dataframe(

spatialpandas/geodataframe.py

@@ -113,8 +113,9 @@ def cx(self):
             self.geometry.array, parent=self
         )
 
-    def build_sindex(self):
-        self.geometry.build_sindex()
+    def build_sindex(self, **kwargs):
+        self.geometry.build_sindex(**kwargs)
+        return self
 
     def _ensure_type(self, obj):
         # Override because a GeoDataFrame operation may result in a regular DataFrame,

spatialpandas/geometry/base.py

@@ -555,9 +555,10 @@ def sindex(self):
             self.build_sindex()
         return self._sindex
 
-    def build_sindex(self):
+    def build_sindex(self, **kwargs):
         if self._sindex is None:
-            self._sindex = HilbertRtree(self.bounds)
+            self._sindex = HilbertRtree(self.bounds, **kwargs)
+        return self
 
     @property
     def cx(self):

spatialpandas/geoseries.py

@@ -76,8 +76,9 @@ def cx(self):
         from .geometry.base import _CoordinateIndexer
         return _CoordinateIndexer(self.array, parent=self)
 
-    def build_sindex(self):
-        self.array.build_sindex()
+    def build_sindex(self, **kwargs):
+        self.array.build_sindex(**kwargs)
+        return self
 
     def intersects_bounds(self, bounds):
         return pd.Series(