Address comments

Author: ahuang11

tests/test_sql.py

@@ -166,48 +166,29 @@ def nan_ds(self):
         },
     ).chunk({"time": 1})
 
-  def test_max_ignores_nan(self, nan_ds):
-    ctx = XarrayContext()
-    ctx.from_dataset("data", nan_ds)
-    result = ctx.sql("SELECT MAX(temp) AS mx FROM data").to_pandas()
-    assert result["mx"].iloc[0] == 8.0
-
-  def test_min_ignores_nan(self, nan_ds):
-    ctx = XarrayContext()
-    ctx.from_dataset("data", nan_ds)
-    result = ctx.sql("SELECT MIN(temp) AS mn FROM data").to_pandas()
-    assert result["mn"].iloc[0] == 1.0
-
-  def test_avg_ignores_nan(self, nan_ds):
+  def test_nan_aggregates(self, nan_ds):
     import numpy as np
 
     ctx = XarrayContext()
     ctx.from_dataset("data", nan_ds)
-    result = ctx.sql("SELECT AVG(temp) AS avg FROM data").to_pandas()
-    expected = np.nanmean([1.0, 2.0, 4.0, 5.0, 7.0, 8.0])
-    assert abs(result["avg"].iloc[0] - expected) < 1e-6
 
-  def test_count_excludes_nan(self, nan_ds):
-    ctx = XarrayContext()
-    ctx.from_dataset("data", nan_ds)
-    result = ctx.sql("SELECT COUNT(temp) AS cnt FROM data").to_pandas()
-    # 8 total cells, 2 are NaN → 6 non-null
-    assert result["cnt"].iloc[0] == 6
-
-  def test_is_null_matches_nan(self, nan_ds):
-    ctx = XarrayContext()
-    ctx.from_dataset("data", nan_ds)
-    result = ctx.sql(
-        "SELECT COUNT(*) AS cnt FROM data WHERE temp IS NULL"
-    ).to_pandas()
-    assert result["cnt"].iloc[0] == 2
-
-  def test_from_dataset_returns_self(self):
-    ds = (
-        xr.tutorial.open_dataset("air_temperature")
-        .isel(time=slice(0, 2), lat=slice(0, 2), lon=slice(0, 2))
-        .chunk({"time": 1})
-    )
-    ctx = XarrayContext()
-    ret = ctx.from_dataset("air", ds)
-    assert ret is ctx
+    # Test multiple aggregates at once:
+    # MAX/MIN/AVG should ignore NaN, COUNT(col) should exclude NaN,
+    # and WHERE col IS NULL should match NaN.
+    query = """
+        SELECT
+            MAX(temp) AS mx,
+            MIN(temp) AS mn,
+            AVG(temp) AS avg,
+            COUNT(temp) AS cnt,
+            COUNT(*) FILTER (WHERE temp IS NULL) AS null_cnt
+        FROM data
+    """
+    result = ctx.sql(query).to_pandas().iloc[0]
+
+    assert result["mx"] == 8.0
+    assert result["mn"] == 1.0
+    expected_avg = np.nanmean([1.0, 2.0, 4.0, 5.0, 7.0, 8.0])
+    assert abs(result["avg"] - expected_avg) < 1e-6
+    assert result["cnt"] == 6
+    assert result["null_cnt"] == 2