Skip to content

Commit

Permalink
Basic reads work without meta.
Browse files Browse the repository at this point in the history
  • Loading branch information
alxmrs committed Jun 23, 2024
1 parent eb01aa6 commit ad2b2bd
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 3 deletions.
8 changes: 5 additions & 3 deletions dask_ee/read.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
'int32': np.int32,
'int64': np.int64,
'int8': np.int8,
'json': dict,
'long': np.int64,
'short': np.int16,
'uint16': np.uint16,
Expand All @@ -39,7 +40,8 @@ def read_ee(
raise NotImplementedError('Auto io_chunks are not implemented yet!')

fc_size, all_info = ee.List([fc.size(), fc]).getInfo()
columns = all_info['columns']
columns = {'geo': 'json'}
columns.update(all_info['columns'])

# TODO(#5): Compare `toList()` to other range operations, like getting all index IDs via `getInfo()`.
pages = [
Expand All @@ -55,12 +57,12 @@ def to_df(page: ee.FeatureCollection) -> pd.DataFrame:
}
)

meta = {k: _BUILTIN_DTYPES[v.lower()] for k, v in columns.items()}
# TODO(alxmrs): Support dask dataframe `meta` via columns.
# meta = {k: _BUILTIN_DTYPES[v.lower()] for k, v in columns.items()}
divisions = tuple(range(0, fc_size, io_chunks))

return dd.from_map(
to_df,
pages,
meta=meta,
divisions=divisions,
)
39 changes: 39 additions & 0 deletions dask_ee/read_integrationtest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
"""Integration tests with Google Earth Engine.
Before running, please authenticate:
```
earthengine authenticate
```
"""
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

print(columns)

self.assertIsNotNone(ddf)
self.assertIsNotNone(head)
self.assertIsInstance(ddf, dd.DataFrame)

print(head)


if __name__ == '__main__':
unittest.main()

0 comments on commit ad2b2bd

Please sign in to comment.