feat: filtering based on multiple variables

[maxrjones]

• Describe the change

This PR adds a where query parameter for filtering rendered data, which is necessary for visualizing TEMPO data in a similar way to worldview.

Testing

• New tests ran with pytest

PR checks

• Standard CI runs automatically on each push.
• To run the CDK synth check, add the run-cdk-checks label to this PR.
• If you push more commits after that run completes, remove and re-add the label to run it again.
• To trigger a dev deployment, add the deploy-dev label. It smoke-tests tiles from the native MUR, virtual MUR, and virtual NLDAS Icechunk stores after deployment.

[hrodmn]

@maxrjones thanks for adding this feature - it is a cool solution to a problem that I have grappled with in the past: how to properly apply a custom mask in a titiler request! I have tried a few things and the only workaround I have found with the current state of things in rio-tiler/titiler is to use the expression parameter to recode the values that I want to mask with numpy.where then provide a custom colormap that excludes the recoded mask values (example notebook). This works for some limited use-cases but is not very robust because the mask happens after warping/resampling so should really just be used with the nearest resampling method.

The where parameter is great because it applies the mask before warping/resampling. I would like to see this feature ported upstream to rio-tiler or titiler someday!

I don't love all of the dask/xarray glue code to keep everything lazy 😎 but I had to do some similar gymnastics in titiler-cmr's custom expression parser for the XarrayBackend so I get it. I don't have a great understanding of how/when to bring dask in to maintain laziness so this was helpful to see.

I will do a line-by-line review but just want to provide my initial thoughts.

[hrodmn]

did you hit Lambda package size limitations that spurred these changes?

[maxrjones]

yes, just adding dask put it ~5MB above the limit 😞

[chuckwondo]

Hi @maxrjones. Sorry for the delay, but I'm reviewing now. I have some suggestions, so give me some time to construct some clear comments.

[chuckwondo]

@maxrjones, I put a few suggestions inline, but my main suggestion is too complex to describe across inline PR comments, so I created a gist: https://gist.github.com/chuckwondo/b151f7f92d54cf01c62da2939238d62b

[maxrjones]

@maxrjones, I put a few suggestions inline, but my main suggestion is too complex to describe across inline PR comments, so I created a gist: gist.github.com/chuckwondo/b151f7f92d54cf01c62da2939238d62b

thanks, @chuckwondo! I took most of your suggestion in 9e26c15 (this PR), with the exception of Make __attrs_post_init__ wrap the opener. I thought it was simpler calling chunk_for_lazy_ops in _apply_where, and the performance difference should be negligable.

[maxrjones]

I don't love all of the dask/xarray glue code to keep everything lazy 😎 but I had to do some similar gymnastics in titiler-cmr's custom expression parser for the XarrayBackend so I get it. I don't have a great understanding of how/when to bring dask in to maintain laziness so this was helpful to see.

I definitely agree with you here! I think in the future we should see about using native Zarr laziness and leaving xarray/dask out of the picture for regularly gridded Icechunk stores. That said, I'd rather have an MVP out and refactor later, as long as we're happy with the API design.

[chuckwondo]

@maxrjones, I put a few suggestions inline, but my main suggestion is too complex to describe across inline PR comments, so I created a gist: gist.github.com/chuckwondo/b151f7f92d54cf01c62da2939238d62b

thanks, @chuckwondo! I took most of your suggestion in 9e26c15 (this PR), with the exception of Make __attrs_post_init__ wrap the opener. I thought it was simpler calling chunk_for_lazy_ops in _apply_where, and the performance difference should be negligable.

Hi @maxrjones, just to clarify, in my gist suggestions, performance isn't the concern. Rather, the reason to drop the re-extraction is provenance/"drift avoidance".

Right now bounds / transform / height / width come from the parent instance's un-chunked extraction, but the pixels come from a second (needless), chunked extraction (both via get_variable). Both extractions happen to align at this point, but there is nothing that guarantees it, so a future change (not in our hands) could cause a drift/provenance problem. Therefore, collapsing them to a single extraction guarantees that drift cannot occur.

In other words, this guarantees that bounds / transform / height / width AND the pixels come from exactly the same extraction, rather than pixels coming from a second extraction that should be effectively the same as the first one, but there is no guarantee that the 2 extractions are truly aligned.

However, that has to happen in the opener, and the reason for that is upstream, not anything in this PR. The inherited titiler.xarray.io.Reader.__attrs_post_init__ opens and extracts back to back:

self.ds = self.opener(self.src_path, **opener_options)
self.input = get_variable(self.ds, self.variable, sel=self.sel)

There is no hook between those two statements that we can leverage from a subclass. The only things a subclass can influence beforehand are self.opener and self.opener_options. That leaves the closure I proposed as a simple, clean way to essentially insert behavior between the 2 lines above by influencing self.opener to guard against any possible future drift between 2 calls to get_variable (by eliminating one of them).

Here are the adjustments I suggest:

1. In chunk_for_lazy_ops, change the return type of chunk_for_lazy_ops to None, rather than xr.Dataset, and remove the return ds at the end, so the function is clearly mutating. (I should have suggested this in the gist to begin with. See the reasoning below.)

2. Then, wrap the opener in __attrs_post_init__, calling chunk_for_lazy_ops purely for the side effect, returning the mutated ds:

   if conditions:
       names = {self.variable, *(c.variable for c in conditions)}
       opener = self.opener
   
       def opener_(src_path: str, **kwargs: Any) -> xr.Dataset:
           ds = opener(src_path, **kwargs)
           chunk_for_lazy_ops(ds, [n for n in names if n in ds])
           return ds
   
       self.opener = opener_

3. In _apply_where, drop the call to chunk_for_lazy_ops, and replace data = self.get_variable(...) with data = self.input. This eliminates the second, unnecessary extraction (get_variable), leaving only the extraction that happens in super().__attrs_post_init__.

Why point 1 is important: my gist has the wrapper return chunk_for_lazy_ops(...) and assigning to self.ds. However, the problem with that is that if chunk_for_lazy_ops gains a call to ds.copy() in the future (for some unforeseen reason), there would be a real leak because the copied value would be returned, but as mentioned in the gist, the copy loses the _close method, becoming uncloseable (the leak). Removing the return value entirely means that nothing can rebind self.ds to a possibly leaky value. That is what makes the closure safe to adopt.

[chuckwondo]

@maxrjones, if you would prefer not to take my latest suggestion in full, then the very least I suggest you do is the bit that removes return ds from chunk_for_lazy_ops and change the return type to None (along with the requisite tweak to the call site that currently expects it to return a value).

[chuckwondo]

@maxrjones, if you would prefer not to take my latest suggestion in full, then the very least I suggest you do is the bit that removes return ds from chunk_for_lazy_ops and change the return type to None (along with the requisite tweak to the call site that currently expects it to return a value).

Note, however, doing this does not address the provenance/drift issue, so I'm encouraging you to take the suggestion in full.

[maxrjones]

thanks for pushing back @chuckwondo. I'll make those changes. I'm also going to break off the independent changes to some smaller PRs.

[maxrjones]

Closing in favor of #166. We can always re-open if others disagree with this decision.