From 7d4954580cc71f61bf7b37c0fed9353a804f7c3b Mon Sep 17 00:00:00 2001
From: Andrew Huang <ahuang@anaconda.com>
Date: Thu, 19 Mar 2026 14:44:45 -0700
Subject: [PATCH 1/3] Add req

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 6bfb1d0..9c6faaf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 test = [
+  "cftime>",
   "pytest",
   "xarray[io]",
   "gcsfs",

From a62ce4870c6312b178ff9a6c9f6450a4125626f4 Mon Sep 17 00:00:00 2001
From: Andrew Huang <ahuang@anaconda.com>
Date: Fri, 20 Mar 2026 15:47:51 -0700
Subject: [PATCH 2/3] Migrate tests

---
 pyproject.toml            |    5 +-
 xarray_sql/df_test.py     |  489 -------------
 xarray_sql/reader_test.py | 1372 -------------------------------------
 xarray_sql/sql_test.py    |  194 ------
 4 files changed, 4 insertions(+), 2056 deletions(-)
 delete mode 100644 xarray_sql/df_test.py
 delete mode 100644 xarray_sql/reader_test.py
 delete mode 100644 xarray_sql/sql_test.py

diff --git a/pyproject.toml b/pyproject.toml
index 9c6faaf..2701dcc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,7 +52,7 @@ features = ["pyo3/extension-module"]
 module-name = "xarray_sql._native"
 
 [tool.setuptools.packages.find]
-exclude = ["demo", "perf_tests"]
+exclude = ["demo", "perf_tests", "tests", "tests.*"]
 
 [tool.pyink]
 line-length = 80
@@ -94,3 +94,6 @@ dev = [
 [tool.uv]
 # Rebuild package when any rust files change
 cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "**/*.rs"}]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py
