diff --git a/pyproject.toml b/pyproject.toml index 6bfb1d0f..78e289d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,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 @@ -93,3 +93,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/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..b173f2e6 --- /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/xarray_sql/df_test.py b/tests/test_df.py similarity index 87% rename from xarray_sql/df_test.py rename to tests/test_df.py index 8eeba0d1..37380aee 100644 --- a/xarray_sql/df_test.py +++ b/tests/test_df.py @@ -6,7 +6,7 @@ import pytest import xarray as xr -from .df import ( +from xarray_sql.df import ( DEFAULT_BATCH_SIZE, _parse_schema, block_slices, @@ -17,82 +17,7 @@ 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}) +from xarray_sql.reader import read_xarray, read_xarray_table def test_explode_cardinality(air): @@ -295,6 +220,12 @@ def test_from_map_batched_basic_functionality(air_small): 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] diff --git a/xarray_sql/reader_test.py b/tests/test_reader.py similarity index 99% rename from xarray_sql/reader_test.py rename to tests/test_reader.py index 6ab64068..d970b4c4 100644 --- a/xarray_sql/reader_test.py +++ b/tests/test_reader.py @@ -23,9 +23,9 @@ import xarray as xr from datafusion import SessionContext -from ._native import LazyArrowStreamTable -from .reader import XarrayRecordBatchReader, read_xarray_table -from .df import _parse_schema +from xarray_sql._native import LazyArrowStreamTable +from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table +from xarray_sql.df import _parse_schema @pytest.fixture diff --git a/xarray_sql/sql_test.py b/tests/test_sql.py similarity index 68% rename from xarray_sql/sql_test.py rename to tests/test_sql.py index 8feedeec..08bde146 100644 --- a/xarray_sql/sql_test.py +++ b/tests/test_sql.py @@ -3,72 +3,7 @@ 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 +from xarray_sql import XarrayContext def test_sanity(air_dataset_small):