Test querying a dataset.
(
self, backend, engine: QueryEngineOptions, parser: QueryParserOptions
)
| 4805 | "backend", ["numpy", pytest.param("dask", marks=[requires_dask])] |
| 4806 | ) |
| 4807 | def test_query( |
| 4808 | self, backend, engine: QueryEngineOptions, parser: QueryParserOptions |
| 4809 | ) -> None: |
| 4810 | """Test querying a dataset.""" |
| 4811 | |
| 4812 | # setup test data |
| 4813 | np.random.seed(42) |
| 4814 | a = np.arange(0, 10, 1) |
| 4815 | b = np.random.randint(0, 100, size=10) |
| 4816 | c = np.linspace(0, 1, 20) |
| 4817 | d = np.random.choice(["foo", "bar", "baz"], size=30, replace=True).astype( |
| 4818 | object |
| 4819 | ) |
| 4820 | aa = DataArray(data=a, dims=["x"], name="a", coords={"a2": ("x", a)}) |
| 4821 | bb = DataArray(data=b, dims=["x"], name="b", coords={"b2": ("x", b)}) |
| 4822 | cc = DataArray(data=c, dims=["y"], name="c", coords={"c2": ("y", c)}) |
| 4823 | dd = DataArray(data=d, dims=["z"], name="d", coords={"d2": ("z", d)}) |
| 4824 | |
| 4825 | if backend == "dask": |
| 4826 | import dask.array as da |
| 4827 | |
| 4828 | aa = aa.copy(data=da.from_array(a, chunks=3)) |
| 4829 | bb = bb.copy(data=da.from_array(b, chunks=3)) |
| 4830 | cc = cc.copy(data=da.from_array(c, chunks=7)) |
| 4831 | dd = dd.copy(data=da.from_array(d, chunks=12)) |
| 4832 | |
| 4833 | # query single dim, single variable |
| 4834 | with raise_if_dask_computes(): |
| 4835 | actual = aa.query(x="a2 > 5", engine=engine, parser=parser) |
| 4836 | expect = aa.isel(x=(a > 5)) |
| 4837 | assert_identical(expect, actual) |
| 4838 | |
| 4839 | # query single dim, single variable, via dict |
| 4840 | with raise_if_dask_computes(): |
| 4841 | actual = aa.query(dict(x="a2 > 5"), engine=engine, parser=parser) |
| 4842 | expect = aa.isel(dict(x=(a > 5))) |
| 4843 | assert_identical(expect, actual) |
| 4844 | |
| 4845 | # query single dim, single variable |
| 4846 | with raise_if_dask_computes(): |
| 4847 | actual = bb.query(x="b2 > 50", engine=engine, parser=parser) |
| 4848 | expect = bb.isel(x=(b > 50)) |
| 4849 | assert_identical(expect, actual) |
| 4850 | |
| 4851 | # query single dim, single variable |
| 4852 | with raise_if_dask_computes(): |
| 4853 | actual = cc.query(y="c2 < .5", engine=engine, parser=parser) |
| 4854 | expect = cc.isel(y=(c < 0.5)) |
| 4855 | assert_identical(expect, actual) |
| 4856 | |
| 4857 | # query single dim, single string variable |
| 4858 | if parser == "pandas": |
| 4859 | # N.B., this query currently only works with the pandas parser |
| 4860 | # xref https://github.com/pandas-dev/pandas/issues/40436 |
| 4861 | with raise_if_dask_computes(): |
| 4862 | actual = dd.query(z='d2 == "bar"', engine=engine, parser=parser) |
| 4863 | expect = dd.isel(z=(d == "bar")) |
| 4864 | assert_identical(expect, actual) |
nothing calls this directly
no test coverage detected