deleted file mode 100644
index 8eeba0d..0000000
--- a/xarray_sql/df_test.py
+++ /dev/null
@@ -1,489 +0,0 @@
-import tracemalloc
-
-import numpy as np
-import pandas as pd
-import pyarrow as pa
-import pytest
-import xarray as xr
-
-from .df import (
-    DEFAULT_BATCH_SIZE,
-    _parse_schema,
-    block_slices,
-    dataset_to_record_batch,
-    explode,
-    from_map,
-    from_map_batched,
-    iter_record_batches,
-    pivot,
-)
-from .reader import read_xarray, read_xarray_table
-
-
-def rand_wx(start: str, end: str) -> xr.Dataset:
-  np.random.seed(42)
-  lat = np.linspace(-90, 90, num=720)
-  lon = np.linspace(-180, 180, num=1440)
-  time = pd.date_range(start, end, freq="h")
-  level = np.array([1000, 500], dtype=np.int32)
-  reference_time = pd.Timestamp(start)
-  temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
-  precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
-  return xr.Dataset(
-      data_vars=dict(
-          temperature=(["lat", "lon", "time", "level"], temperature),
-          precipitation=(["lat", "lon", "time", "level"], precipitation),
-      ),
-      coords=dict(
-          lat=lat,
-          lon=lon,
-          time=time,
-          level=level,
-          reference_time=reference_time,
-      ),
-      attrs=dict(description="Random weather."),
-  )
-
-
-def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
-  """Create a large xarray dataset for memory testing."""
-  np.random.seed(42)
-
-  time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
-  lat = np.linspace(-90, 90, lat_points)
-  lon = np.linspace(-180, 180, lon_points)
-
-  temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
-  precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
-
-  return xr.Dataset(
-      {
-          "temperature": (["time", "lat", "lon"], temp_data),
-          "precipitation": (["time", "lat", "lon"], precip_data),
-      },
-      coords={"time": time, "lat": lat, "lon": lon},
-  )
-
-
-def adding_function(x, y):
-  """Simple function that adds two values and returns a DataFrame."""
-  result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
-  return result
-
-
-@pytest.fixture
-def air():
-  ds = xr.tutorial.open_dataset("air_temperature")
-  chunks = {"time": 240}
-  return ds.chunk(chunks)
-
-
-@pytest.fixture
-def air_small(air):
-  return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
-      {"time": 240}
-  )
-
-
-@pytest.fixture
-def randwx():
-  return rand_wx("1995-01-13T00", "1995-01-13T01")
-
-
-@pytest.fixture
-def large_ds():
-  return create_large_dataset().chunk({"time": 25})
-
-
-def test_explode_cardinality(air):
-  dss = explode(air)
-  assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()])
-
-
-def test_explode_dim_sizes_one(air):
-  chunks = {"time": 240}
-  ds = next(iter(explode(air)))
-  for k, v in chunks.items():
-    assert k in ds.dims
-    assert v == ds.sizes[k]
-
-
-def test_explode_data_equal_one_first(air):
-  ds = next(iter(explode(air)))
-  iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()}
-  assert air.isel(iselection).equals(ds)
-
-
-def test_explode_data_equal_one_last(air):
-  dss = list(explode(air))
-  ds = dss[-1]
-
-  # For the last chunk, we need to calculate where it actually starts
-  # The original logic slice(0, s) only works for the first chunk
-  iselection = {}
-  for dim in ds.dims:
-    # Get chunk boundaries
-    chunk_bounds = np.cumsum((0,) + air.chunks[dim])
-    # Last chunk index
-    last_chunk_idx = len(air.chunks[dim]) - 1
-    # Calculate actual start and end positions
-    start = chunk_bounds[last_chunk_idx]
-    end = chunk_bounds[last_chunk_idx + 1]
-    iselection[dim] = slice(start, end)
-
-  assert air.isel(iselection).equals(ds)
-
-
-def test_from_map_basic():
-  def make_df(x):
-    return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]})
-
-  result = from_map(make_df, [1, 2, 3])
-  assert isinstance(result, pa.Table)
-  assert len(result) == 6
-  assert result.column_names == ["value", "index"]
-
-
-def test_from_map_multiple_iterables():
-  def add_values(x, y):
-    return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]})
-
-  result = from_map(add_values, [1, 2], [10, 20])
-  assert isinstance(result, pa.Table)
-  assert len(result) == 2
-
-  df = result.to_pandas()
-  assert list(df["sum"]) == [11, 22]
-
-
-def test_from_map_with_args():
-  def multiply_and_add(x, multiplier, add_value):
-    return pd.DataFrame({"result": [x * multiplier + add_value]})
-
-  result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10))
-  assert isinstance(result, pa.Table)
-  assert len(result) == 3
-
-  df = result.to_pandas()
-  assert list(df["result"]) == [12, 14, 16]
-
-
-def test_from_map_with_pyarrow_tables():
-  def make_arrow_table(x):
-    df = pd.DataFrame({"value": [x]})
-    return pa.Table.from_pandas(df)
-
-  result = from_map(make_arrow_table, [1, 2, 3])
-  assert isinstance(result, pa.Table)
-  assert len(result) == 3
-
-
-def test_iter_record_batches_splits_into_multiple_batches(air_small):
-  """iter_record_batches should emit >1 batch when partition exceeds batch_size."""
-  schema = _parse_schema(air_small)
-  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
-  ds_block = air_small.isel(block)
-  total_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes]))
-
-  small_batch = 16  # force many small batches
-  batches = list(iter_record_batches(ds_block, schema, batch_size=small_batch))
-
-  assert len(batches) == -(-total_rows // small_batch)  # ceiling division
-  assert all(b.num_rows <= small_batch for b in batches)
-  assert sum(b.num_rows for b in batches) == total_rows
-
-
-def test_iter_record_batches_matches_dataset_to_record_batch(air_small):
-  """Concatenating all iter_record_batches output must equal dataset_to_record_batch."""
-  schema = _parse_schema(air_small)
-  dim_cols = [f.name for f in schema if f.name in air_small.dims]
-  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
-  ds_block = air_small.isel(block)
-
-  batches = list(iter_record_batches(ds_block, schema, batch_size=16))
-  actual_df = (
-      pa.Table.from_batches(batches)
-      .to_pandas()
-      .sort_values(dim_cols)
-      .reset_index(drop=True)
-  )
-  expected_df = (
-      dataset_to_record_batch(ds_block, schema)
-      .to_pandas()
-      .sort_values(dim_cols)
-      .reset_index(drop=True)
-  )
-  pd.testing.assert_frame_equal(actual_df, expected_df)
-
-
-def test_iter_record_batches_default_batch_size():
-  """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch."""
-  ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2))
-  schema = _parse_schema(ds)
-  total_rows = int(np.prod([ds.sizes[d] for d in ds.sizes]))
-  assert total_rows <= DEFAULT_BATCH_SIZE, "fixture too large — adjust isel"
-  batches = list(iter_record_batches(ds, schema))
-  assert len(batches) == 1
-  assert batches[0].num_rows == total_rows
-
-
-def test_dataset_to_record_batch_matches_pivot(air_small):
-  """dataset_to_record_batch should contain the same rows as pivot.
-
-  Row ordering may differ (pivot uses ds.dims key order; dataset_to_record_batch
-  uses the data variable's own dim order). Both orderings are valid for SQL, so
-  we sort by the coordinate columns before comparing.
-  """
-  schema = _parse_schema(air_small)
-  dim_cols = [f.name for f in schema if f.name in air_small.dims]
-  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
-
-  for block in blocks:
-    ds_block = air_small.isel(block)
-    actual_df = (
-        dataset_to_record_batch(ds_block, schema)
-        .to_pandas()
-        .sort_values(dim_cols)
-        .reset_index(drop=True)
-    )
-    expected_df = (
-        pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema)
-        .to_pandas()
-        .sort_values(dim_cols)
-        .reset_index(drop=True)
-    )
-
-    pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False)
-
-
-def test_dataset_to_record_batch_column_order(air_small):
-  """Output column order must match schema (dims first, then data vars)."""
-  schema = _parse_schema(air_small)
-  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
-  batch = dataset_to_record_batch(air_small.isel(block), schema)
-  assert batch.schema.names == schema.names
-
-
-def test_dataset_to_record_batch_row_count(air_small):
-  """Row count must equal the product of the block dimension sizes."""
-  schema = _parse_schema(air_small)
-  chunks = {"time": 4, "lat": 3, "lon": 4}
-  for block in block_slices(air_small, chunks=chunks):
-    ds_block = air_small.isel(block)
-    expected_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes]))
-    batch = dataset_to_record_batch(ds_block, schema)
-    assert batch.num_rows == expected_rows
-
-
-def test_from_map_batched_basic_functionality(air_small):
-  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
-
-  first_block_df = pivot(air_small.isel(blocks[0]))
-  expected_schema = pa.Schema.from_pandas(first_block_df)
-
-  reader = from_map_batched(
-      pivot, [air_small.isel(block) for block in blocks], schema=expected_schema
-  )
-
-  assert isinstance(reader, pa.RecordBatchReader)
-  assert reader.schema == expected_schema
-
-  batches = list(reader)
-  assert len(batches) > 0
-  for batch in batches:
-    assert batch.schema == expected_schema
-    assert len(batch) > 0
-
-
-def test_from_map_batched_multiple_iterables():
-  x_values = [1, 2, 3, 4, 5]
-  y_values = [10, 20, 30, 40, 50]
-
-  expected_schema = pa.schema(
-      [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())]
-  )
-
-  reader = from_map_batched(
-      adding_function, x_values, y_values, schema=expected_schema
-  )
-  table = reader.read_all()
-  df = table.to_pandas()
-
-  expected_df = pd.DataFrame(
-      {
-          "x": x_values,
-          "y": y_values,
-          "sum": [x + y for x, y in zip(x_values, y_values)],
-      }
-  )
-  pd.testing.assert_frame_equal(df, expected_df)
-
-
-def test_from_map_batched_with_args_and_kwargs():
-  def multiply_and_add(x, multiplier, offset=0):
-    return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]})
-
-  values = [1, 2, 3]
-  expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())])
-
-  reader = from_map_batched(
-      multiply_and_add, values, args=(2,), offset=5, schema=expected_schema
-  )
-  table = reader.read_all()
-  df = table.to_pandas()
-
-  expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]})
-  pd.testing.assert_frame_equal(df, expected_df)
-
-
-def test_from_map_batched_empty_iterables():
-  empty_schema = pa.schema([("value", pa.int64())])
-
-  reader = from_map_batched(
-      lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema
-  )
-  batches = list(reader)
-  assert len(batches) == 0
-
-
-def test_from_map_batched_consistency_with_regular_map(air_small):
-  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3}))
-  datasets = [air_small.isel(block) for block in blocks]
-
-  first_df = pivot(datasets[0])
-  schema = pa.Schema.from_pandas(first_df)
-
-  reader = from_map_batched(pivot, datasets, schema=schema)
-  batched_table = reader.read_all()
-
-  regular_dfs = [pivot(ds) for ds in datasets]
-  regular_table = pa.Table.from_pandas(
-      pd.concat(regular_dfs, ignore_index=True)
-  )
-
-  assert batched_table.schema == regular_table.schema
-  assert len(batched_table) == len(regular_table)
-
-  batched_df = (
-      batched_table.to_pandas()
-      .sort_values(["time", "lat", "lon"])
-      .reset_index(drop=True)
-  )
-  regular_df = (
-      regular_table.to_pandas()
-      .sort_values(["time", "lat", "lon"])
-      .reset_index(drop=True)
-  )
-
-  pd.testing.assert_frame_equal(batched_df, regular_df)
-
-
-def test_from_map_batched_integration_with_datafusion_via_read_xarray():
-  air = xr.tutorial.open_dataset("air_temperature")
-  air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15))
-  air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8})
-
-  arrow_stream = read_xarray(
-      air_chunked, chunks={"time": 25, "lat": 5, "lon": 8}
-  )
-
-  assert hasattr(arrow_stream, "schema")
-  assert hasattr(arrow_stream, "__iter__")
-
-  table = arrow_stream.read_all()
-  assert len(table) > 0
-
-  expected_columns = {"time", "lat", "lon", "air"}
-  actual_columns = set(table.column_names)
-  assert expected_columns.issubset(actual_columns)
-
-
-def test_read_xarray_loads_one_chunk_at_a_time(large_ds):
-  tracemalloc.start()
-  iterable = read_xarray(large_ds)
-  first_size, first_peak = tracemalloc.get_traced_memory()
-  tracemalloc.reset_peak()
-
-  sizes, peaks = [], []
-
-  first_chunk = large_ds.isel(next(block_slices(large_ds)))
-  chunk_size = first_chunk.nbytes
-
-  # Creating the iterator should be inexpensive -- less than one chunk.
-  # We multiply by constant factors because chunks have additional overhead
-  assert first_size < chunk_size * 3
-  assert first_peak < chunk_size * 6
-
-  for it in iterable:
-    _ = it
-    cur_size, cur_peak = tracemalloc.get_traced_memory()
-    tracemalloc.reset_peak()
-    sizes.append(cur_size)
-    peaks.append(cur_peak)
-
-  for size in sizes:
-    # Observed range: 1.59–1.83× chunk_size.
-    # iter_record_batches holds data-variable arrays (≈1× chunk) while
-    # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk).
-    assert chunk_size * 1.3 < size, f"size {size} unexpectedly low"
-    assert chunk_size * 2.2 > size, f"size {size} unexpectedly high"
-
-  for peak in peaks:
-    # Observed range: 1.84–3.28× chunk_size.
-    # Peak includes data arrays + Arrow batch + temporary coordinate index
-    # arrays; the first batch of each chunk is highest (Dask compute overhead).
-    assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low"
-    assert chunk_size * 4.0 > peak, f"peak {peak} unexpectedly high"
-
-  assert max(peaks) < large_ds.nbytes
-
-  tracemalloc.stop()
-
-
-def test_read_xarray_table_memory_bounds(large_ds):
-  """read_xarray_table should not materialise data at registration time.
-
-  Registration should only hold coordinate arrays and Rust partition metadata
-  (no data variables).  Peak memory during a full-table query should be a
-  small fraction of the whole dataset — i.e. partitions are processed without
-  loading all of them simultaneously.
-  """
-  from datafusion import SessionContext
-
-  first_chunk = large_ds.isel(next(block_slices(large_ds)))
-  chunk_size = first_chunk.nbytes
-
-  # --- Registration phase ---
-  tracemalloc.start()
-  table = read_xarray_table(large_ds)
-  reg_size, reg_peak = tracemalloc.get_traced_memory()
-  tracemalloc.reset_peak()
-
-  # The lazy generator only materialises coord arrays (~O(dim sizes)) and
-  # factory closure objects — no data arrays.  Both metrics should be well
-  # below one chunk of data.
-  assert reg_size < chunk_size, (
-      f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: "
-      "data may have been loaded eagerly"
-  )
-  assert (
-      reg_peak < chunk_size * 2
-  ), f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})"
-
-  # --- Query phase ---
-  ctx = SessionContext()
-  ctx.register_table("weather", table)
-  ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM weather").collect()
-  _, query_peak = tracemalloc.get_traced_memory()
-
-  # tracemalloc measures Python-heap allocations, which include Arrow
-  # buffer copies and object overhead on top of the raw data.  The
-  # observed peak is typically 1.1–1.5× the raw dataset size; we use
-  # 2× as a generous bound that would still catch catastrophic regressions
-  # (e.g. loading all partitions twice simultaneously).
-  assert query_peak < large_ds.nbytes * 2, (
-      f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: "
-      "may be holding excessive data in memory"
-  )
-
-  tracemalloc.stop()
diff --git a/xarray_sql/reader_test.py b/xarray_sql/reader_test.py
deleted file mode 100644
index 6ab6406..0000000
--- a/xarray_sql/reader_test.py
+++ /dev/null
@@ -1,1372 +0,0 @@
-"""Tests for XarrayRecordBatchReader lazy streaming behavior.
-
-These tests verify that XarrayRecordBatchReader provides true lazy evaluation:
-- No data iteration during reader creation
-- No data iteration during DataFusion table registration (using LazyArrowStreamTable)
-- Data iteration ONLY occurs during query execution (collect())
-
-The lazy streaming is achieved via the Rust LazyArrowStreamTable class which
-implements the __datafusion_table_provider__ protocol using StreamingTable.
-
-Additional tests verify:
-- True streaming with bounded memory (batches processed incrementally)
-- Back-pressure behavior (producer pauses when consumer is slow)
-- Error propagation through the stream
-"""
-
-import threading
-import time
-import numpy as np
-import pandas as pd
-import pyarrow as pa
-import pytest
-import xarray as xr
-from datafusion import SessionContext
-
-from ._native import LazyArrowStreamTable
-from .reader import XarrayRecordBatchReader, read_xarray_table
-from .df import _parse_schema
-
-
-@pytest.fixture
-def small_ds():
-  """Create a small dataset for testing."""
-  np.random.seed(42)
-  time = pd.date_range("2020-01-01", periods=100, freq="h")
-  lat = np.linspace(-90, 90, 10)
-  lon = np.linspace(-180, 180, 10)
-
-  data = np.random.rand(100, 10, 10).astype(np.float32)
-
-  return xr.Dataset(
-      {"temperature": (["time", "lat", "lon"], data)},
-      coords={"time": time, "lat": lat, "lon": lon},
-  )
-
-
-class IterationTracker:
-  """Tracks when iteration occurs for testing lazy evaluation.
-
-  The callback signature is ``(block, projection_names)`` where
-  ``projection_names`` is the list of column names requested by the query
-  (``None`` when no projection pushdown occurred, e.g. for
-  ``XarrayRecordBatchReader`` or a ``SELECT *`` query).
-  """
-
-  def __init__(self):
-    self.iteration_count = 0
-    self.blocks_seen = []
-    self.projections_seen = []
-
-  def __call__(self, block, projection_names=None):
-    self.iteration_count += 1
-    self.blocks_seen.append(block)
-    self.projections_seen.append(projection_names)
-
-  def reset(self):
-    self.iteration_count = 0
-    self.blocks_seen = []
-    self.projections_seen = []
-
-
-class TestXarrayRecordBatchReaderCreation:
-  """Tests that reader creation does NOT trigger data iteration."""
-
-  def test_reader_creation_does_not_iterate(self, small_ds):
-    """Creating a reader should NOT iterate through any data."""
-    tracker = IterationTracker()
-
-    reader = XarrayRecordBatchReader(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    assert tracker.iteration_count == 0, (
-        f"Expected 0 iterations during reader creation, "
-        f"but got {tracker.iteration_count}"
-    )
-
-  def test_schema_access_does_not_iterate(self, small_ds):
-    """Accessing the schema should NOT trigger iteration."""
-    tracker = IterationTracker()
-
-    reader = XarrayRecordBatchReader(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    # Access schema
-    _ = reader.schema
-    _ = reader.__arrow_c_schema__()
-
-    assert tracker.iteration_count == 0, (
-        f"Expected 0 iterations when accessing schema, "
-        f"but got {tracker.iteration_count}"
-    )
-
-
-class TestDataFusionRegistration:
-  """Tests that DataFusion table registration does NOT trigger iteration.
-
-  These tests use read_xarray_table with register_table()
-  to achieve true lazy evaluation.
-  """
-
-  def test_register_table_does_not_iterate(self, small_ds):
-    """Registering a LazyArrowStreamTable should NOT iterate data.
-
-    This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps
-    a factory and implements __datafusion_table_provider__ with StreamingTable,
-    ensuring data is only read during query execution.
-    """
-    tracker = IterationTracker()
-
-    # Use read_xarray_table which creates a factory-based table
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    assert tracker.iteration_count == 0, (
-        f"LAZY EVALUATION FAILED: Expected 0 iterations during "
-        f"register_table(), but got {tracker.iteration_count}."
-    )
-
-  def test_sql_planning_does_not_iterate(self, small_ds):
-    """Creating a SQL query plan should NOT iterate data."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Create a query but don't execute it
-    query = ctx.sql("SELECT AVG(temperature) FROM test_table")
-
-    # Just creating the query shouldn't iterate
-    assert tracker.iteration_count == 0, (
-        f"Expected 0 iterations during SQL planning, "
-        f"but got {tracker.iteration_count}. "
-        f"DataFusion may be scanning data during query planning."
-    )
-
-
-class TestDataFusionCollect:
-  """Tests that data iteration ONLY occurs during collect().
-
-  These tests use read_xarray_table to verify lazy evaluation.
-  """
-
-  def test_collect_triggers_iteration(self, small_ds):
-    """collect() should trigger data iteration."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Verify no iteration yet (lazy registration)
-    iteration_before_collect = tracker.iteration_count
-    assert (
-        iteration_before_collect == 0
-    ), "Should have 0 iterations before collect"
-
-    # Now collect - this SHOULD iterate
-    result = ctx.sql("SELECT * FROM test_table LIMIT 10").collect()
-
-    assert tracker.iteration_count > 0, (
-        f"Expected iterations during collect(), but got 0. "
-        f"Data was never read!"
-    )
-    assert (
-        tracker.iteration_count > iteration_before_collect
-    ), f"Expected more iterations after collect()"
-
-  def test_full_query_iterates_all_blocks(self, small_ds):
-    """A query that reads all data should iterate all blocks."""
-    tracker = IterationTracker()
-
-    chunks = {"time": 25}
-    table = read_xarray_table(
-        small_ds,
-        chunks=chunks,
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Run a query that needs to scan all data
-    result = ctx.sql("SELECT COUNT(*) FROM test_table").collect()
-
-    # With time=100 and chunks=25, we expect 4 blocks
-    expected_blocks = 100 // 25
-    assert tracker.iteration_count == expected_blocks, (
-        f"Expected {expected_blocks} block iterations, "
-        f"but got {tracker.iteration_count}"
-    )
-
-  def test_aggregation_query_iterates_correctly(self, small_ds):
-    """Aggregation queries should iterate all necessary blocks."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Run aggregation
-    result = ctx.sql(
-        "SELECT lat, AVG(temperature) as avg_temp "
-        "FROM test_table GROUP BY lat"
-    ).collect()
-
-    # Should have iterated some blocks
-    assert tracker.iteration_count > 0
-    assert len(result) > 0
-
-
-class TestLazyEvaluationEndToEnd:
-  """End-to-end tests verifying lazy evaluation through the full pipeline.
-
-  These tests use read_xarray_table to achieve true lazy evaluation.
-  """
-
-  def test_lazy_evaluation_sequence(self, small_ds):
-    """Verify the exact sequence of lazy evaluation stages.
-
-    This is the comprehensive test that proves true lazy evaluation:
-    1. Table creation: 0 iterations
-    2. Table registration: 0 iterations
-    3. Query planning: 0 iterations
-    4. collect(): N iterations (where N = number of blocks)
-    """
-    tracker = IterationTracker()
-
-    # Stage 1: Table creation (with factory)
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-    iterations_after_table = tracker.iteration_count
-    assert iterations_after_table == 0, (
-        f"Stage 1 FAILED: Table creation triggered "
-        f"{iterations_after_table} iterations"
-    )
-
-    # Stage 2: Table registration
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-    iterations_after_registration = tracker.iteration_count
-    assert iterations_after_registration == 0, (
-        f"Stage 2 FAILED: Table registration triggered "
-        f"{iterations_after_registration} iterations"
-    )
-
-    # Stage 3: Query planning
-    query = ctx.sql("SELECT * FROM test_table")
-    iterations_after_planning = tracker.iteration_count
-    assert iterations_after_planning == 0, (
-        f"Stage 3 FAILED: Query planning triggered "
-        f"{iterations_after_planning} iterations"
-    )
-
-    # Stage 4: collect() - NOW iteration should happen
-    result = query.collect()
-    iterations_after_collect = tracker.iteration_count
-    assert (
-        iterations_after_collect > 0
-    ), f"Stage 4 FAILED: collect() triggered 0 iterations - no data was read!"
-
-    # Verify we got the expected number of blocks (100 time steps / 25 = 4)
-    expected_blocks = 4
-    assert (
-        iterations_after_collect == expected_blocks
-    ), f"Expected {expected_blocks} iterations, got {iterations_after_collect}"
-
-  def test_multiple_queries_on_same_table(self, small_ds):
-    """Same table can be queried multiple times with fresh iteration each time."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 50},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # First query
-    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
-    first_query_iterations = tracker.iteration_count
-    assert first_query_iterations > 0, "First query should iterate"
-
-    # Second query on same table - should iterate again
-    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
-    second_query_iterations = tracker.iteration_count
-    assert (
-        second_query_iterations > first_query_iterations
-    ), "Second query should trigger additional iterations"
-
-  def test_stream_consumed_error(self, small_ds):
-    """Once consumed, a single XarrayRecordBatchReader should not be reusable."""
-    reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25})
-
-    # Consume the reader by converting to a PyArrow reader and reading
-    import pyarrow as pa
-
-    pa_reader = pa.RecordBatchReader.from_stream(reader)
-    _ = pa_reader.read_all()
-
-    # Reader is now consumed, calling __arrow_c_stream__ again should fail
-    with pytest.raises(RuntimeError, match="already consumed"):
-      reader.__arrow_c_stream__()
-
-
-class TestDataIntegrity:
-  """Tests that verify data correctness alongside lazy evaluation.
-
-  These tests use read_xarray_table for lazy streaming.
-  """
-
-  def test_query_results_are_correct(self, small_ds):
-    """Verify that lazy evaluation produces correct results."""
-    table = read_xarray_table(small_ds, chunks={"time": 25})
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Get count
-    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
-    count = result[0].to_pandas()["cnt"].iloc[0]
-
-    # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows
-    expected_count = 100 * 10 * 10
-    assert (
-        count == expected_count
-    ), f"Expected {expected_count} rows, got {count}"
-
-  def test_aggregation_results_are_correct(self, small_ds):
-    """Verify aggregation produces correct results."""
-    table = read_xarray_table(small_ds, chunks={"time": 25})
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Get average temperature
-    result = ctx.sql(
-        "SELECT AVG(temperature) as avg_temp FROM test_table"
-    ).collect()
-    avg_temp = result[0].to_pandas()["avg_temp"].iloc[0]
-
-    # With seed 42 and random data in [0, 1), average should be ~0.5
-    assert (
-        0.4 < avg_temp < 0.6
-    ), f"Expected average temperature ~0.5, got {avg_temp}"
-
-
-class TestPyArrowInterop:
-  """Tests for PyArrow interoperability."""
-
-  def test_from_stream_does_not_iterate(self, small_ds):
-    """pa.RecordBatchReader.from_stream() should not iterate."""
-    tracker = IterationTracker()
-
-    reader = XarrayRecordBatchReader(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    # Create PyArrow reader from our stream
-    pa_reader = pa.RecordBatchReader.from_stream(reader)
-
-    assert tracker.iteration_count == 0, (
-        f"Expected 0 iterations when creating PyArrow reader, "
-        f"but got {tracker.iteration_count}"
-    )
-
-  def test_pyarrow_iteration_triggers_callbacks(self, small_ds):
-    """Iterating via PyArrow should trigger our callbacks."""
-    tracker = IterationTracker()
-
-    reader = XarrayRecordBatchReader(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    pa_reader = pa.RecordBatchReader.from_stream(reader)
-
-    # Now iterate
-    for batch in pa_reader:
-      pass
-
-    assert (
-        tracker.iteration_count == 4
-    ), f"Expected 4 iterations, got {tracker.iteration_count}"
-
-  def test_read_all_iterates_all(self, small_ds):
-    """read_all() should iterate through all blocks."""
-    tracker = IterationTracker()
-
-    reader = XarrayRecordBatchReader(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    pa_reader = pa.RecordBatchReader.from_stream(reader)
-    table = pa_reader.read_all()
-
-    assert tracker.iteration_count == 4
-    assert len(table) == 100 * 10 * 10
-
-
-class StreamingTracker:
-  """Tracks timing of batch iterations to verify streaming behavior.
-
-  This tracker records when each batch is processed, allowing us to verify
-  that batches are streamed incrementally rather than all loaded at once.
-  """
-
-  def __init__(self):
-    self.batch_times = []
-    self.batch_count = 0
-    self._lock = threading.Lock()
-
-  def __call__(self, block, projection_names=None):
-    with self._lock:
-      self.batch_times.append(time.monotonic())
-      self.batch_count += 1
-
-  def reset(self):
-    with self._lock:
-      self.batch_times = []
-      self.batch_count = 0
-
-  @property
-  def max_concurrent_batches_estimate(self):
-    """Estimate max batches that could have been in memory simultaneously.
-
-    If all batches are loaded at once, all batch_times will be very close.
-    If streaming works correctly, batch_times should be spread out.
-    """
-    if len(self.batch_times) < 2:
-      return len(self.batch_times)
-
-    # Sort times and look at gaps
-    sorted_times = sorted(self.batch_times)
-    # If times are spread out, streaming is working
-    # If all times are within a tiny window, all batches loaded at once
-    total_duration = sorted_times[-1] - sorted_times[0]
-
-    # If the spread is very small compared to number of batches,
-    # batches were likely all loaded at once
-    return len(self.batch_times)
-
-
-class TestStreamingBehavior:
-  """Tests that verify true streaming with bounded memory.
-
-  These tests ensure that the Rust implementation streams batches through
-  a bounded channel rather than loading all data into memory at once.
-  """
-
-  def test_batches_processed_incrementally(self, small_ds):
-    """Verify batches are processed one at a time, not all at once.
-
-    This test uses a callback that tracks when each batch is processed.
-    With true streaming, batches should be processed incrementally.
-    """
-    tracker = StreamingTracker()
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Run query that scans all data
-    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
-
-    # All 4 batches should have been processed
-    assert (
-        tracker.batch_count == 4
-    ), f"Expected 4 batches, got {tracker.batch_count}"
-
-  def test_all_partitions_processed(self, small_ds):
-    """Verify that all partitions are processed (order may vary with parallelism)."""
-    blocks_seen = []
-
-    def track_order(block, projection_names=None):
-      # Record the time slice for ordering verification
-      blocks_seen.append(block.get("time", None))
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=track_order,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-    ctx.sql("SELECT * FROM test_table").collect()
-
-    # Should have 4 blocks/partitions
-    assert len(blocks_seen) == 4
-
-    # All blocks should be present (though order may vary due to parallelism)
-    # Extract start positions and verify they cover all expected ranges
-    starts = sorted([b.start for b in blocks_seen])
-    expected_starts = [0, 25, 50, 75]
-    assert (
-        starts == expected_starts
-    ), f"Expected partition starts {expected_starts}, got {starts}"
-
-  def test_large_dataset_streams_correctly(self):
-    """Test streaming with a larger dataset to verify memory behavior.
-
-    This test creates a dataset with many blocks to verify that
-    streaming works correctly at scale.
-    """
-    # Create a dataset with 20 blocks
-    np.random.seed(42)
-    time = pd.date_range("2020-01-01", periods=200, freq="h")
-    lat = np.linspace(-90, 90, 10)
-    lon = np.linspace(-180, 180, 10)
-
-    data = np.random.rand(200, 10, 10).astype(np.float32)
-
-    large_ds = xr.Dataset(
-        {"temperature": (["time", "lat", "lon"], data)},
-        coords={"time": time, "lat": lat, "lon": lon},
-    )
-
-    tracker = StreamingTracker()
-
-    # Use small chunks to create many blocks
-    table = read_xarray_table(
-        large_ds,
-        chunks={"time": 10},  # 200 / 10 = 20 blocks
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # Run a query that needs all data
-    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
-    count = result[0].to_pandas()["cnt"].iloc[0]
-
-    # Verify all blocks were processed
-    assert (
-        tracker.batch_count == 20
-    ), f"Expected 20 batches for large dataset, got {tracker.batch_count}"
-
-    # Verify data integrity
-    expected_count = 200 * 10 * 10
-    assert (
-        count == expected_count
-    ), f"Expected {expected_count} rows, got {count}"
-
-
-class TestBoundedMemoryBehavior:
-  """Tests that verify memory usage remains bounded during streaming.
-
-  The key property we're testing: only a small number of batches should
-  be in memory at once (the channel buffer size, which is 4), not the
-  entire dataset.
-
-  These tests verify that:
-  1. Many batches can be processed without loading all into memory
-  2. Production times are spread out (indicating back-pressure)
-  3. Large datasets complete successfully (memory doesn't explode)
-  """
-
-  def test_many_batches_stream_successfully(self):
-    """Verify streaming works with many more batches than buffer size.
-
-    With buffer size = 4, if we have 16 batches and streaming works,
-    the query should complete successfully. If all batches were loaded
-    at once (no streaming), this would use 4x more memory.
-    """
-    # Create dataset with 16 batches (4x buffer size)
-    np.random.seed(42)
-    time_coord = pd.date_range("2020-01-01", periods=160, freq="h")
-    lat = np.linspace(-90, 90, 5)
-    lon = np.linspace(-180, 180, 5)
-    data = np.random.rand(160, 5, 5).astype(np.float32)
-
-    ds = xr.Dataset(
-        {"temperature": (["time", "lat", "lon"], data)},
-        coords={"time": time_coord, "lat": lat, "lon": lon},
-    )
-
-    tracker = StreamingTracker()
-
-    # 16 batches (160 / 10 = 16)
-    table = read_xarray_table(
-        ds,
-        chunks={"time": 10},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
-    count = result[0].to_pandas()["cnt"].iloc[0]
-
-    # All 16 batches should have been processed
-    assert (
-        tracker.batch_count == 16
-    ), f"Expected 16 batches, got {tracker.batch_count}"
-
-    # Verify data integrity
-    expected = 160 * 5 * 5
-    assert count == expected, f"Expected {expected} rows, got {count}"
-
-  def test_production_times_spread_out(self):
-    """Verify batch production is spread over time, not instant.
-
-    If back-pressure works, later batches can only be produced after
-    earlier batches have been consumed. Production times should span
-    a non-zero duration.
-    """
-    np.random.seed(123)
-    time_coord = pd.date_range("2020-01-01", periods=100, freq="h")
-    lat = np.linspace(-90, 90, 5)
-    lon = np.linspace(-180, 180, 5)
-    data = np.random.rand(100, 5, 5).astype(np.float32)
-
-    ds = xr.Dataset(
-        {"temperature": (["time", "lat", "lon"], data)},
-        coords={"time": time_coord, "lat": lat, "lon": lon},
-    )
-
-    tracker = StreamingTracker()
-
-    # 10 batches, more than buffer size of 4
-    table = read_xarray_table(
-        ds,
-        chunks={"time": 10},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
-
-    # All 10 batches should be produced
-    assert tracker.batch_count == 10
-
-    # Production should span some time (not all instant)
-    sorted_times = sorted(tracker.batch_times)
-    production_span = sorted_times[-1] - sorted_times[0]
-
-    # With streaming and back-pressure, production_span should be > 0
-    # (If all batches were produced simultaneously, span would be ~0)
-    assert production_span >= 0, "Production span should be non-negative"
-
-  def test_large_batch_count_completes(self):
-    """Verify that processing many batches completes successfully.
-
-    This is a stress test: 50 batches is well above the buffer size of 4.
-    If streaming works correctly, this should complete without memory issues.
-    """
-    np.random.seed(456)
-    time_coord = pd.date_range("2020-01-01", periods=500, freq="h")
-    lat = np.linspace(-90, 90, 10)
-    lon = np.linspace(-180, 180, 10)
-    data = np.random.rand(500, 10, 10).astype(np.float32)
-
-    ds = xr.Dataset(
-        {"temperature": (["time", "lat", "lon"], data)},
-        coords={"time": time_coord, "lat": lat, "lon": lon},
-    )
-
-    tracker = StreamingTracker()
-
-    # 50 batches (500 / 10 = 50)
-    table = read_xarray_table(
-        ds,
-        chunks={"time": 10},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
-    count = result[0].to_pandas()["cnt"].iloc[0]
-
-    # All 50 batches processed
-    assert (
-        tracker.batch_count == 50
-    ), f"Expected 50 batches, got {tracker.batch_count}"
-
-    # Data integrity
-    expected = 500 * 10 * 10
-    assert count == expected, f"Expected {expected} rows, got {count}"
-
-  def test_aggregation_with_many_batches(self):
-    """Verify aggregation queries work correctly with many batches.
-
-    GROUP BY queries require processing all data, making them a good
-    test for streaming behavior. Uses collect() to verify that parallel
-    aggregation returns complete results (fixed in DataFusion 52+).
-    """
-    np.random.seed(789)
-    time_coord = pd.date_range("2020-01-01", periods=120, freq="h")
-    # Use integer lat/lon to avoid floating point grouping issues
-    lat = np.array([0, 1, 2, 3, 4], dtype=np.float64)
-    lon = np.array([0, 1, 2, 3, 4], dtype=np.float64)
-    data = np.random.rand(120, 5, 5).astype(np.float32)
-
-    ds = xr.Dataset(
-        {"temperature": (["time", "lat", "lon"], data)},
-        coords={"time": time_coord, "lat": lat, "lon": lon},
-    )
-
-    tracker = StreamingTracker()
-
-    # 12 partitions (one per chunk)
-    table = read_xarray_table(
-        ds,
-        chunks={"time": 10},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # GROUP BY requires scanning all data; collect() must return complete results
-    df = ctx.sql(
-        "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat"
-    ).to_pandas()
-
-    assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}"
-
-    # All partitions processed
-    assert (
-        tracker.batch_count == 12
-    ), f"Expected 12 partitions, got {tracker.batch_count}"
-
-
-class TestErrorPropagation:
-  """Tests that verify errors are properly propagated through the stream.
-
-  These tests ensure that errors during batch reading surface to the user
-  rather than being silently swallowed.
-  """
-
-  def test_factory_error_propagates(self):
-    """Errors from the factory function should propagate to the user."""
-
-    def failing_factory():
-      raise ValueError("Factory intentionally failed")
-
-    schema = pa.schema([("value", pa.int64())])
-    # partitions is an iterable of (factory, metadata_dict) pairs
-    table = LazyArrowStreamTable([(failing_factory, {})], schema)
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # The error should surface when we try to collect
-    with pytest.raises(Exception) as exc_info:
-      ctx.sql("SELECT * FROM test_table").collect()
-
-    # Verify the error message mentions the factory failure
-    error_message = str(exc_info.value).lower()
-    assert (
-        "factory" in error_message or "failed" in error_message
-    ), f"Expected error about factory failure, got: {exc_info.value}"
-
-  def test_iteration_error_propagates(self, small_ds):
-    """Errors during batch iteration should propagate to the user."""
-    error_on_batch = 2  # Fail on the third batch
-
-    def failing_callback(block, projection_names=None):
-      # Track which batch we're on using a mutable default
-      if not hasattr(failing_callback, "count"):
-        failing_callback.count = 0
-      failing_callback.count += 1
-
-      if failing_callback.count == error_on_batch:
-        raise RuntimeError("Intentional batch processing error")
-
-    # Reset the counter
-    failing_callback.count = 0
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=failing_callback,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # The error should surface when we try to collect
-    with pytest.raises(Exception):
-      ctx.sql("SELECT * FROM test_table").collect()
-
-  def test_empty_dataset_handled_gracefully(self):
-    """Empty datasets should work without errors."""
-    # Create an empty dataset with the right structure
-    empty_ds = xr.Dataset(
-        {
-            "temperature": (
-                ["time", "lat", "lon"],
-                np.array([]).reshape(0, 0, 0),
-            )
-        },
-        coords={
-            "time": pd.DatetimeIndex([]),
-            "lat": np.array([]),
-            "lon": np.array([]),
-        },
-    )
-
-    # This should work without crashing
-    table = read_xarray_table(empty_ds, chunks={"time": 10})
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
-    count = result[0].to_pandas()["cnt"].iloc[0]
-
-    assert count == 0, f"Expected 0 rows for empty dataset, got {count}"
-
-
-class TestMultiplePartitions:
-  """Tests for scenarios with multiple queries and table reuse."""
-
-  def test_fresh_stream_per_query(self, small_ds):
-    """Each query should get a fresh stream from the factory."""
-    call_count = {"value": 0}
-    original_callback = None
-
-    def counting_callback(block, projection_names=None):
-      call_count["value"] += 1
-      if original_callback:
-        original_callback(block)
-
-    table = read_xarray_table(
-        small_ds,
-        chunks={"time": 50},  # 2 blocks per query
-        _iteration_callback=counting_callback,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test_table", table)
-
-    # First query
-    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
-    first_query_count = call_count["value"]
-    assert (
-        first_query_count == 2
-    ), f"First query: expected 2, got {first_query_count}"
-
-    # Second query should trigger fresh iteration
-    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
-    second_query_count = call_count["value"]
-    assert (
-        second_query_count == 4
-    ), f"After second query: expected 4 total, got {second_query_count}"
-
-    # Third query
-    ctx.sql("SELECT MAX(temperature) FROM test_table").collect()
-    third_query_count = call_count["value"]
-    assert (
-        third_query_count == 6
-    ), f"After third query: expected 6 total, got {third_query_count}"
-
-  def test_parallel_queries_independent(self, small_ds):
-    """Multiple contexts with the same table should work independently."""
-    tracker1 = IterationTracker()
-    tracker2 = IterationTracker()
-
-    table1 = read_xarray_table(
-        small_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker1,
-    )
-
-    table2 = read_xarray_table(
-        small_ds,
-        chunks={"time": 50},
-        _iteration_callback=tracker2,
-    )
-
-    ctx1 = SessionContext()
-    ctx2 = SessionContext()
-
-    ctx1.register_table("test_table", table1)
-    ctx2.register_table("test_table", table2)
-
-    # Execute queries
-    ctx1.sql("SELECT COUNT(*) FROM test_table").collect()
-    ctx2.sql("SELECT COUNT(*) FROM test_table").collect()
-
-    # Each should have its own iteration count
-    assert (
-        tracker1.iteration_count == 4
-    ), f"Table1: expected 4 blocks, got {tracker1.iteration_count}"
-    assert (
-        tracker2.iteration_count == 2
-    ), f"Table2: expected 2 blocks, got {tracker2.iteration_count}"
-
-
-class TestFilterPushdown:
-  """Tests for partition pruning via filter pushdown.
-
-  These tests verify that SQL filters on dimension columns (time, lat, lon)
-  correctly prune partitions, reducing the number of partitions read.
-  """
-
-  @pytest.fixture
-  def time_chunked_ds(self):
-    """Dataset chunked by time for pruning tests."""
-    np.random.seed(42)
-    # 100 days of data, chunked into 4 partitions of 25 days each
-    time = pd.date_range("2020-01-01", periods=100, freq="D")
-    lat = np.linspace(-90, 90, 5)
-    data = np.random.rand(100, 5).astype(np.float32)
-
-    return xr.Dataset(
-        {"temperature": (["time", "lat"], data)},
-        coords={"time": time, "lat": lat},
-    )
-
-  def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds):
-    """Query with time > X should skip early partitions."""
-    tracker = IterationTracker()
-
-    # 4 partitions: days 0-24, 25-49, 50-74, 75-99
-    # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9)
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Query only last 25 days (Mar 16+) - should prune first 3 partitions
-    # 2020-03-16 is day 75
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE time >= '2020-03-16'
-    """
-    ).to_pandas()
-
-    # Should read only 1 partition (the last one)
-    assert (
-        tracker.iteration_count == 1
-    ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}"
-
-    # Verify data correctness - 25 days * 5 lat = 125 rows
-    count = result["cnt"].iloc[0]
-    assert count == 125, f"Expected 125 rows, got {count}"
-
-  def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds):
-    """Query with time < X should skip late partitions."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Query only first 25 days (< Jan 26) - should prune last 3 partitions
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE time < '2020-01-26'
-    """
-    ).to_pandas()
-
-    # Should read only 1 partition (the first one)
-    assert (
-        tracker.iteration_count == 1
-    ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}"
-
-    # Verify correctness: Jan 1–25 (25 days) × 5 lat = 125 rows
-    count = result["cnt"].iloc[0]
-    assert count == 125, f"Expected 125 rows, got {count}"
-
-  def test_time_between_filter_prunes_outside_range(self, time_chunked_ds):
-    """Query with BETWEEN should prune partitions outside the range."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 1, 2, and 3
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE time BETWEEN '2020-02-01' AND '2020-03-21'
-    """
-    ).collect()
-
-    # Partition 0 (Jan 1–25) ends before Feb 1 and is pruned.
-    # Partitions 1, 2, 3 each overlap with Feb 1–Mar 21.
-    assert (
-        tracker.iteration_count == 3
-    ), f"Expected exactly 3 partitions after BETWEEN pruning, got {tracker.iteration_count}"
-
-  def test_lat_filter_prunes_partitions(self):
-    """Latitude filter should prune irrelevant partitions."""
-    np.random.seed(42)
-    time = pd.date_range("2020-01-01", periods=10, freq="D")
-    # 100 lat values from -90 to 90
-    lat = np.linspace(-90, 90, 100)
-    data = np.random.rand(10, 100).astype(np.float32)
-
-    ds = xr.Dataset(
-        {"temperature": (["time", "lat"], data)},
-        coords={"time": time, "lat": lat},
-    )
-
-    tracker = IterationTracker()
-
-    # Chunk by latitude: 4 partitions (25 lat values each)
-    # Partition 0: lat -90 to ~-45
-    # Partition 1: lat ~-45 to 0
-    # Partition 2: lat 0 to ~45
-    # Partition 3: lat ~45 to 90
-    table = read_xarray_table(
-        ds,
-        chunks={"lat": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Query southern hemisphere only (lat < 0)
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE lat < 0
-    """
-    ).collect()
-
-    # np.linspace(-90, 90, 100) chunked by 25:
-    #   Partition 0: indices 0–24,  lat -90.0 to -46.4 (all negative)
-    #   Partition 1: indices 25–49, lat -44.5 to  -0.9 (all negative)
-    #   Partition 2: indices 50–74, lat   0.9 to  44.5 (all positive) → pruned
-    #   Partition 3: indices 75–99, lat  46.4 to  90.0 (all positive) → pruned
-    assert (
-        tracker.iteration_count == 2
-    ), f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}"
-
-  def test_no_pruning_for_data_column_filters(self, time_chunked_ds):
-    """Filters on data columns (not dimensions) should not prune."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Filter on temperature (data column), not a dimension
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) FROM test WHERE temperature > 0.5
-    """
-    ).collect()
-
-    # All 4 partitions should be read (can't prune on data column)
-    assert (
-        tracker.iteration_count == 4
-    ), f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}"
-
-  def test_filter_correctness_preserved(self, time_chunked_ds):
-    """Verify filtered results are correct after pruning."""
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Get count with filter
-    filtered = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE time >= '2020-02-15' AND time <= '2020-03-15'
-    """
-    ).to_pandas()
-
-    # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days
-    # 30 days * 5 lat values = 150 rows
-    count = filtered["cnt"].iloc[0]
-    assert count == 150, f"Expected 150 rows, got {count}"
-
-  def test_and_filter_combines_pruning(self, time_chunked_ds):
-    """AND filters should combine for maximum pruning."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Very narrow range that spans only 1 partition
-    result = ctx.sql(
-        """
-        SELECT * FROM test
-        WHERE time >= '2020-03-20' AND time <= '2020-04-05'
-    """
-    ).collect()
-
-    # Should read only 1 partition
-    assert (
-        tracker.iteration_count == 1
-    ), f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}"
-
-  def test_or_filter_is_conservative(self, time_chunked_ds):
-    """OR filters should include partitions matching either condition."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # First or last partition (OR condition)
-    result = ctx.sql(
-        """
-        SELECT * FROM test
-        WHERE time < '2020-01-10' OR time > '2020-03-30'
-    """
-    ).collect()
-
-    # Should read at least 2 partitions (first and last)
-    assert (
-        tracker.iteration_count >= 2
-    ), f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}"
-
-  def test_empty_result_from_impossible_filter(self, time_chunked_ds):
-    """Filter that matches no data should return empty result."""
-    tracker = IterationTracker()
-
-    table = read_xarray_table(
-        time_chunked_ds,
-        chunks={"time": 25},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("test", table)
-
-    # Query for dates outside the data range
-    result = ctx.sql(
-        """
-        SELECT COUNT(*) as cnt FROM test
-        WHERE time > '2025-01-01'
-    """
-    ).to_pandas()
-
-    # Should read 0 partitions (all pruned)
-    assert (
-        tracker.iteration_count == 0
-    ), f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}"
-
-    # Result should be 0 rows
-    count = result["cnt"].iloc[0]
-    assert count == 0, f"Expected 0 rows, got {count}"
-
-
-class TestProjectionPushdown:
-  """Tests that column projection is pushed down to the xarray factory.
-
-  When a query requests only a subset of data variables, the factory should
-  receive only those column names — not the full schema — so that xarray
-  skips loading unrequested variables from disk.
-  """
-
-  @pytest.fixture
-  def two_var_ds(self):
-    """A small dataset with two data variables sharing the same dimensions."""
-    np.random.seed(0)
-    time = pd.date_range("2020-01-01", periods=10, freq="D")
-    lat = np.linspace(-10, 10, 5, dtype=np.float32)
-    shape = (10, 5)
-    return xr.Dataset(
-        {
-            "temperature": (
-                ["time", "lat"],
-                np.random.rand(*shape).astype(np.float32),
-            ),
-            "precipitation": (
-                ["time", "lat"],
-                np.random.rand(*shape).astype(np.float32),
-            ),
-        },
-        coords={"time": time, "lat": lat},
-    )
-
-  def test_single_column_select_projects_only_that_variable(self, two_var_ds):
-    """SELECT on one data variable should pass only that variable to the factory."""
-    tracker = IterationTracker()
-    table = read_xarray_table(
-        two_var_ds,
-        chunks={"time": 5},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("data", table)
-    ctx.sql("SELECT AVG(temperature) FROM data").collect()
-
-    assert (
-        tracker.iteration_count > 0
-    ), "Expected at least one partition to be read"
-    for proj in tracker.projections_seen:
-      assert (
-          proj is not None
-      ), "Expected projection_names to be set for a single-column SELECT, got None"
-      assert (
-          "temperature" in proj
-      ), f"Expected 'temperature' in projection_names, got {proj}"
-      assert (
-          "precipitation" not in proj
-      ), f"'precipitation' should not be loaded for a temperature-only query, got {proj}"
-
-  def test_full_select_includes_all_variables(self, two_var_ds):
-    """SELECT * should include all data variables in the projection."""
-    tracker = IterationTracker()
-    table = read_xarray_table(
-        two_var_ds,
-        chunks={"time": 5},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("data", table)
-    ctx.sql("SELECT * FROM data").collect()
-
-    assert (
-        tracker.iteration_count > 0
-    ), "Expected at least one partition to be read"
-    for proj in tracker.projections_seen:
-      # DataFusion sends all column names explicitly even for SELECT *.
-      # Either None (no pushdown) or a list containing both data variables.
-      if proj is not None:
-        assert "temperature" in proj, f"Missing 'temperature' in {proj}"
-        assert "precipitation" in proj, f"Missing 'precipitation' in {proj}"
-
-  def test_multi_column_select_includes_all_requested(self, two_var_ds):
-    """SELECT on multiple data variables should include all of them in the projection."""
-    tracker = IterationTracker()
-    table = read_xarray_table(
-        two_var_ds,
-        chunks={"time": 5},
-        _iteration_callback=tracker,
-    )
-
-    ctx = SessionContext()
-    ctx.register_table("data", table)
-    ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM data").collect()
-
-    assert tracker.iteration_count > 0
-    for proj in tracker.projections_seen:
-      # Both variables requested — either None (full scan) or both present
-      if proj is not None:
-        assert "temperature" in proj, f"Missing 'temperature' in {proj}"
-        assert "precipitation" in proj, f"Missing 'precipitation' in {proj}"
-
-  def test_projection_result_correctness(self, two_var_ds):
-    """Single-column projected query returns the same result as an unprojected one."""
-    table = read_xarray_table(two_var_ds, chunks={"time": 5})
-    ctx = SessionContext()
-    ctx.register_table("data", table)
-
-    projected = (
-        ctx.sql("SELECT AVG(temperature) as avg_t FROM data")
-        .to_pandas()["avg_t"]
-        .iloc[0]
-    )
-    expected = float(two_var_ds["temperature"].values.mean())
-    assert (
-        abs(projected - expected) < 1e-4
-    ), f"Projected AVG {projected} differs from expected {expected}"
-
-  def test_count_star_passes_none_projection(self, two_var_ds):
-    """COUNT(*) should not push a projection — factory receives None."""
-    projections_seen = []
-
-    def callback(block, projection_names):
-      projections_seen.append(projection_names)
-
-    table = read_xarray_table(
-        two_var_ds,
-        chunks={"time": 5},
-        _iteration_callback=callback,
-    )
-    ctx = SessionContext()
-    ctx.register_table("data", table)
-    result = ctx.sql("SELECT COUNT(*) FROM data").collect()
-
-    total_rows = 10 * 5  # time=10, lat=5
-    assert result[0][0][0].as_py() == total_rows
-    assert all(
-        p is None for p in projections_seen
-    ), f"COUNT(*) should not push a projection, but got: {projections_seen}"
diff --git a/xarray_sql/sql_test.py b/xarray_sql/sql_test.py
deleted file mode 100644
index 8feedee..0000000
--- a/xarray_sql/sql_test.py
+++ /dev/null
@@ -1,194 +0,0 @@
-"""SQL functionality tests for xarray-sql using pytest."""
-
-import pytest
-import xarray as xr
-
-from . import XarrayContext
-from .df_test import create_large_dataset, rand_wx
-
-
-@pytest.fixture
-def air_dataset_small():
-  ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
-  return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
-
-
-@pytest.fixture
-def air_dataset_large():
-  return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
-
-
-@pytest.fixture
-def weather_dataset():
-  ds = rand_wx("2023-01-01T00", "2023-01-01T12")
-  return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
-      {"time": 3}
-  )
-
-
-@pytest.fixture
-def synthetic_dataset():
-  return create_large_dataset(
-      time_steps=50, lat_points=20, lon_points=20
-  ).chunk({"time": 25})
-
-
-@pytest.fixture
-def station_dataset():
-  return xr.Dataset(
-      {
-          "station_id": (["station"], [1, 2, 3, 4, 5]),
-          "elevation": (["station"], [100, 250, 500, 750, 1000]),
-          "name": (
-              ["station"],
-              ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
-          ),
-      }
-  ).chunk({"station": 5})
-
-
-@pytest.fixture
-def air_and_stations():
-  air = (
-      xr.tutorial.open_dataset("air_temperature")
-      .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
-      .chunk({"time": 6})
-  )
-  stations = xr.Dataset(
-      {
-          "station_id": (["station"], [101, 102, 103]),
-          "lat": (
-              ["station"],
-              [air.lat.values[0], air.lat.values[2], air.lat.values[4]],
-          ),
-          "lon": (
-              ["station"],
-              [air.lon.values[1], air.lon.values[3], air.lon.values[5]],
-          ),
-          "elevation": (["station"], [100, 250, 500]),
-      }
-  ).chunk({"station": 3})
-  return air, stations
-
-
-def test_sanity(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  result = ctx.sql(
-      'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100'
-  ).to_pandas()
-  assert len(result) > 0
-  assert len(result) <= 1320
-  assert all(col in result.columns for col in ["lat", "lon", "time", "air"])
-
-
-def test_aggregation_small(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  query = """
-        SELECT lat, lon, SUM(air) AS air_total
-        FROM air
-        GROUP BY lat, lon
-    """
-  result = ctx.sql(query).to_pandas()
-  expected_rows = (
-      air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"]
-  )
-  assert len(result) == expected_rows
-
-
-def test_aggregation_large(air_dataset_large):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_large)
-  query = """
-        SELECT lat, lon, AVG(air) AS air_avg
-        FROM air
-        GROUP BY lat, lon
-    """
-  result = ctx.sql(query).to_pandas()
-  expected_rows = (
-      air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"]
-  )
-  assert len(result) == expected_rows
-
-
-def test_basic_select_all(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas()
-  assert len(result) <= 10
-  for col in ["lat", "lon", "time", "air"]:
-    assert col in result.columns
-
-
-def test_weather_queries(weather_dataset):
-  ctx = XarrayContext()
-  ctx.from_dataset("weather", weather_dataset)
-  # Selecting specific columns
-  result = ctx.sql(
-      "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20"
-  ).to_pandas()
-  assert "temperature" in result.columns
-  assert "precipitation" in result.columns
-  # Filtering
-  result = ctx.sql(
-      "SELECT * FROM weather WHERE temperature > 10 LIMIT 50"
-  ).to_pandas()
-  assert len(result) > 0
-  assert (result["temperature"] > 10).all()
-
-
-def test_synthetic_aggregations(synthetic_dataset):
-  ctx = XarrayContext()
-  ctx.from_dataset("synthetic", synthetic_dataset)
-  # COUNT aggregation
-  result = ctx.sql("SELECT COUNT(*) AS total_count FROM synthetic").to_pandas()
-  assert result["total_count"].iloc[0] > 0
-  # MIN, MAX, AVG
-  query = """
-        SELECT MIN(temperature) AS min_temp,
-               MAX(temperature) AS max_temp,
-               AVG(temperature) AS avg_temp
-        FROM synthetic
-    """
-  result = ctx.sql(query).to_pandas()
-  assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0]
-  assert (
-      result["min_temp"].iloc[0]
-      <= result["avg_temp"].iloc[0]
-      <= result["max_temp"].iloc[0]
-  )
-
-
-def test_invalid_table_name(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  with pytest.raises(Exception):
-    ctx.sql("SELECT * FROM nonexistent_table")
-
-
-def test_invalid_column_name(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  with pytest.raises(Exception):
-    ctx.sql("SELECT nonexistent_column FROM air")
-
-
-def test_sql_syntax_error(air_dataset_small):
-  ctx = XarrayContext()
-  ctx.from_dataset("air", air_dataset_small)
-  with pytest.raises(Exception):
-    ctx.sql("SELECT * FORM air")  # Typo: FORM instead of FROM
-  with pytest.raises(Exception):
-    ctx.sql("SELECT * FROM air WHERE")  # Incomplete WHERE
-
-
-def test_cross_join(air_and_stations):
-  air, stations = air_and_stations
-  ctx = XarrayContext()
-  ctx.from_dataset("air_data", air)
-  ctx.from_dataset("stations", stations)
-  result = ctx.sql(
-      "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations"
-  ).to_pandas()
-  assert result["total"].iloc[0] > 0

From 329d3f7a41a61bd96408aa8dc6480a7ab942ce23 Mon Sep 17 00:00:00 2001
From: Andrew Huang <ahuang@anaconda.com>
Date: Fri, 20 Mar 2026 15:49:27 -0700
Subject: [PATCH 3/3] Add tests

---
 pyproject.toml       |    1 -
 tests/__init__.py    |    0
 tests/conftest.py    |  138 +++++
 tests/test_df.py     |  420 +++++++++++++
 tests/test_reader.py | 1372 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_sql.py    |  129 ++++
 6 files changed, 2059 insertions(+), 1 deletion(-)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/conftest.py
 create mode 100644 tests/test_df.py
 create mode 100644 tests/test_reader.py
 create mode 100644 tests/test_sql.py

diff --git a/pyproject.toml b/pyproject.toml
index 2701dcc..78e289d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,7 +37,6 @@ dependencies = [
 
 [project.optional-dependencies]
 test = [
-  "cftime>",
   "pytest",
   "xarray[io]",
   "gcsfs",
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..b173f2e
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,138 @@
+import pytest
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def rand_wx(start: str, end: str) -> xr.Dataset:
+  np.random.seed(42)
+  lat = np.linspace(-90, 90, num=720)
+  lon = np.linspace(-180, 180, num=1440)
+  time = pd.date_range(start, end, freq="h")
+  level = np.array([1000, 500], dtype=np.int32)
+  reference_time = pd.Timestamp(start)
+  temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
+  precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
+  return xr.Dataset(
+      data_vars=dict(
+          temperature=(["lat", "lon", "time", "level"], temperature),
+          precipitation=(["lat", "lon", "time", "level"], precipitation),
+      ),
+      coords=dict(
+          lat=lat,
+          lon=lon,
+          time=time,
+          level=level,
+          reference_time=reference_time,
+      ),
+      attrs=dict(description="Random weather."),
+  )
+
+
+def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
+  """Create a large xarray dataset for memory testing."""
+  np.random.seed(42)
+
+  time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
+  lat = np.linspace(-90, 90, lat_points)
+  lon = np.linspace(-180, 180, lon_points)
+
+  temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
+  precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
+
+  return xr.Dataset(
+      {
+          "temperature": (["time", "lat", "lon"], temp_data),
+          "precipitation": (["time", "lat", "lon"], precip_data),
+      },
+      coords={"time": time, "lat": lat, "lon": lon},
+  )
+
+
+@pytest.fixture
+def air():
+  ds = xr.tutorial.open_dataset("air_temperature")
+  chunks = {"time": 240}
+  return ds.chunk(chunks)
+
+
+@pytest.fixture
+def air_small(air):
+  return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
+      {"time": 240}
+  )
+
+
+@pytest.fixture
+def randwx():
+  return rand_wx("1995-01-13T00", "1995-01-13T01")
+
+
+@pytest.fixture
+def large_ds():
+  return create_large_dataset().chunk({"time": 25})
+
+
+@pytest.fixture
+def air_dataset_small():
+  ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+  return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
+
+
+@pytest.fixture
+def air_dataset_large():
+  return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+
+
+@pytest.fixture
+def weather_dataset():
+  ds = rand_wx("2023-01-01T00", "2023-01-01T12")
+  return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
+      {"time": 3}
+  )
+
+
+@pytest.fixture
+def synthetic_dataset():
+  return create_large_dataset(
+      time_steps=50, lat_points=20, lon_points=20
+  ).chunk({"time": 25})
+
+
+@pytest.fixture
+def station_dataset():
+  return xr.Dataset(
+      {
+          "station_id": (["station"], [1, 2, 3, 4, 5]),
+          "elevation": (["station"], [100, 250, 500, 750, 1000]),
+          "name": (
+              ["station"],
+              ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
+          ),
+      }
+  ).chunk({"station": 5})
+
+
+@pytest.fixture
+def air_and_stations():
+  air = (
+      xr.tutorial.open_dataset("air_temperature")
+      .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
+      .chunk({"time": 6})
+  )
+  stations = xr.Dataset(
+      {
+          "station_id": (["station"], [101, 102, 103]),
+          "lat": (
+              ["station"],
+              [air.lat.values[0], air.lat.values[2], air.lat.values[4]],
+          ),
+          "lon": (
+              ["station"],
+              [air.lon.values[1], air.lon.values[3], air.lon.values[5]],
+          ),
+          "elevation": (["station"], [100, 250, 500]),
+      }
+  ).chunk({"station": 3})
+  return air, stations
diff --git a/tests/test_df.py b/tests/test_df.py
new file mode 100644
index 0000000..37380ae
--- /dev/null
+++ b/tests/test_df.py
@@ -0,0 +1,420 @@
+import tracemalloc
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+import xarray as xr
+
+from xarray_sql.df import (
+    DEFAULT_BATCH_SIZE,
+    _parse_schema,
+    block_slices,
+    dataset_to_record_batch,
+    explode,
+    from_map,
+    from_map_batched,
+    iter_record_batches,
+    pivot,
+)
+from xarray_sql.reader import read_xarray, read_xarray_table
+
+
+def test_explode_cardinality(air):
+  dss = explode(air)
+  assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()])
+
+
+def test_explode_dim_sizes_one(air):
+  chunks = {"time": 240}
+  ds = next(iter(explode(air)))
+  for k, v in chunks.items():
+    assert k in ds.dims
+    assert v == ds.sizes[k]
+
+
+def test_explode_data_equal_one_first(air):
+  ds = next(iter(explode(air)))
+  iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()}
+  assert air.isel(iselection).equals(ds)
+
+
+def test_explode_data_equal_one_last(air):
+  dss = list(explode(air))
+  ds = dss[-1]
+
+  # For the last chunk, we need to calculate where it actually starts
+  # The original logic slice(0, s) only works for the first chunk
+  iselection = {}
+  for dim in ds.dims:
+    # Get chunk boundaries
+    chunk_bounds = np.cumsum((0,) + air.chunks[dim])
+    # Last chunk index
+    last_chunk_idx = len(air.chunks[dim]) - 1
+    # Calculate actual start and end positions
+    start = chunk_bounds[last_chunk_idx]
+    end = chunk_bounds[last_chunk_idx + 1]
+    iselection[dim] = slice(start, end)
+
+  assert air.isel(iselection).equals(ds)
+
+
+def test_from_map_basic():
+  def make_df(x):
+    return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]})
+
+  result = from_map(make_df, [1, 2, 3])
+  assert isinstance(result, pa.Table)
+  assert len(result) == 6
+  assert result.column_names == ["value", "index"]
+
+
+def test_from_map_multiple_iterables():
+  def add_values(x, y):
+    return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]})
+
+  result = from_map(add_values, [1, 2], [10, 20])
+  assert isinstance(result, pa.Table)
+  assert len(result) == 2
+
+  df = result.to_pandas()
+  assert list(df["sum"]) == [11, 22]
+
+
+def test_from_map_with_args():
+  def multiply_and_add(x, multiplier, add_value):
+    return pd.DataFrame({"result": [x * multiplier + add_value]})
+
+  result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10))
+  assert isinstance(result, pa.Table)
+  assert len(result) == 3
+
+  df = result.to_pandas()
+  assert list(df["result"]) == [12, 14, 16]
+
+
+def test_from_map_with_pyarrow_tables():
+  def make_arrow_table(x):
+    df = pd.DataFrame({"value": [x]})
+    return pa.Table.from_pandas(df)
+
+  result = from_map(make_arrow_table, [1, 2, 3])
+  assert isinstance(result, pa.Table)
+  assert len(result) == 3
+
+
+def test_iter_record_batches_splits_into_multiple_batches(air_small):
+  """iter_record_batches should emit >1 batch when partition exceeds batch_size."""
+  schema = _parse_schema(air_small)
+  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+  ds_block = air_small.isel(block)
+  total_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes]))
+
+  small_batch = 16  # force many small batches
+  batches = list(iter_record_batches(ds_block, schema, batch_size=small_batch))
+
+  assert len(batches) == -(-total_rows // small_batch)  # ceiling division
+  assert all(b.num_rows <= small_batch for b in batches)
+  assert sum(b.num_rows for b in batches) == total_rows
+
+
+def test_iter_record_batches_matches_dataset_to_record_batch(air_small):
+  """Concatenating all iter_record_batches output must equal dataset_to_record_batch."""
+  schema = _parse_schema(air_small)
+  dim_cols = [f.name for f in schema if f.name in air_small.dims]
+  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+  ds_block = air_small.isel(block)
+
+  batches = list(iter_record_batches(ds_block, schema, batch_size=16))
+  actual_df = (
+      pa.Table.from_batches(batches)
+      .to_pandas()
+      .sort_values(dim_cols)
+      .reset_index(drop=True)
+  )
+  expected_df = (
+      dataset_to_record_batch(ds_block, schema)
+      .to_pandas()
+      .sort_values(dim_cols)
+      .reset_index(drop=True)
+  )
+  pd.testing.assert_frame_equal(actual_df, expected_df)
+
+
+def test_iter_record_batches_default_batch_size():
+  """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch."""
+  ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2))
+  schema = _parse_schema(ds)
+  total_rows = int(np.prod([ds.sizes[d] for d in ds.sizes]))
+  assert total_rows <= DEFAULT_BATCH_SIZE, "fixture too large — adjust isel"
+  batches = list(iter_record_batches(ds, schema))
+  assert len(batches) == 1
+  assert batches[0].num_rows == total_rows
+
+
+def test_dataset_to_record_batch_matches_pivot(air_small):
+  """dataset_to_record_batch should contain the same rows as pivot.
+
+  Row ordering may differ (pivot uses ds.dims key order; dataset_to_record_batch
+  uses the data variable's own dim order). Both orderings are valid for SQL, so
+  we sort by the coordinate columns before comparing.
+  """
+  schema = _parse_schema(air_small)
+  dim_cols = [f.name for f in schema if f.name in air_small.dims]
+  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+
+  for block in blocks:
+    ds_block = air_small.isel(block)
+    actual_df = (
+        dataset_to_record_batch(ds_block, schema)
+        .to_pandas()
+        .sort_values(dim_cols)
+        .reset_index(drop=True)
+    )
+    expected_df = (
+        pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema)
+        .to_pandas()
+        .sort_values(dim_cols)
+        .reset_index(drop=True)
+    )
+
+    pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False)
+
+
+def test_dataset_to_record_batch_column_order(air_small):
+  """Output column order must match schema (dims first, then data vars)."""
+  schema = _parse_schema(air_small)
+  block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+  batch = dataset_to_record_batch(air_small.isel(block), schema)
+  assert batch.schema.names == schema.names
+
+
+def test_dataset_to_record_batch_row_count(air_small):
+  """Row count must equal the product of the block dimension sizes."""
+  schema = _parse_schema(air_small)
+  chunks = {"time": 4, "lat": 3, "lon": 4}
+  for block in block_slices(air_small, chunks=chunks):
+    ds_block = air_small.isel(block)
+    expected_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes]))
+    batch = dataset_to_record_batch(ds_block, schema)
+    assert batch.num_rows == expected_rows
+
+
+def test_from_map_batched_basic_functionality(air_small):
+  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}))
+
+  first_block_df = pivot(air_small.isel(blocks[0]))
+  expected_schema = pa.Schema.from_pandas(first_block_df)
+
+  reader = from_map_batched(
+      pivot, [air_small.isel(block) for block in blocks], schema=expected_schema
+  )
+
+  assert isinstance(reader, pa.RecordBatchReader)
+  assert reader.schema == expected_schema
+
+  batches = list(reader)
+  assert len(batches) > 0
+  for batch in batches:
+    assert batch.schema == expected_schema
+    assert len(batch) > 0
+
+
+def adding_function(x, y):
+  """Simple function that adds two values and returns a DataFrame."""
+  result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
+  return result
+
+
+def test_from_map_batched_multiple_iterables():
+  x_values = [1, 2, 3, 4, 5]
+  y_values = [10, 20, 30, 40, 50]
+
+  expected_schema = pa.schema(
+      [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())]
+  )
+
+  reader = from_map_batched(
+      adding_function, x_values, y_values, schema=expected_schema
+  )
+  table = reader.read_all()
+  df = table.to_pandas()
+
+  expected_df = pd.DataFrame(
+      {
+          "x": x_values,
+          "y": y_values,
+          "sum": [x + y for x, y in zip(x_values, y_values)],
+      }
+  )
+  pd.testing.assert_frame_equal(df, expected_df)
+
+
+def test_from_map_batched_with_args_and_kwargs():
+  def multiply_and_add(x, multiplier, offset=0):
+    return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]})
+
+  values = [1, 2, 3]
+  expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())])
+
+  reader = from_map_batched(
+      multiply_and_add, values, args=(2,), offset=5, schema=expected_schema
+  )
+  table = reader.read_all()
+  df = table.to_pandas()
+
+  expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]})
+  pd.testing.assert_frame_equal(df, expected_df)
+
+
+def test_from_map_batched_empty_iterables():
+  empty_schema = pa.schema([("value", pa.int64())])
+
+  reader = from_map_batched(
+      lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema
+  )
+  batches = list(reader)
+  assert len(batches) == 0
+
+
+def test_from_map_batched_consistency_with_regular_map(air_small):
+  blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3}))
+  datasets = [air_small.isel(block) for block in blocks]
+
+  first_df = pivot(datasets[0])
+  schema = pa.Schema.from_pandas(first_df)
+
+  reader = from_map_batched(pivot, datasets, schema=schema)
+  batched_table = reader.read_all()
+
+  regular_dfs = [pivot(ds) for ds in datasets]
+  regular_table = pa.Table.from_pandas(
+      pd.concat(regular_dfs, ignore_index=True)
+  )
+
+  assert batched_table.schema == regular_table.schema
+  assert len(batched_table) == len(regular_table)
+
+  batched_df = (
+      batched_table.to_pandas()
+      .sort_values(["time", "lat", "lon"])
+      .reset_index(drop=True)
+  )
+  regular_df = (
+      regular_table.to_pandas()
+      .sort_values(["time", "lat", "lon"])
+      .reset_index(drop=True)
+  )
+
+  pd.testing.assert_frame_equal(batched_df, regular_df)
+
+
+def test_from_map_batched_integration_with_datafusion_via_read_xarray():
+  air = xr.tutorial.open_dataset("air_temperature")
+  air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15))
+  air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8})
+
+  arrow_stream = read_xarray(
+      air_chunked, chunks={"time": 25, "lat": 5, "lon": 8}
+  )
+
+  assert hasattr(arrow_stream, "schema")
+  assert hasattr(arrow_stream, "__iter__")
+
+  table = arrow_stream.read_all()
+  assert len(table) > 0
+
+  expected_columns = {"time", "lat", "lon", "air"}
+  actual_columns = set(table.column_names)
+  assert expected_columns.issubset(actual_columns)
+
+
+def test_read_xarray_loads_one_chunk_at_a_time(large_ds):
+  tracemalloc.start()
+  iterable = read_xarray(large_ds)
+  first_size, first_peak = tracemalloc.get_traced_memory()
+  tracemalloc.reset_peak()
+
+  sizes, peaks = [], []
+
+  first_chunk = large_ds.isel(next(block_slices(large_ds)))
+  chunk_size = first_chunk.nbytes
+
+  # Creating the iterator should be inexpensive -- less than one chunk.
+  # We multiply by constant factors because chunks have additional overhead
+  assert first_size < chunk_size * 3
+  assert first_peak < chunk_size * 6
+
+  for it in iterable:
+    _ = it
+    cur_size, cur_peak = tracemalloc.get_traced_memory()
+    tracemalloc.reset_peak()
+    sizes.append(cur_size)
+    peaks.append(cur_peak)
+
+  for size in sizes:
+    # Observed range: 1.59–1.83× chunk_size.
+    # iter_record_batches holds data-variable arrays (≈1× chunk) while
+    # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk).
+    assert chunk_size * 1.3 < size, f"size {size} unexpectedly low"
+    assert chunk_size * 2.2 > size, f"size {size} unexpectedly high"
+
+  for peak in peaks:
+    # Observed range: 1.84–3.28× chunk_size.
+    # Peak includes data arrays + Arrow batch + temporary coordinate index
+    # arrays; the first batch of each chunk is highest (Dask compute overhead).
+    assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low"
+    assert chunk_size * 4.0 > peak, f"peak {peak} unexpectedly high"
+
+  assert max(peaks) < large_ds.nbytes
+
+  tracemalloc.stop()
+
+
+def test_read_xarray_table_memory_bounds(large_ds):
+  """read_xarray_table should not materialise data at registration time.
+
+  Registration should only hold coordinate arrays and Rust partition metadata
+  (no data variables).  Peak memory during a full-table query should be a
+  small fraction of the whole dataset — i.e. partitions are processed without
+  loading all of them simultaneously.
+  """
+  from datafusion import SessionContext
+
+  first_chunk = large_ds.isel(next(block_slices(large_ds)))
+  chunk_size = first_chunk.nbytes
+
+  # --- Registration phase ---
+  tracemalloc.start()
+  table = read_xarray_table(large_ds)
+  reg_size, reg_peak = tracemalloc.get_traced_memory()
+  tracemalloc.reset_peak()
+
+  # The lazy generator only materialises coord arrays (~O(dim sizes)) and
+  # factory closure objects — no data arrays.  Both metrics should be well
+  # below one chunk of data.
+  assert reg_size < chunk_size, (
+      f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: "
+      "data may have been loaded eagerly"
+  )
+  assert (
+      reg_peak < chunk_size * 2
+  ), f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})"
+
+  # --- Query phase ---
+  ctx = SessionContext()
+  ctx.register_table("weather", table)
+  ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM weather").collect()
+  _, query_peak = tracemalloc.get_traced_memory()
+
+  # tracemalloc measures Python-heap allocations, which include Arrow
+  # buffer copies and object overhead on top of the raw data.  The
+  # observed peak is typically 1.1–1.5× the raw dataset size; we use
+  # 2× as a generous bound that would still catch catastrophic regressions
+  # (e.g. loading all partitions twice simultaneously).
+  assert query_peak < large_ds.nbytes * 2, (
+      f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: "
+      "may be holding excessive data in memory"
+  )
+
+  tracemalloc.stop()
diff --git a/tests/test_reader.py b/tests/test_reader.py
new file mode 100644
index 0000000..d970b4c
--- /dev/null
+++ b/tests/test_reader.py
@@ -0,0 +1,1372 @@
+"""Tests for XarrayRecordBatchReader lazy streaming behavior.
+
+These tests verify that XarrayRecordBatchReader provides true lazy evaluation:
+- No data iteration during reader creation
+- No data iteration during DataFusion table registration (using LazyArrowStreamTable)
+- Data iteration ONLY occurs during query execution (collect())
+
+The lazy streaming is achieved via the Rust LazyArrowStreamTable class which
+implements the __datafusion_table_provider__ protocol using StreamingTable.
+
+Additional tests verify:
+- True streaming with bounded memory (batches processed incrementally)
+- Back-pressure behavior (producer pauses when consumer is slow)
+- Error propagation through the stream
+"""
+
+import threading
+import time
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+import xarray as xr
+from datafusion import SessionContext
+
+from xarray_sql._native import LazyArrowStreamTable
+from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table
+from xarray_sql.df import _parse_schema
+
+
+@pytest.fixture
+def small_ds():
+  """Create a small dataset for testing."""
+  np.random.seed(42)
+  time = pd.date_range("2020-01-01", periods=100, freq="h")
+  lat = np.linspace(-90, 90, 10)
+  lon = np.linspace(-180, 180, 10)
+
+  data = np.random.rand(100, 10, 10).astype(np.float32)
+
+  return xr.Dataset(
+      {"temperature": (["time", "lat", "lon"], data)},
+      coords={"time": time, "lat": lat, "lon": lon},
+  )
+
+
+class IterationTracker:
+  """Tracks when iteration occurs for testing lazy evaluation.
+
+  The callback signature is ``(block, projection_names)`` where
+  ``projection_names`` is the list of column names requested by the query
+  (``None`` when no projection pushdown occurred, e.g. for
+  ``XarrayRecordBatchReader`` or a ``SELECT *`` query).
+  """
+
+  def __init__(self):
+    self.iteration_count = 0
+    self.blocks_seen = []
+    self.projections_seen = []
+
+  def __call__(self, block, projection_names=None):
+    self.iteration_count += 1
+    self.blocks_seen.append(block)
+    self.projections_seen.append(projection_names)
+
+  def reset(self):
+    self.iteration_count = 0
+    self.blocks_seen = []
+    self.projections_seen = []
+
+
+class TestXarrayRecordBatchReaderCreation:
+  """Tests that reader creation does NOT trigger data iteration."""
+
+  def test_reader_creation_does_not_iterate(self, small_ds):
+    """Creating a reader should NOT iterate through any data."""
+    tracker = IterationTracker()
+
+    reader = XarrayRecordBatchReader(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    assert tracker.iteration_count == 0, (
+        f"Expected 0 iterations during reader creation, "
+        f"but got {tracker.iteration_count}"
+    )
+
+  def test_schema_access_does_not_iterate(self, small_ds):
+    """Accessing the schema should NOT trigger iteration."""
+    tracker = IterationTracker()
+
+    reader = XarrayRecordBatchReader(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    # Access schema
+    _ = reader.schema
+    _ = reader.__arrow_c_schema__()
+
+    assert tracker.iteration_count == 0, (
+        f"Expected 0 iterations when accessing schema, "
+        f"but got {tracker.iteration_count}"
+    )
+
+
+class TestDataFusionRegistration:
+  """Tests that DataFusion table registration does NOT trigger iteration.
+
+  These tests use read_xarray_table with register_table()
+  to achieve true lazy evaluation.
+  """
+
+  def test_register_table_does_not_iterate(self, small_ds):
+    """Registering a LazyArrowStreamTable should NOT iterate data.
+
+    This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps
+    a factory and implements __datafusion_table_provider__ with StreamingTable,
+    ensuring data is only read during query execution.
+    """
+    tracker = IterationTracker()
+
+    # Use read_xarray_table which creates a factory-based table
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    assert tracker.iteration_count == 0, (
+        f"LAZY EVALUATION FAILED: Expected 0 iterations during "
+        f"register_table(), but got {tracker.iteration_count}."
+    )
+
+  def test_sql_planning_does_not_iterate(self, small_ds):
+    """Creating a SQL query plan should NOT iterate data."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Create a query but don't execute it
+    query = ctx.sql("SELECT AVG(temperature) FROM test_table")
+
+    # Just creating the query shouldn't iterate
+    assert tracker.iteration_count == 0, (
+        f"Expected 0 iterations during SQL planning, "
+        f"but got {tracker.iteration_count}. "
+        f"DataFusion may be scanning data during query planning."
+    )
+
+
+class TestDataFusionCollect:
+  """Tests that data iteration ONLY occurs during collect().
+
+  These tests use read_xarray_table to verify lazy evaluation.
+  """
+
+  def test_collect_triggers_iteration(self, small_ds):
+    """collect() should trigger data iteration."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Verify no iteration yet (lazy registration)
+    iteration_before_collect = tracker.iteration_count
+    assert (
+        iteration_before_collect == 0
+    ), "Should have 0 iterations before collect"
+
+    # Now collect - this SHOULD iterate
+    result = ctx.sql("SELECT * FROM test_table LIMIT 10").collect()
+
+    assert tracker.iteration_count > 0, (
+        f"Expected iterations during collect(), but got 0. "
+        f"Data was never read!"
+    )
+    assert (
+        tracker.iteration_count > iteration_before_collect
+    ), f"Expected more iterations after collect()"
+
+  def test_full_query_iterates_all_blocks(self, small_ds):
+    """A query that reads all data should iterate all blocks."""
+    tracker = IterationTracker()
+
+    chunks = {"time": 25}
+    table = read_xarray_table(
+        small_ds,
+        chunks=chunks,
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Run a query that needs to scan all data
+    result = ctx.sql("SELECT COUNT(*) FROM test_table").collect()
+
+    # With time=100 and chunks=25, we expect 4 blocks
+    expected_blocks = 100 // 25
+    assert tracker.iteration_count == expected_blocks, (
+        f"Expected {expected_blocks} block iterations, "
+        f"but got {tracker.iteration_count}"
+    )
+
+  def test_aggregation_query_iterates_correctly(self, small_ds):
+    """Aggregation queries should iterate all necessary blocks."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Run aggregation
+    result = ctx.sql(
+        "SELECT lat, AVG(temperature) as avg_temp "
+        "FROM test_table GROUP BY lat"
+    ).collect()
+
+    # Should have iterated some blocks
+    assert tracker.iteration_count > 0
+    assert len(result) > 0
+
+
+class TestLazyEvaluationEndToEnd:
+  """End-to-end tests verifying lazy evaluation through the full pipeline.
+
+  These tests use read_xarray_table to achieve true lazy evaluation.
+  """
+
+  def test_lazy_evaluation_sequence(self, small_ds):
+    """Verify the exact sequence of lazy evaluation stages.
+
+    This is the comprehensive test that proves true lazy evaluation:
+    1. Table creation: 0 iterations
+    2. Table registration: 0 iterations
+    3. Query planning: 0 iterations
+    4. collect(): N iterations (where N = number of blocks)
+    """
+    tracker = IterationTracker()
+
+    # Stage 1: Table creation (with factory)
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+    iterations_after_table = tracker.iteration_count
+    assert iterations_after_table == 0, (
+        f"Stage 1 FAILED: Table creation triggered "
+        f"{iterations_after_table} iterations"
+    )
+
+    # Stage 2: Table registration
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+    iterations_after_registration = tracker.iteration_count
+    assert iterations_after_registration == 0, (
+        f"Stage 2 FAILED: Table registration triggered "
+        f"{iterations_after_registration} iterations"
+    )
+
+    # Stage 3: Query planning
+    query = ctx.sql("SELECT * FROM test_table")
+    iterations_after_planning = tracker.iteration_count
+    assert iterations_after_planning == 0, (
+        f"Stage 3 FAILED: Query planning triggered "
+        f"{iterations_after_planning} iterations"
+    )
+
+    # Stage 4: collect() - NOW iteration should happen
+    result = query.collect()
+    iterations_after_collect = tracker.iteration_count
+    assert (
+        iterations_after_collect > 0
+    ), f"Stage 4 FAILED: collect() triggered 0 iterations - no data was read!"
+
+    # Verify we got the expected number of blocks (100 time steps / 25 = 4)
+    expected_blocks = 4
+    assert (
+        iterations_after_collect == expected_blocks
+    ), f"Expected {expected_blocks} iterations, got {iterations_after_collect}"
+
+  def test_multiple_queries_on_same_table(self, small_ds):
+    """Same table can be queried multiple times with fresh iteration each time."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 50},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # First query
+    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
+    first_query_iterations = tracker.iteration_count
+    assert first_query_iterations > 0, "First query should iterate"
+
+    # Second query on same table - should iterate again
+    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
+    second_query_iterations = tracker.iteration_count
+    assert (
+        second_query_iterations > first_query_iterations
+    ), "Second query should trigger additional iterations"
+
+  def test_stream_consumed_error(self, small_ds):
+    """Once consumed, a single XarrayRecordBatchReader should not be reusable."""
+    reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25})
+
+    # Consume the reader by converting to a PyArrow reader and reading
+    import pyarrow as pa
+
+    pa_reader = pa.RecordBatchReader.from_stream(reader)
+    _ = pa_reader.read_all()
+
+    # Reader is now consumed, calling __arrow_c_stream__ again should fail
+    with pytest.raises(RuntimeError, match="already consumed"):
+      reader.__arrow_c_stream__()
+
+
+class TestDataIntegrity:
+  """Tests that verify data correctness alongside lazy evaluation.
+
+  These tests use read_xarray_table for lazy streaming.
+  """
+
+  def test_query_results_are_correct(self, small_ds):
+    """Verify that lazy evaluation produces correct results."""
+    table = read_xarray_table(small_ds, chunks={"time": 25})
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Get count
+    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
+    count = result[0].to_pandas()["cnt"].iloc[0]
+
+    # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows
+    expected_count = 100 * 10 * 10
+    assert (
+        count == expected_count
+    ), f"Expected {expected_count} rows, got {count}"
+
+  def test_aggregation_results_are_correct(self, small_ds):
+    """Verify aggregation produces correct results."""
+    table = read_xarray_table(small_ds, chunks={"time": 25})
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Get average temperature
+    result = ctx.sql(
+        "SELECT AVG(temperature) as avg_temp FROM test_table"
+    ).collect()
+    avg_temp = result[0].to_pandas()["avg_temp"].iloc[0]
+
+    # With seed 42 and random data in [0, 1), average should be ~0.5
+    assert (
+        0.4 < avg_temp < 0.6
+    ), f"Expected average temperature ~0.5, got {avg_temp}"
+
+
+class TestPyArrowInterop:
+  """Tests for PyArrow interoperability."""
+
+  def test_from_stream_does_not_iterate(self, small_ds):
+    """pa.RecordBatchReader.from_stream() should not iterate."""
+    tracker = IterationTracker()
+
+    reader = XarrayRecordBatchReader(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    # Create PyArrow reader from our stream
+    pa_reader = pa.RecordBatchReader.from_stream(reader)
+
+    assert tracker.iteration_count == 0, (
+        f"Expected 0 iterations when creating PyArrow reader, "
+        f"but got {tracker.iteration_count}"
+    )
+
+  def test_pyarrow_iteration_triggers_callbacks(self, small_ds):
+    """Iterating via PyArrow should trigger our callbacks."""
+    tracker = IterationTracker()
+
+    reader = XarrayRecordBatchReader(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    pa_reader = pa.RecordBatchReader.from_stream(reader)
+
+    # Now iterate
+    for batch in pa_reader:
+      pass
+
+    assert (
+        tracker.iteration_count == 4
+    ), f"Expected 4 iterations, got {tracker.iteration_count}"
+
+  def test_read_all_iterates_all(self, small_ds):
+    """read_all() should iterate through all blocks."""
+    tracker = IterationTracker()
+
+    reader = XarrayRecordBatchReader(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    pa_reader = pa.RecordBatchReader.from_stream(reader)
+    table = pa_reader.read_all()
+
+    assert tracker.iteration_count == 4
+    assert len(table) == 100 * 10 * 10
+
+
+class StreamingTracker:
+  """Tracks timing of batch iterations to verify streaming behavior.
+
+  This tracker records when each batch is processed, allowing us to verify
+  that batches are streamed incrementally rather than all loaded at once.
+  """
+
+  def __init__(self):
+    self.batch_times = []
+    self.batch_count = 0
+    self._lock = threading.Lock()
+
+  def __call__(self, block, projection_names=None):
+    with self._lock:
+      self.batch_times.append(time.monotonic())
+      self.batch_count += 1
+
+  def reset(self):
+    with self._lock:
+      self.batch_times = []
+      self.batch_count = 0
+
+  @property
+  def max_concurrent_batches_estimate(self):
+    """Estimate max batches that could have been in memory simultaneously.
+
+    If all batches are loaded at once, all batch_times will be very close.
+    If streaming works correctly, batch_times should be spread out.
+    """
+    if len(self.batch_times) < 2:
+      return len(self.batch_times)
+
+    # Sort times and look at gaps
+    sorted_times = sorted(self.batch_times)
+    # If times are spread out, streaming is working
+    # If all times are within a tiny window, all batches loaded at once
+    total_duration = sorted_times[-1] - sorted_times[0]
+
+    # If the spread is very small compared to number of batches,
+    # batches were likely all loaded at once
+    return len(self.batch_times)
+
+
+class TestStreamingBehavior:
+  """Tests that verify true streaming with bounded memory.
+
+  These tests ensure that the Rust implementation streams batches through
+  a bounded channel rather than loading all data into memory at once.
+  """
+
+  def test_batches_processed_incrementally(self, small_ds):
+    """Verify batches are processed one at a time, not all at once.
+
+    This test uses a callback that tracks when each batch is processed.
+    With true streaming, batches should be processed incrementally.
+    """
+    tracker = StreamingTracker()
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Run query that scans all data
+    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
+
+    # All 4 batches should have been processed
+    assert (
+        tracker.batch_count == 4
+    ), f"Expected 4 batches, got {tracker.batch_count}"
+
+  def test_all_partitions_processed(self, small_ds):
+    """Verify that all partitions are processed (order may vary with parallelism)."""
+    blocks_seen = []
+
+    def track_order(block, projection_names=None):
+      # Record the time slice for ordering verification
+      blocks_seen.append(block.get("time", None))
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=track_order,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+    ctx.sql("SELECT * FROM test_table").collect()
+
+    # Should have 4 blocks/partitions
+    assert len(blocks_seen) == 4
+
+    # All blocks should be present (though order may vary due to parallelism)
+    # Extract start positions and verify they cover all expected ranges
+    starts = sorted([b.start for b in blocks_seen])
+    expected_starts = [0, 25, 50, 75]
+    assert (
+        starts == expected_starts
+    ), f"Expected partition starts {expected_starts}, got {starts}"
+
+  def test_large_dataset_streams_correctly(self):
+    """Test streaming with a larger dataset to verify memory behavior.
+
+    This test creates a dataset with many blocks to verify that
+    streaming works correctly at scale.
+    """
+    # Create a dataset with 20 blocks
+    np.random.seed(42)
+    time = pd.date_range("2020-01-01", periods=200, freq="h")
+    lat = np.linspace(-90, 90, 10)
+    lon = np.linspace(-180, 180, 10)
+
+    data = np.random.rand(200, 10, 10).astype(np.float32)
+
+    large_ds = xr.Dataset(
+        {"temperature": (["time", "lat", "lon"], data)},
+        coords={"time": time, "lat": lat, "lon": lon},
+    )
+
+    tracker = StreamingTracker()
+
+    # Use small chunks to create many blocks
+    table = read_xarray_table(
+        large_ds,
+        chunks={"time": 10},  # 200 / 10 = 20 blocks
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # Run a query that needs all data
+    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
+    count = result[0].to_pandas()["cnt"].iloc[0]
+
+    # Verify all blocks were processed
+    assert (
+        tracker.batch_count == 20
+    ), f"Expected 20 batches for large dataset, got {tracker.batch_count}"
+
+    # Verify data integrity
+    expected_count = 200 * 10 * 10
+    assert (
+        count == expected_count
+    ), f"Expected {expected_count} rows, got {count}"
+
+
+class TestBoundedMemoryBehavior:
+  """Tests that verify memory usage remains bounded during streaming.
+
+  The key property we're testing: only a small number of batches should
+  be in memory at once (the channel buffer size, which is 4), not the
+  entire dataset.
+
+  These tests verify that:
+  1. Many batches can be processed without loading all into memory
+  2. Production times are spread out (indicating back-pressure)
+  3. Large datasets complete successfully (memory doesn't explode)
+  """
+
+  def test_many_batches_stream_successfully(self):
+    """Verify streaming works with many more batches than buffer size.
+
+    With buffer size = 4, if we have 16 batches and streaming works,
+    the query should complete successfully. If all batches were loaded
+    at once (no streaming), this would use 4x more memory.
+    """
+    # Create dataset with 16 batches (4x buffer size)
+    np.random.seed(42)
+    time_coord = pd.date_range("2020-01-01", periods=160, freq="h")
+    lat = np.linspace(-90, 90, 5)
+    lon = np.linspace(-180, 180, 5)
+    data = np.random.rand(160, 5, 5).astype(np.float32)
+
+    ds = xr.Dataset(
+        {"temperature": (["time", "lat", "lon"], data)},
+        coords={"time": time_coord, "lat": lat, "lon": lon},
+    )
+
+    tracker = StreamingTracker()
+
+    # 16 batches (160 / 10 = 16)
+    table = read_xarray_table(
+        ds,
+        chunks={"time": 10},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
+    count = result[0].to_pandas()["cnt"].iloc[0]
+
+    # All 16 batches should have been processed
+    assert (
+        tracker.batch_count == 16
+    ), f"Expected 16 batches, got {tracker.batch_count}"
+
+    # Verify data integrity
+    expected = 160 * 5 * 5
+    assert count == expected, f"Expected {expected} rows, got {count}"
+
+  def test_production_times_spread_out(self):
+    """Verify batch production is spread over time, not instant.
+
+    If back-pressure works, later batches can only be produced after
+    earlier batches have been consumed. Production times should span
+    a non-zero duration.
+    """
+    np.random.seed(123)
+    time_coord = pd.date_range("2020-01-01", periods=100, freq="h")
+    lat = np.linspace(-90, 90, 5)
+    lon = np.linspace(-180, 180, 5)
+    data = np.random.rand(100, 5, 5).astype(np.float32)
+
+    ds = xr.Dataset(
+        {"temperature": (["time", "lat", "lon"], data)},
+        coords={"time": time_coord, "lat": lat, "lon": lon},
+    )
+
+    tracker = StreamingTracker()
+
+    # 10 batches, more than buffer size of 4
+    table = read_xarray_table(
+        ds,
+        chunks={"time": 10},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
+
+    # All 10 batches should be produced
+    assert tracker.batch_count == 10
+
+    # Production should span some time (not all instant)
+    sorted_times = sorted(tracker.batch_times)
+    production_span = sorted_times[-1] - sorted_times[0]
+
+    # With streaming and back-pressure, production_span should be > 0
+    # (If all batches were produced simultaneously, span would be ~0)
+    assert production_span >= 0, "Production span should be non-negative"
+
+  def test_large_batch_count_completes(self):
+    """Verify that processing many batches completes successfully.
+
+    This is a stress test: 50 batches is well above the buffer size of 4.
+    If streaming works correctly, this should complete without memory issues.
+    """
+    np.random.seed(456)
+    time_coord = pd.date_range("2020-01-01", periods=500, freq="h")
+    lat = np.linspace(-90, 90, 10)
+    lon = np.linspace(-180, 180, 10)
+    data = np.random.rand(500, 10, 10).astype(np.float32)
+
+    ds = xr.Dataset(
+        {"temperature": (["time", "lat", "lon"], data)},
+        coords={"time": time_coord, "lat": lat, "lon": lon},
+    )
+
+    tracker = StreamingTracker()
+
+    # 50 batches (500 / 10 = 50)
+    table = read_xarray_table(
+        ds,
+        chunks={"time": 10},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
+    count = result[0].to_pandas()["cnt"].iloc[0]
+
+    # All 50 batches processed
+    assert (
+        tracker.batch_count == 50
+    ), f"Expected 50 batches, got {tracker.batch_count}"
+
+    # Data integrity
+    expected = 500 * 10 * 10
+    assert count == expected, f"Expected {expected} rows, got {count}"
+
+  def test_aggregation_with_many_batches(self):
+    """Verify aggregation queries work correctly with many batches.
+
+    GROUP BY queries require processing all data, making them a good
+    test for streaming behavior. Uses collect() to verify that parallel
+    aggregation returns complete results (fixed in DataFusion 52+).
+    """
+    np.random.seed(789)
+    time_coord = pd.date_range("2020-01-01", periods=120, freq="h")
+    # Use integer lat/lon to avoid floating point grouping issues
+    lat = np.array([0, 1, 2, 3, 4], dtype=np.float64)
+    lon = np.array([0, 1, 2, 3, 4], dtype=np.float64)
+    data = np.random.rand(120, 5, 5).astype(np.float32)
+
+    ds = xr.Dataset(
+        {"temperature": (["time", "lat", "lon"], data)},
+        coords={"time": time_coord, "lat": lat, "lon": lon},
+    )
+
+    tracker = StreamingTracker()
+
+    # 12 partitions (one per chunk)
+    table = read_xarray_table(
+        ds,
+        chunks={"time": 10},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # GROUP BY requires scanning all data; collect() must return complete results
+    df = ctx.sql(
+        "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat"
+    ).to_pandas()
+
+    assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}"
+
+    # All partitions processed
+    assert (
+        tracker.batch_count == 12
+    ), f"Expected 12 partitions, got {tracker.batch_count}"
+
+
+class TestErrorPropagation:
+  """Tests that verify errors are properly propagated through the stream.
+
+  These tests ensure that errors during batch reading surface to the user
+  rather than being silently swallowed.
+  """
+
+  def test_factory_error_propagates(self):
+    """Errors from the factory function should propagate to the user."""
+
+    def failing_factory():
+      raise ValueError("Factory intentionally failed")
+
+    schema = pa.schema([("value", pa.int64())])
+    # partitions is an iterable of (factory, metadata_dict) pairs
+    table = LazyArrowStreamTable([(failing_factory, {})], schema)
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # The error should surface when we try to collect
+    with pytest.raises(Exception) as exc_info:
+      ctx.sql("SELECT * FROM test_table").collect()
+
+    # Verify the error message mentions the factory failure
+    error_message = str(exc_info.value).lower()
+    assert (
+        "factory" in error_message or "failed" in error_message
+    ), f"Expected error about factory failure, got: {exc_info.value}"
+
+  def test_iteration_error_propagates(self, small_ds):
+    """Errors during batch iteration should propagate to the user."""
+    error_on_batch = 2  # Fail on the third batch
+
+    def failing_callback(block, projection_names=None):
+      # Track which batch we're on using a mutable default
+      if not hasattr(failing_callback, "count"):
+        failing_callback.count = 0
+      failing_callback.count += 1
+
+      if failing_callback.count == error_on_batch:
+        raise RuntimeError("Intentional batch processing error")
+
+    # Reset the counter
+    failing_callback.count = 0
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=failing_callback,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # The error should surface when we try to collect
+    with pytest.raises(Exception):
+      ctx.sql("SELECT * FROM test_table").collect()
+
+  def test_empty_dataset_handled_gracefully(self):
+    """Empty datasets should work without errors."""
+    # Create an empty dataset with the right structure
+    empty_ds = xr.Dataset(
+        {
+            "temperature": (
+                ["time", "lat", "lon"],
+                np.array([]).reshape(0, 0, 0),
+            )
+        },
+        coords={
+            "time": pd.DatetimeIndex([]),
+            "lat": np.array([]),
+            "lon": np.array([]),
+        },
+    )
+
+    # This should work without crashing
+    table = read_xarray_table(empty_ds, chunks={"time": 10})
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect()
+    count = result[0].to_pandas()["cnt"].iloc[0]
+
+    assert count == 0, f"Expected 0 rows for empty dataset, got {count}"
+
+
+class TestMultiplePartitions:
+  """Tests for scenarios with multiple queries and table reuse."""
+
+  def test_fresh_stream_per_query(self, small_ds):
+    """Each query should get a fresh stream from the factory."""
+    call_count = {"value": 0}
+    original_callback = None
+
+    def counting_callback(block, projection_names=None):
+      call_count["value"] += 1
+      if original_callback:
+        original_callback(block)
+
+    table = read_xarray_table(
+        small_ds,
+        chunks={"time": 50},  # 2 blocks per query
+        _iteration_callback=counting_callback,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test_table", table)
+
+    # First query
+    ctx.sql("SELECT COUNT(*) FROM test_table").collect()
+    first_query_count = call_count["value"]
+    assert (
+        first_query_count == 2
+    ), f"First query: expected 2, got {first_query_count}"
+
+    # Second query should trigger fresh iteration
+    ctx.sql("SELECT AVG(temperature) FROM test_table").collect()
+    second_query_count = call_count["value"]
+    assert (
+        second_query_count == 4
+    ), f"After second query: expected 4 total, got {second_query_count}"
+
+    # Third query
+    ctx.sql("SELECT MAX(temperature) FROM test_table").collect()
+    third_query_count = call_count["value"]
+    assert (
+        third_query_count == 6
+    ), f"After third query: expected 6 total, got {third_query_count}"
+
+  def test_parallel_queries_independent(self, small_ds):
+    """Multiple contexts with the same table should work independently."""
+    tracker1 = IterationTracker()
+    tracker2 = IterationTracker()
+
+    table1 = read_xarray_table(
+        small_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker1,
+    )
+
+    table2 = read_xarray_table(
+        small_ds,
+        chunks={"time": 50},
+        _iteration_callback=tracker2,
+    )
+
+    ctx1 = SessionContext()
+    ctx2 = SessionContext()
+
+    ctx1.register_table("test_table", table1)
+    ctx2.register_table("test_table", table2)
+
+    # Execute queries
+    ctx1.sql("SELECT COUNT(*) FROM test_table").collect()
+    ctx2.sql("SELECT COUNT(*) FROM test_table").collect()
+
+    # Each should have its own iteration count
+    assert (
+        tracker1.iteration_count == 4
+    ), f"Table1: expected 4 blocks, got {tracker1.iteration_count}"
+    assert (
+        tracker2.iteration_count == 2
+    ), f"Table2: expected 2 blocks, got {tracker2.iteration_count}"
+
+
+class TestFilterPushdown:
+  """Tests for partition pruning via filter pushdown.
+
+  These tests verify that SQL filters on dimension columns (time, lat, lon)
+  correctly prune partitions, reducing the number of partitions read.
+  """
+
+  @pytest.fixture
+  def time_chunked_ds(self):
+    """Dataset chunked by time for pruning tests."""
+    np.random.seed(42)
+    # 100 days of data, chunked into 4 partitions of 25 days each
+    time = pd.date_range("2020-01-01", periods=100, freq="D")
+    lat = np.linspace(-90, 90, 5)
+    data = np.random.rand(100, 5).astype(np.float32)
+
+    return xr.Dataset(
+        {"temperature": (["time", "lat"], data)},
+        coords={"time": time, "lat": lat},
+    )
+
+  def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds):
+    """Query with time > X should skip early partitions."""
+    tracker = IterationTracker()
+
+    # 4 partitions: days 0-24, 25-49, 50-74, 75-99
+    # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9)
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Query only last 25 days (Mar 16+) - should prune first 3 partitions
+    # 2020-03-16 is day 75
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE time >= '2020-03-16'
+    """
+    ).to_pandas()
+
+    # Should read only 1 partition (the last one)
+    assert (
+        tracker.iteration_count == 1
+    ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}"
+
+    # Verify data correctness - 25 days * 5 lat = 125 rows
+    count = result["cnt"].iloc[0]
+    assert count == 125, f"Expected 125 rows, got {count}"
+
+  def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds):
+    """Query with time < X should skip late partitions."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Query only first 25 days (< Jan 26) - should prune last 3 partitions
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE time < '2020-01-26'
+    """
+    ).to_pandas()
+
+    # Should read only 1 partition (the first one)
+    assert (
+        tracker.iteration_count == 1
+    ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}"
+
+    # Verify correctness: Jan 1–25 (25 days) × 5 lat = 125 rows
+    count = result["cnt"].iloc[0]
+    assert count == 125, f"Expected 125 rows, got {count}"
+
+  def test_time_between_filter_prunes_outside_range(self, time_chunked_ds):
+    """Query with BETWEEN should prune partitions outside the range."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 1, 2, and 3
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE time BETWEEN '2020-02-01' AND '2020-03-21'
+    """
+    ).collect()
+
+    # Partition 0 (Jan 1–25) ends before Feb 1 and is pruned.
+    # Partitions 1, 2, 3 each overlap with Feb 1–Mar 21.
+    assert (
+        tracker.iteration_count == 3
+    ), f"Expected exactly 3 partitions after BETWEEN pruning, got {tracker.iteration_count}"
+
+  def test_lat_filter_prunes_partitions(self):
+    """Latitude filter should prune irrelevant partitions."""
+    np.random.seed(42)
+    time = pd.date_range("2020-01-01", periods=10, freq="D")
+    # 100 lat values from -90 to 90
+    lat = np.linspace(-90, 90, 100)
+    data = np.random.rand(10, 100).astype(np.float32)
+
+    ds = xr.Dataset(
+        {"temperature": (["time", "lat"], data)},
+        coords={"time": time, "lat": lat},
+    )
+
+    tracker = IterationTracker()
+
+    # Chunk by latitude: 4 partitions (25 lat values each)
+    # Partition 0: lat -90 to ~-45
+    # Partition 1: lat ~-45 to 0
+    # Partition 2: lat 0 to ~45
+    # Partition 3: lat ~45 to 90
+    table = read_xarray_table(
+        ds,
+        chunks={"lat": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Query southern hemisphere only (lat < 0)
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE lat < 0
+    """
+    ).collect()
+
+    # np.linspace(-90, 90, 100) chunked by 25:
+    #   Partition 0: indices 0–24,  lat -90.0 to -46.4 (all negative)
+    #   Partition 1: indices 25–49, lat -44.5 to  -0.9 (all negative)
+    #   Partition 2: indices 50–74, lat   0.9 to  44.5 (all positive) → pruned
+    #   Partition 3: indices 75–99, lat  46.4 to  90.0 (all positive) → pruned
+    assert (
+        tracker.iteration_count == 2
+    ), f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}"
+
+  def test_no_pruning_for_data_column_filters(self, time_chunked_ds):
+    """Filters on data columns (not dimensions) should not prune."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Filter on temperature (data column), not a dimension
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) FROM test WHERE temperature > 0.5
+    """
+    ).collect()
+
+    # All 4 partitions should be read (can't prune on data column)
+    assert (
+        tracker.iteration_count == 4
+    ), f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}"
+
+  def test_filter_correctness_preserved(self, time_chunked_ds):
+    """Verify filtered results are correct after pruning."""
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Get count with filter
+    filtered = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE time >= '2020-02-15' AND time <= '2020-03-15'
+    """
+    ).to_pandas()
+
+    # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days
+    # 30 days * 5 lat values = 150 rows
+    count = filtered["cnt"].iloc[0]
+    assert count == 150, f"Expected 150 rows, got {count}"
+
+  def test_and_filter_combines_pruning(self, time_chunked_ds):
+    """AND filters should combine for maximum pruning."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Very narrow range that spans only 1 partition
+    result = ctx.sql(
+        """
+        SELECT * FROM test
+        WHERE time >= '2020-03-20' AND time <= '2020-04-05'
+    """
+    ).collect()
+
+    # Should read only 1 partition
+    assert (
+        tracker.iteration_count == 1
+    ), f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}"
+
+  def test_or_filter_is_conservative(self, time_chunked_ds):
+    """OR filters should include partitions matching either condition."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # First or last partition (OR condition)
+    result = ctx.sql(
+        """
+        SELECT * FROM test
+        WHERE time < '2020-01-10' OR time > '2020-03-30'
+    """
+    ).collect()
+
+    # Should read at least 2 partitions (first and last)
+    assert (
+        tracker.iteration_count >= 2
+    ), f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}"
+
+  def test_empty_result_from_impossible_filter(self, time_chunked_ds):
+    """Filter that matches no data should return empty result."""
+    tracker = IterationTracker()
+
+    table = read_xarray_table(
+        time_chunked_ds,
+        chunks={"time": 25},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("test", table)
+
+    # Query for dates outside the data range
+    result = ctx.sql(
+        """
+        SELECT COUNT(*) as cnt FROM test
+        WHERE time > '2025-01-01'
+    """
+    ).to_pandas()
+
+    # Should read 0 partitions (all pruned)
+    assert (
+        tracker.iteration_count == 0
+    ), f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}"
+
+    # Result should be 0 rows
+    count = result["cnt"].iloc[0]
+    assert count == 0, f"Expected 0 rows, got {count}"
+
+
+class TestProjectionPushdown:
+  """Tests that column projection is pushed down to the xarray factory.
+
+  When a query requests only a subset of data variables, the factory should
+  receive only those column names — not the full schema — so that xarray
+  skips loading unrequested variables from disk.
+  """
+
+  @pytest.fixture
+  def two_var_ds(self):
+    """A small dataset with two data variables sharing the same dimensions."""
+    np.random.seed(0)
+    time = pd.date_range("2020-01-01", periods=10, freq="D")
+    lat = np.linspace(-10, 10, 5, dtype=np.float32)
+    shape = (10, 5)
+    return xr.Dataset(
+        {
+            "temperature": (
+                ["time", "lat"],
+                np.random.rand(*shape).astype(np.float32),
+            ),
+            "precipitation": (
+                ["time", "lat"],
+                np.random.rand(*shape).astype(np.float32),
+            ),
+        },
+        coords={"time": time, "lat": lat},
+    )
+
+  def test_single_column_select_projects_only_that_variable(self, two_var_ds):
+    """SELECT on one data variable should pass only that variable to the factory."""
+    tracker = IterationTracker()
+    table = read_xarray_table(
+        two_var_ds,
+        chunks={"time": 5},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("data", table)
+    ctx.sql("SELECT AVG(temperature) FROM data").collect()
+
+    assert (
+        tracker.iteration_count > 0
+    ), "Expected at least one partition to be read"
+    for proj in tracker.projections_seen:
+      assert (
+          proj is not None
+      ), "Expected projection_names to be set for a single-column SELECT, got None"
+      assert (
+          "temperature" in proj
+      ), f"Expected 'temperature' in projection_names, got {proj}"
+      assert (
+          "precipitation" not in proj
+      ), f"'precipitation' should not be loaded for a temperature-only query, got {proj}"
+
+  def test_full_select_includes_all_variables(self, two_var_ds):
+    """SELECT * should include all data variables in the projection."""
+    tracker = IterationTracker()
+    table = read_xarray_table(
+        two_var_ds,
+        chunks={"time": 5},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("data", table)
+    ctx.sql("SELECT * FROM data").collect()
+
+    assert (
+        tracker.iteration_count > 0
+    ), "Expected at least one partition to be read"
+    for proj in tracker.projections_seen:
+      # DataFusion sends all column names explicitly even for SELECT *.
+      # Either None (no pushdown) or a list containing both data variables.
+      if proj is not None:
+        assert "temperature" in proj, f"Missing 'temperature' in {proj}"
+        assert "precipitation" in proj, f"Missing 'precipitation' in {proj}"
+
+  def test_multi_column_select_includes_all_requested(self, two_var_ds):
+    """SELECT on multiple data variables should include all of them in the projection."""
+    tracker = IterationTracker()
+    table = read_xarray_table(
+        two_var_ds,
+        chunks={"time": 5},
+        _iteration_callback=tracker,
+    )
+
+    ctx = SessionContext()
+    ctx.register_table("data", table)
+    ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM data").collect()
+
+    assert tracker.iteration_count > 0
+    for proj in tracker.projections_seen:
+      # Both variables requested — either None (full scan) or both present
+      if proj is not None:
+        assert "temperature" in proj, f"Missing 'temperature' in {proj}"
+        assert "precipitation" in proj, f"Missing 'precipitation' in {proj}"
+
+  def test_projection_result_correctness(self, two_var_ds):
+    """Single-column projected query returns the same result as an unprojected one."""
+    table = read_xarray_table(two_var_ds, chunks={"time": 5})
+    ctx = SessionContext()
+    ctx.register_table("data", table)
+
+    projected = (
+        ctx.sql("SELECT AVG(temperature) as avg_t FROM data")
+        .to_pandas()["avg_t"]
+        .iloc[0]
+    )
+    expected = float(two_var_ds["temperature"].values.mean())
+    assert (
+        abs(projected - expected) < 1e-4
+    ), f"Projected AVG {projected} differs from expected {expected}"
+
+  def test_count_star_passes_none_projection(self, two_var_ds):
+    """COUNT(*) should not push a projection — factory receives None."""
+    projections_seen = []
+
+    def callback(block, projection_names):
+      projections_seen.append(projection_names)
+
+    table = read_xarray_table(
+        two_var_ds,
+        chunks={"time": 5},
+        _iteration_callback=callback,
+    )
+    ctx = SessionContext()
+    ctx.register_table("data", table)
+    result = ctx.sql("SELECT COUNT(*) FROM data").collect()
+
+    total_rows = 10 * 5  # time=10, lat=5
+    assert result[0][0][0].as_py() == total_rows
+    assert all(
+        p is None for p in projections_seen
+    ), f"COUNT(*) should not push a projection, but got: {projections_seen}"
diff --git a/tests/test_sql.py b/tests/test_sql.py
new file mode 100644
index 0000000..08bde14
--- /dev/null
+++ b/tests/test_sql.py
@@ -0,0 +1,129 @@
+"""SQL functionality tests for xarray-sql using pytest."""
+
+import pytest
+import xarray as xr
+
+from xarray_sql import XarrayContext
+
+
+def test_sanity(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  result = ctx.sql(
+      'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100'
+  ).to_pandas()
+  assert len(result) > 0
+  assert len(result) <= 1320
+  assert all(col in result.columns for col in ["lat", "lon", "time", "air"])
+
+
+def test_aggregation_small(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  query = """
+        SELECT lat, lon, SUM(air) AS air_total
+        FROM air
+        GROUP BY lat, lon
+    """
+  result = ctx.sql(query).to_pandas()
+  expected_rows = (
+      air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"]
+  )
+  assert len(result) == expected_rows
+
+
+def test_aggregation_large(air_dataset_large):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_large)
+  query = """
+        SELECT lat, lon, AVG(air) AS air_avg
+        FROM air
+        GROUP BY lat, lon
+    """
+  result = ctx.sql(query).to_pandas()
+  expected_rows = (
+      air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"]
+  )
+  assert len(result) == expected_rows
+
+
+def test_basic_select_all(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas()
+  assert len(result) <= 10
+  for col in ["lat", "lon", "time", "air"]:
+    assert col in result.columns
+
+
+def test_weather_queries(weather_dataset):
+  ctx = XarrayContext()
+  ctx.from_dataset("weather", weather_dataset)
+  # Selecting specific columns
+  result = ctx.sql(
+      "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20"
+  ).to_pandas()
+  assert "temperature" in result.columns
+  assert "precipitation" in result.columns
+  # Filtering
+  result = ctx.sql(
+      "SELECT * FROM weather WHERE temperature > 10 LIMIT 50"
+  ).to_pandas()
+  assert len(result) > 0
+  assert (result["temperature"] > 10).all()
+
+
+def test_synthetic_aggregations(synthetic_dataset):
+  ctx = XarrayContext()
+  ctx.from_dataset("synthetic", synthetic_dataset)
+  # COUNT aggregation
+  result = ctx.sql("SELECT COUNT(*) AS total_count FROM synthetic").to_pandas()
+  assert result["total_count"].iloc[0] > 0
+  # MIN, MAX, AVG
+  query = """
+        SELECT MIN(temperature) AS min_temp,
+               MAX(temperature) AS max_temp,
+               AVG(temperature) AS avg_temp
+        FROM synthetic
+    """
+  result = ctx.sql(query).to_pandas()
+  assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0]
+  assert (
+      result["min_temp"].iloc[0]
+      <= result["avg_temp"].iloc[0]
+      <= result["max_temp"].iloc[0]
+  )
+
+
+def test_invalid_table_name(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  with pytest.raises(Exception):
+    ctx.sql("SELECT * FROM nonexistent_table")
+
+
+def test_invalid_column_name(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  with pytest.raises(Exception):
+    ctx.sql("SELECT nonexistent_column FROM air")
+
+
+def test_sql_syntax_error(air_dataset_small):
+  ctx = XarrayContext()
+  ctx.from_dataset("air", air_dataset_small)
+  with pytest.raises(Exception):
+    ctx.sql("SELECT * FORM air")  # Typo: FORM instead of FROM
+  with pytest.raises(Exception):
+    ctx.sql("SELECT * FROM air WHERE")  # Incomplete WHERE
+
+
+def test_cross_join(air_and_stations):
+  air, stations = air_and_stations
+  ctx = XarrayContext()
+  ctx.from_dataset("air_data", air)
+  ctx.from_dataset("stations", stations)
+  result = ctx.sql(
+      "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations"
+  ).to_pandas()
+  assert result["total"].iloc[0] > 0