feat: Allow users to configure the internal thread pool (#11)

xee/ext.py

@@ -89,6 +89,12 @@ class EarthEngineStore(common.AbstractDataStore):
 
   DEFAULT_MASK_VALUE = np.iinfo(np.int32).max
 
+  def __init__(self, executor_kwargs=None):
+    if executor_kwargs is None:
+      self.executor_kwargs = {}
+    else:
+      self.executor_kwargs = executor_kwargs
+
   @classmethod
   def open(
       cls,
@@ -634,7 +640,7 @@ def reduce_bands(x, acc):
     return target_image
 
   def _raw_indexing_method(
-      self, key: tuple[Union[int, slice], ...]
+      self, key: tuple[Union[int, slice], ...] 
   ) -> np.typing.ArrayLike:
     key, squeeze_axes = self._key_to_slices(key)
 
@@ -682,8 +688,9 @@ def _raw_indexing_method(
         for _ in range(shape[0])
     ]
 
-    # TODO(#11): Allow users to configure this via kwargs.
-    with concurrent.futures.ThreadPoolExecutor() as pool:
+
+    # Pass executor_kwargs to ThreadPoolExecutor
+    with concurrent.futures.ThreadPoolExecutor(**self.executor_kwargs) as pool:
       for (i, j, k), arr in pool.map(
           self._make_tile, self._tile_indexes(key[0], bbox)
       ):
@@ -772,6 +779,7 @@ def open_dataset(
       primary_dim_name: Optional[str] = None,
       primary_dim_property: Optional[str] = None,
       ee_mask_value: Optional[float] = None,
+      executor_kwargs: Optional[dict] = None,
   ) -> xarray.Dataset:
     """Open an Earth Engine ImageCollection as an Xarray Dataset.
 
@@ -859,6 +867,7 @@ def open_dataset(
         primary_dim_name=primary_dim_name,
         primary_dim_property=primary_dim_property,
         mask_value=ee_mask_value,
+        executor_kwargs=executor_kwargs,
     )
 
     store_entrypoint = backends_store.StoreBackendEntrypoint()

xee/ext_integration_test.py

@@ -358,6 +358,18 @@ def test_data_sanity_check(self):
     self.assertNotEqual(temperature_2m.min(), 0.0)
     self.assertNotEqual(temperature_2m.max(), 0.0)
 
+  def test_open_dataset_with_executor_kwargs(self):
+    executor_kwargs = {'max_workers': 2}
+    ds = self.entry.open_dataset(
+        'ee://LANDSAT/LC08/C01/T1',
+        drop_variables=tuple(f'B{i}' for i in range(3, 12)),
+        scale=25.0,  
+        n_images=3,
+        executor_kwargs=executor_kwargs,
+    )
+
+    self.assertEqual(ds.thread_pool.max_workers, executor_kwargs['max_workers'])
+
 
 if __name__ == '__main__':
   absltest.main()