Repository: alxmrs/dask-ee (language: python)

Issue:
Implement MVP
As discussed in google/Xee#153, I'll put together a PR with the core functionality of loading Earth Engine data into a Dask dataframe via a top-level `read_ee` function.

We settled on `dask-geopandas` as an optional dependency to enable geospatial processing on deferred dataframes, but I'm leaning towards omitting that from the MVP until geopandas/dask-geopandas#285 is merged. The alternative would be to pin to an older Dask before the incompatible query planning system was added or disable query planning in the [Dask configuration](https://dask-geopandas.readthedocs.io/en/stable/docs/reference/api/dask_geopandas.from_dask_dataframe.html#dask_geopandas.from_dask_dataframe). Given that there's a fix in progress, I'm in favor of just postponing that feature until it's released and then pinning to that.

That doesn't stop us from optionally using GeoDataFrames for the underlying chunks in the Dask dataframe, and if a user *needs* to run geospatial operations without computing they can always set their own Dask config and wrap the `read_ee` output in [`from_dask_dataframe`](https://dask-geopandas.readthedocs.io/en/stable/docs/reference/api/dask_geopandas.from_dask_dataframe.html#dask_geopandas.from_dask_dataframe).

Let me know if you have strong feelings on that @alxmrs - otherwise I'll move ahead with that plan.


Discussion / hints:
That plan sounds correct. Thanks for looking into the limitation, I wasn’t aware. 

Fix patch:
diff --git a/dask_ee/__init__.py b/dask_ee/__init__.py
index f2166d7..677febd 100644
--- a/dask_ee/__init__.py
+++ b/dask_ee/__init__.py
@@ -1,1 +1,1 @@
-# TODO(alxmrs): https://docs.google.com/document/d/1Ltl6XrZ_uGD2J7OW1roUsxhpFxx5QvNeInjuONIkmL8/edit#heading=h.uj7plawe7x7e
+from .read import read_ee
diff --git a/dask_ee/read.py b/dask_ee/read.py
new file mode 100644
index 0000000..fd5c21f
--- /dev/null
+++ b/dask_ee/read.py
@@ -0,0 +1,70 @@
+# Special thanks to @aazuspan for help with the implementation
+import typing as t
+
+import dask.dataframe as dd
+import ee
+import numpy as np
+import pandas as pd
+
+# Order is in appearance of types in the EE documentation. This looks alphabetical.
+_BUILTIN_DTYPES = {
+    'Byte': np.uint8,
+    'Double': np.float64,
+    'Float': np.float32,
+    'Int': np.int32,
+    'Int16': np.int16,
+    'Int32': np.int32,
+    'Int64': np.int64,
+    'Int8': np.int8,
+    'Json': np.object_,  # added to handle GeoJSON columns.
+    'Long': np.int64,
+    'Short': np.int16,
+    'Uint16': np.uint16,
+    'Uint32': np.uint32,
+    'Uint8': np.uint8,
+    'String': np.str_,
+}
+
+
+# TODO(#4): Support 'auto' chunks, where we calculate the maximum allowed page size given the number of
+#  bytes in each row.
+def read_ee(
+    fc: ee.FeatureCollection, io_chunks: t.Union[int, t.Literal['auto']] = 5_000
+) -> dd.DataFrame:
+
+  if io_chunks == 'auto':
+    raise NotImplementedError('Auto `io_chunks` are not implemented yet!')
+
+  # Make all the getInfo() calls at once, up front.
+  fc_size, all_info = ee.List([fc.size(), fc.limit(0)]).getInfo()
+
+  columns = {'geo': 'Json'}
+  columns.update(all_info['columns'])
+  del columns['system:index']
+
+  divisions = tuple(range(0, fc_size, io_chunks))
+
+  # TODO(#5): Compare `toList()` to other range operations, like getting all index IDs via `getInfo()`.
+  pages = [ee.FeatureCollection(fc.toList(io_chunks, i)) for i in divisions]
+  # Get the remainder, if it exists. `io_chunks` are not likely to evenly partition the data.
+  d, r = divmod(fc_size, io_chunks)
+  if r != 0:
+    pages.append(ee.FeatureCollection(fc.toList(r, d)))
+    divisions += (fc_size,)
+
+  def to_df(page: ee.FeatureCollection) -> pd.DataFrame:
+    return ee.data.computeFeatures(
+        {
+            'expression': page,
+            'fileFormat': 'PANDAS_DATAFRAME',
+        }
+    )
+
+  meta = {k: _BUILTIN_DTYPES[v] for k, v in columns.items()}
+
+  return dd.from_map(
+      to_df,
+      pages,
+      meta=meta,
+      divisions=divisions,
+  )
diff --git a/dask_ee/write.py b/dask_ee/write.py
new file mode 100644
index 0000000..afd98ca
--- /dev/null
+++ b/dask_ee/write.py
@@ -0,0 +1,5 @@
+# TODO(alxmrs): This is open for design and implementation.
+
+
+def to_ee(ddf, *args, **kwargs):
+  raise NotImplementedError('This has not yet been designed.')
diff --git a/pyproject.toml b/pyproject.toml
index 7cf5ce7..118fa77 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,6 +25,7 @@ classifiers = [
 ]
 dependencies = [
   "earthengine-api>=0.1.374",
+  "pandas",
   "dask",
 ]
 


