Search code, repositories, users, issues, pull requests...

Copy link

@pvrraju pvrraju commented Sep 5, 2025

Hi @alxmrs,

Added lazy loading (issue #44):

lazy_load=True option

Defers metadata RPCs → faster large dataset opens

Backward compatible + tested

can you review my PR?

Mimic xr.open_zarr(..., chunks=None) to provide a way for lazy, fast open_dataset() calls. #44


 Implement lazy loading to defer metadata RPCs until data access time ( g...

660100d

...oogle#44)

@pvrraju pvrraju mentioned this pull request

Sep 8, 2025

Open

alxmrs

alxmrs reviewed

xee/ext.py Outdated

self._info_cache = {}

# Perform minimal RPCs if lazy loading is enabled

if getattr(self, 'lazy_load', False):

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of this eager getattr, I think we should just pass lazy_load as an argument to this function.

alxmrs

alxmrs reviewed

xee/ext.py Outdated

return self._info_cache

# Full metadata loading if not lazy

if not self._info_cache or len(self._info_cache) < 5: # Check if we have full metadata

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like this magic number. Is there a more principled way to check if the cache needs updating?

alxmrs

alxmrs reviewed

xee/ext.py Outdated

Comment on lines 383 to 388

columns = ['system:id', self.primary_dim_property]

properties = (

self.image_collection.reduceColumns(

ee.Reducer.toList().repeat(len(columns)), columns

).get('list')

).getInfo()

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Idea: maybe we can break this out into a function and call it here and in get_info as needed.

alxmrs

alxmrs reviewed

xee/ext_integration_test.py

Comment on lines +583 to +586

# Verify that lazy opening is faster than regular opening

self.assertLess(lazy_open_time, regular_open_time,

f"Lazy loading ({lazy_open_time:.2f}s) should be faster than regular loading ({regular_open_time:.2f}s)")

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I'm really happy to see this test.

alxmrs

alxmrs reviewed

xee/ext_integration_test.py Outdated

'1992年10月05日', '1992年10月06日') # Using a smaller date range for the test

# Open dataset with lazy loading

start_time = time.time()

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's use time.perf_counter() to capture time instead (and at the places below).

alxmrs

alxmrs reviewed

xee/ext_integration_test.py

Comment on lines +588 to +599

# Verify that both datasets have the same structure

self.assertEqual(lazy_ds.dims, regular_ds.dims)

self.assertEqual(list(lazy_ds.data_vars), list(regular_ds.data_vars))

# Access data and verify it's the same

var_name = list(lazy_ds.data_vars)[0]

lazy_data = lazy_ds[var_name].isel(time=0).values

regular_data = regular_ds[var_name].isel(time=0).values

# Both should have same shape and data should not be all zeros or NaNs

self.assertEqual(lazy_data.shape, regular_data.shape)

self.assertTrue(np.allclose(lazy_data, regular_data, equal_nan=True))

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks this is good to have in the test.

alxmrs

alxmrs reviewed

xee/ext.py

Comment on lines +297 to +298

if not hasattr(self, '_info_cache'):

self._info_cache = {}

Copy link

Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm concerned that we have two levels of caching: this private dictionary and the functools.cached_property level. Is there any way we can use one system instead of both? I anticipate cache invalidation problems down the line.

One way we could clean this system up is by breaking out what info we get into multiple functions (that are also cached): e.g. a helper fn gets the first line essential stuff below. Since it's cached, we'd need less if...else logic to manage state; that would be happen in Python's memoizer decorator -- we'd make use of the cache just by calling the helper functions.

WDYT about this approach?

alxmrs

alxmrs suggested changes

Copy link

Contributor

@alxmrs alxmrs left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Provided specific and high-level feedback.

In addition, I have a general warning: I don't have much control over whether PRs like this will actually get merged into Xee; patches like this require secondary review to ensure that the code is compatible with Google's internal use of Xee, and I'm not sure if there is capacity on the relevant teams to do these types of reviews right now.


 Address reviewer feedback for lazy loading implementation

01abbbd