diff --git a/perf_tests/compute_air.py b/perf_tests/compute_air.py index 476cefe5..a564fbd2 100755 --- a/perf_tests/compute_air.py +++ b/perf_tests/compute_air.py @@ -1,13 +1,13 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql if __name__ == "__main__": air = xr.tutorial.open_dataset("air_temperature") chunks = {"time": 240} air = air.chunk(chunks) - df = qr.read_xarray(air).compute() + df = xql.read_xarray(air).read_pandas() print(len(df)) diff --git a/perf_tests/groupby_air.py b/perf_tests/groupby_air.py index c9f4d8ec..b54d7bf9 100755 --- a/perf_tests/groupby_air.py +++ b/perf_tests/groupby_air.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 +from datafusion import SessionContext import xarray as xr -import xarray_sql as qr -from dask_sql import Context +import xarray_sql as xql if __name__ == "__main__": @@ -13,12 +13,12 @@ time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) ).chunk(chunks) - df = qr.read_xarray(air_small) + df = xql.read_xarray_table(air_small) - c = Context() - c.create_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = c.sql( + query = ctx.sql( """ SELECT "lat", "lon", SUM("air") as air_total @@ -29,10 +29,9 @@ """ ) - result = query.compute() + result = query.collect() - expected = air_small.dims["lat"] * air_small.dims["lon"] - assert ( - len(result) == expected - ), f"Length must be {expected}, but was {len(result)}." + expected = air_small.sizes["lat"] * air_small.sizes["lon"] + actual = sum(len(batch) for batch in result) + assert actual == expected, f"Length must be {expected}, but was {actual}." print(expected) diff --git a/perf_tests/groupby_air_full.py b/perf_tests/groupby_air_full.py index e54ad139..57663d84 100755 --- a/perf_tests/groupby_air_full.py +++ b/perf_tests/groupby_air_full.py @@ -1,21 +1,20 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr -from dask_sql import Context - +import xarray_sql as xql +from datafusion import SessionContext if __name__ == "__main__": air = xr.tutorial.open_dataset("air_temperature") chunks = {"time": 240} air = air.chunk(chunks) - df = qr.read_xarray(air) + df = xql.read_xarray_table(air) - c = Context() - c.create_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = c.sql( + query = ctx.sql( """ SELECT "lat", "lon", SUM("air") as air_total @@ -26,10 +25,10 @@ """ ) - result = query.compute() + result = query.collect() + + expected = air.sizes["lat"] * air.sizes["lon"] + actual = sum(len(batch) for batch in result) - expected = air.dims["lat"] * air.dims["lon"] - assert ( - len(result) == expected - ), f"Length must be {expected}, but was {len(result)}." + assert actual == expected, f"Length must be {expected}, but was {actual}." print(expected) diff --git a/perf_tests/open_era5.py b/perf_tests/open_era5.py index 27e7edcc..e9f0b88c 100755 --- a/perf_tests/open_era5.py +++ b/perf_tests/open_era5.py @@ -1,15 +1,15 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql # Requires authenticating with GCP era5_ds = xr.open_zarr( "gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2", chunks={"time": 240, "level": 1}, ) -era5_wind_df = qr.read_xarray( +era5_wind_df = xql.read_xarray( era5_ds[["u_component_of_wind", "v_component_of_wind"]] ) -print(era5_wind_df.columns) +print(era5_wind_df.schema) diff --git a/perf_tests/sanity.py b/perf_tests/sanity.py index 2d51a684..e89df822 100755 --- a/perf_tests/sanity.py +++ b/perf_tests/sanity.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import xarray as xr -import xarray_sql as qr +import xarray_sql as xql if __name__ == "__main__": air = xr.tutorial.open_dataset("air_temperature") @@ -11,6 +11,6 @@ time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) ).chunk(chunks) - df = qr.read_xarray(air_small).compute() + df = xql.read_xarray(air_small).read_pandas() print(len(df))