fix: perf_tests scripts to use pyarrow or LazyArrowStreamTable

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))

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)

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)

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)

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))