Test patch:
diff --git a/dask_ee/default_test.py b/dask_ee/default_test.py
deleted file mode 100644
index 722e4f9..0000000
--- a/dask_ee/default_test.py
+++ /dev/null
@@ -1,14 +0,0 @@
-import unittest
-
-
-class ImportTest(unittest.TestCase):
-
-  def test_can_import_module(self):
-    try:
-      import dask_ee
-    except ModuleNotFoundError:
-      self.fail('Cannot import `dask_ee`.')
-
-
-if __name__ == '__main__':
-  unittest.main()
diff --git a/dask_ee/read_integrationtest.py b/dask_ee/read_integrationtest.py
new file mode 100644
index 0000000..ca794f6
--- /dev/null
+++ b/dask_ee/read_integrationtest.py
@@ -0,0 +1,50 @@
+"""Integration tests with Google Earth Engine.
+
+Before running, please authenticate:
+```
+earthengine authenticate
+```
+"""
+
+import cProfile
+import pstats
+import unittest
+
+import dask.dataframe as dd
+import ee
+
+import dask_ee
+
+
+class ReadIntegrationTests(unittest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    ee.Initialize()
+
+  def test_reads_dask_dataframe(self):
+    fc = ee.FeatureCollection("WRI/GPPD/power_plants")
+    ddf = dask_ee.read_ee(fc)
+
+    head = ddf.head()
+    columns = ddf.columns
+
+    self.assertIsNotNone(ddf)
+    self.assertIsNotNone(head)
+    self.assertIsInstance(ddf, dd.DataFrame)
+    self.assertEqual(ddf.compute().shape, (28_664, 23))
+
+    print(columns)
+    print(head)
+
+  def test_prof__read_ee(self):
+    fc = ee.FeatureCollection("WRI/GPPD/power_plants")
+    with cProfile.Profile() as pr:
+      _ = dask_ee.read_ee(fc)
+
+      # Modified version of `pr.print_stats()`.
+      pstats.Stats(pr).sort_stats("cumtime").print_stats()
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/dask_ee/read_test.py b/dask_ee/read_test.py
new file mode 100644
index 0000000..1e4b17f
--- /dev/null
+++ b/dask_ee/read_test.py
@@ -0,0 +1,14 @@
+import unittest
+
+
+class ReadFeatureCollections(unittest.TestCase):
+
+  def test_can_import_read_op(self):
+    try:
+      from dask_ee import read_ee
+    except ModuleNotFoundError:
+      self.fail('Cannot import `read_ee` function.')
+
+
+if __name__ == '__main__':
+  unittest.main()