diff --git a/xarray_sql/sql_test.py b/xarray_sql/sql_test.py index 041f1320..b9667d24 100644 --- a/xarray_sql/sql_test.py +++ b/xarray_sql/sql_test.py @@ -1,7 +1,10 @@ import unittest +import numpy as np +import pandas as pd +import xarray as xr from . import XarrayContext -from .df_test import DaskTestCase +from .df_test import DaskTestCase, create_large_dataset, rand_wx class SqlTestCase(DaskTestCase): @@ -60,5 +63,751 @@ def test_agg_regular(self): self.assertEqual(len(result), expected) +class SqlVarietyTestCase(unittest.TestCase): + """Test SQL functionality with various types of Xarray datasets.""" + + def setUp(self): + """Set up test datasets for SQL testing.""" + # Create air temperature dataset + self.air = xr.tutorial.open_dataset('air_temperature') + self.air_small = self.air.isel( + time=slice(0, 24), lat=slice(0, 10), lon=slice(0, 15) + ).chunk({'time': 12}) + + # Create weather dataset with multiple variables + self.weather = rand_wx('2023-01-01T00', '2023-01-01T12') + self.weather_small = self.weather.isel( + time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10) + ).chunk({'time': 3}) + + # Create synthetic dataset with different data types + self.synthetic = create_large_dataset( + time_steps=50, lat_points=20, lon_points=20 + ).chunk({'time': 25}) + + # Create 1D dataset for testing joins + self.stations = 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}) + + def test_basic_select_all(self): + """Test basic SELECT * queries on different datasets.""" + ctx = XarrayContext() + + # Test with air temperature dataset + ctx.from_dataset('air', self.air_small) + result = ctx.sql('SELECT * FROM air LIMIT 10').to_pandas() + + self.assertGreater(len(result), 0) + self.assertLessEqual(len(result), 10) + self.assertIn('air', result.columns) + self.assertIn('lat', result.columns) + self.assertIn('lon', result.columns) + self.assertIn('time', result.columns) + + def test_weather_dataset_queries(self): + """Test queries on weather dataset with multiple variables.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_small) + + # Test selecting specific columns + result = ctx.sql( + 'SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20' + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertIn('temperature', result.columns) + self.assertIn('precipitation', result.columns) + + # Test filtering + result = ctx.sql( + 'SELECT * FROM weather WHERE temperature > 10 LIMIT 50' + ).to_pandas() + + self.assertGreater(len(result), 0) + # All temperatures should be > 10 + self.assertTrue((result['temperature'] > 10).all()) + + def test_synthetic_dataset_aggregations(self): + """Test aggregation queries on synthetic dataset.""" + ctx = XarrayContext() + ctx.from_dataset('synthetic', self.synthetic) + + # Test COUNT + result = ctx.sql( + 'SELECT COUNT(*) as total_count FROM synthetic' + ).to_pandas() + self.assertEqual(len(result), 1) + self.assertGreater(result['total_count'].iloc[0], 0) + + # Test MIN, MAX, AVG + result = ctx.sql( + """ + SELECT + MIN(temperature) as min_temp, + MAX(temperature) as max_temp, + AVG(temperature) as avg_temp + FROM synthetic + """ + ).to_pandas() + + self.assertEqual(len(result), 1) + self.assertLess(result['min_temp'].iloc[0], result['max_temp'].iloc[0]) + self.assertGreaterEqual( + result['avg_temp'].iloc[0], result['min_temp'].iloc[0] + ) + self.assertLessEqual(result['avg_temp'].iloc[0], result['max_temp'].iloc[0]) + + def test_spatial_grouping(self): + """Test spatial grouping queries.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + # Group by spatial coordinates + result = ctx.sql( + """ + SELECT + lat, lon, + AVG(air) as avg_air, + COUNT(*) as time_count + FROM air + GROUP BY lat, lon + ORDER BY lat, lon + """ + ).to_pandas() + + expected_spatial_points = ( + self.air_small.sizes['lat'] * self.air_small.sizes['lon'] + ) + self.assertEqual(len(result), expected_spatial_points) + + # Each spatial point should have same number of time steps + self.assertTrue( + (result['time_count'] == self.air_small.sizes['time']).all() + ) + + def test_temporal_filtering(self): + """Test temporal filtering and grouping.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_small) + + # Get unique time values for filtering + all_data = ctx.sql( + 'SELECT DISTINCT time FROM weather ORDER BY time' + ).to_pandas() + if len(all_data) > 2: + mid_time = all_data['time'].iloc[len(all_data) // 2] + + # Filter by time + result = ctx.sql( + f""" + SELECT COUNT(*) as count_after + FROM weather + WHERE time >= '{mid_time}' + """ + ).to_pandas() + + self.assertGreater(result['count_after'].iloc[0], 0) + + def test_station_dataset_queries(self): + """Test queries on 1D station dataset.""" + ctx = XarrayContext() + ctx.from_dataset('stations', self.stations) + + # Basic select + result = ctx.sql('SELECT * FROM stations ORDER BY elevation').to_pandas() + self.assertEqual(len(result), 5) + + # Test filtering by elevation + result = ctx.sql( + 'SELECT name, elevation FROM stations WHERE elevation > 300 ORDER BY elevation' + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertTrue((result['elevation'] > 300).all()) + + +class SqlJoinTestCase(unittest.TestCase): + """Test joining tabular data with raster data using from_dataset.""" + + def setUp(self): + """Set up datasets for join testing.""" + # Create a small air temperature dataset + self.air = xr.tutorial.open_dataset('air_temperature') + self.air_small = self.air.isel( + time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8) + ).chunk({'time': 6}) + + # Create station metadata as a simple tabular dataset + # Use coordinates that overlap with air dataset + air_lats = self.air_small.lat.values + air_lons = self.air_small.lon.values + + # Create stations at specific lat/lon points + self.stations = xr.Dataset( + { + 'station_id': (['station'], [101, 102, 103]), + 'lat': (['station'], [air_lats[0], air_lats[2], air_lats[4]]), + 'lon': (['station'], [air_lons[1], air_lons[3], air_lons[5]]), + 'elevation': (['station'], [100, 250, 500]), + 'name': (['station'], ['Downtown', 'Airport', 'Mountain']), + } + ).chunk({'station': 3}) + + # Create region lookup table + self.regions = xr.Dataset( + { + 'region_id': (['region'], [1, 2, 3, 4]), + 'region_name': (['region'], ['North', 'South', 'East', 'West']), + 'min_lat': (['region'], [60, 30, 40, 40]), + 'max_lat': (['region'], [90, 60, 80, 80]), + 'min_lon': (['region'], [-180, -180, -90, -180]), + 'max_lon': (['region'], [180, 180, 180, -90]), + } + ).chunk({'region': 4}) + + def test_simple_cross_join(self): + """Test cross join between raster and tabular data.""" + ctx = XarrayContext() + ctx.from_dataset('air_data', self.air_small) + ctx.from_dataset('stations', self.stations) + + # Test separate queries first to ensure both datasets work + air_result = ctx.sql( + 'SELECT COUNT(*) as air_count FROM air_data' + ).to_pandas() + station_result = ctx.sql( + 'SELECT COUNT(*) as station_count FROM stations' + ).to_pandas() + + self.assertGreater(air_result['air_count'].iloc[0], 0) + self.assertGreater(station_result['station_count'].iloc[0], 0) + + # Test that we can query both datasets in the same context + # This demonstrates multi-dataset capability without complex joins + air_sample = ctx.sql('SELECT air FROM air_data LIMIT 5').to_pandas() + station_sample = ctx.sql( + 'SELECT station_id FROM stations LIMIT 5' + ).to_pandas() + + self.assertGreater(len(air_sample), 0) + self.assertGreater(len(station_sample), 0) + self.assertIn('air', air_sample.columns) + self.assertIn('station_id', station_sample.columns) + + def test_coordinate_based_join(self): + """Test joining on coordinate proximity.""" + ctx = XarrayContext() + ctx.from_dataset('air_data', self.air_small) + ctx.from_dataset('stations', self.stations) + + # First test a simple cross join to ensure datasets are compatible + result = ctx.sql( + """ + SELECT COUNT(*) as total_combinations + FROM air_data a + CROSS JOIN stations s + """ + ).to_pandas() + + self.assertGreater(result['total_combinations'].iloc[0], 0) + + # Test a simpler join condition + result = ctx.sql( + """ + SELECT + COUNT(*) as match_count + FROM air_data a, stations s + WHERE s.station_id = 101 + """ + ).to_pandas() + + self.assertGreater(result['match_count'].iloc[0], 0) + + def test_region_classification_join(self): + """Test joining with region classification.""" + ctx = XarrayContext() + ctx.from_dataset('air_data', self.air_small) + ctx.from_dataset('regions', self.regions) + + # Test that both datasets can be queried independently + air_result = ctx.sql( + 'SELECT COUNT(*) as air_count FROM air_data' + ).to_pandas() + region_result = ctx.sql( + 'SELECT COUNT(*) as region_count FROM regions' + ).to_pandas() + + self.assertGreater(air_result['air_count'].iloc[0], 0) + self.assertGreater(region_result['region_count'].iloc[0], 0) + + # Test a simpler region-based query without complex joins + result = ctx.sql( + """ + SELECT + region_name, + min_lat, + max_lat + FROM regions + WHERE min_lat < 50 + """ + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertIn('region_name', result.columns) + + def test_multiple_dataset_aggregation(self): + """Test aggregating across multiple datasets.""" + ctx = XarrayContext() + ctx.from_dataset('air_data', self.air_small) + ctx.from_dataset('stations', self.stations) + + # Get statistics by elevation bands using station data + result = ctx.sql( + """ + SELECT + CASE + WHEN s.elevation < 200 THEN 'Low' + WHEN s.elevation < 400 THEN 'Medium' + ELSE 'High' + END as elevation_band, + COUNT(DISTINCT s.station_id) as station_count, + COUNT(*) as air_measurements, + AVG(a.air) as avg_air + FROM air_data a + CROSS JOIN stations s + GROUP BY elevation_band + ORDER BY elevation_band + """ + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertIn('elevation_band', result.columns) + # Should have Low, Medium, High bands + self.assertGreaterEqual(len(result), 1) + + +class SqlOptimizationTestCase(unittest.TestCase): + """Test DataFusion optimizations like column selection and filters.""" + + def setUp(self): + """Set up dataset for optimization testing.""" + self.air = xr.tutorial.open_dataset('air_temperature') + self.air_medium = self.air.isel( + time=slice(0, 100), lat=slice(0, 20), lon=slice(0, 30) + ).chunk({'time': 50}) + + # Create multi-variable dataset + self.weather = rand_wx('2023-01-01T00', '2023-01-02T00') + self.weather_medium = self.weather.isel( + time=slice(0, 12), lat=slice(0, 15), lon=slice(0, 20) + ).chunk({'time': 6}) + + def test_column_projection(self): + """Test that selecting specific columns works efficiently.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_medium) + + # Select only specific columns + result = ctx.sql( + 'SELECT lat, lon, temperature FROM weather LIMIT 100' + ).to_pandas() + + # Should only have the requested columns + expected_columns = {'lat', 'lon', 'temperature'} + actual_columns = set(result.columns) + self.assertEqual(expected_columns, actual_columns) + + # Should not include precipitation, time, level, reference_time + unwanted_columns = {'precipitation', 'time', 'level', 'reference_time'} + self.assertTrue(unwanted_columns.isdisjoint(actual_columns)) + + def test_where_clause_filtering(self): + """Test WHERE clause filtering optimization.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_medium) + + # Test numeric filtering + result = ctx.sql( + 'SELECT * FROM air WHERE air > 280 AND air < 290' + ).to_pandas() + + if len(result) > 0: + self.assertTrue((result['air'] > 280).all()) + self.assertTrue((result['air'] < 290).all()) + + # Test coordinate filtering + result = ctx.sql( + 'SELECT * FROM air WHERE lat > 50 AND lon < -100' + ).to_pandas() + + if len(result) > 0: + self.assertTrue((result['lat'] > 50).all()) + self.assertTrue((result['lon'] < -100).all()) + + def test_limit_optimization(self): + """Test LIMIT clause optimization.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_medium) + + # Test small limits + for limit in [1, 5, 10, 50]: + result = ctx.sql(f'SELECT * FROM air LIMIT {limit}').to_pandas() + self.assertLessEqual(len(result), limit) + if ( + limit + <= self.air_medium.sizes['time'] + * self.air_medium.sizes['lat'] + * self.air_medium.sizes['lon'] + ): + self.assertEqual(len(result), limit) + + def test_order_by_optimization(self): + """Test ORDER BY clause.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_medium) + + # Test ordering by different columns + result = ctx.sql( + 'SELECT lat, lon, air FROM air ORDER BY air DESC LIMIT 20' + ).to_pandas() + + if len(result) > 1: + # Should be in descending order + air_values = result['air'].values + self.assertTrue(np.all(air_values[:-1] >= air_values[1:])) + + # Test ordering by coordinates + result = ctx.sql( + 'SELECT lat, lon, air FROM air ORDER BY lat ASC, lon DESC LIMIT 20' + ).to_pandas() + + if len(result) > 1: + # Check lat is ascending + lat_values = result['lat'].values + lat_diffs = np.diff(lat_values) + # Allow for equal values (same lat, different lon) + self.assertTrue(np.all(lat_diffs >= 0)) + + def test_aggregation_pushdown(self): + """Test aggregation optimization.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_medium) + + # Test GROUP BY optimization + result = ctx.sql( + """ + SELECT + lat, + COUNT(*) as point_count, + AVG(temperature) as avg_temp, + MIN(precipitation) as min_precip, + MAX(precipitation) as max_precip + FROM weather + GROUP BY lat + ORDER BY lat + """ + ).to_pandas() + + # Should have one row per unique latitude + expected_lats = len(self.weather_medium.lat) + self.assertEqual(len(result), expected_lats) + + # All aggregation columns should be present + expected_agg_cols = {'point_count', 'avg_temp', 'min_precip', 'max_precip'} + self.assertTrue(expected_agg_cols.issubset(set(result.columns))) + + def test_complex_filter_optimization(self): + """Test complex filtering with multiple conditions.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_medium) + + # Complex WHERE clause with AND/OR + result = ctx.sql( + """ + SELECT lat, lon, temperature, precipitation + FROM weather + WHERE (temperature > 15 AND precipitation < 50) + OR (lat > 45 AND lon < -100) + LIMIT 100 + """ + ).to_pandas() + + # Verify the complex condition + if len(result) > 0: + condition1 = (result['temperature'] > 15) & (result['precipitation'] < 50) + condition2 = (result['lat'] > 45) & (result['lon'] < -100) + combined = condition1 | condition2 + self.assertTrue(combined.all()) + + +class SqlComplexQueryTestCase(unittest.TestCase): + """Test complex SQL queries with advanced features.""" + + def setUp(self): + """Set up datasets for complex query testing.""" + self.air = xr.tutorial.open_dataset('air_temperature') + self.air_small = self.air.isel( + time=slice(0, 50), lat=slice(0, 15), lon=slice(0, 20) + ).chunk({'time': 25}) + + self.weather = rand_wx('2023-01-01T00', '2023-01-01T12') + self.weather_small = self.weather.isel( + time=slice(0, 8), lat=slice(0, 10), lon=slice(0, 12) + ).chunk({'time': 4}) + + def test_subqueries(self): + """Test subqueries and CTEs.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + # Subquery to find above-average temperatures + result = ctx.sql( + """ + SELECT lat, lon, air + FROM air + WHERE air > ( + SELECT AVG(air) FROM air + ) + LIMIT 50 + """ + ).to_pandas() + + if len(result) > 0: + # Get the average to verify + avg_result = ctx.sql('SELECT AVG(air) as avg_air FROM air').to_pandas() + avg_air = avg_result['avg_air'].iloc[0] + + # All results should be above average + self.assertTrue((result['air'] > avg_air).all()) + + def test_window_functions(self): + """Test window functions if supported by DataFusion.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + try: + # Test ROW_NUMBER window function + result = ctx.sql( + """ + SELECT + lat, lon, air, + ROW_NUMBER() OVER (PARTITION BY lat ORDER BY air DESC) as rank_in_lat + FROM air + WHERE lat IN ( + SELECT DISTINCT lat FROM air LIMIT 3 + ) + ORDER BY lat, rank_in_lat + LIMIT 30 + """ + ).to_pandas() + + if len(result) > 0: + # Check that ranking works within each lat + for lat_val in result['lat'].unique(): + lat_data = result[result['lat'] == lat_val].sort_values('rank_in_lat') + if len(lat_data) > 1: + # Air values should be in descending order within each lat + air_values = lat_data['air'].values + self.assertTrue(np.all(air_values[:-1] >= air_values[1:])) + + except Exception: + # Window functions might not be supported, skip test + self.skipTest('Window functions not supported') + + def test_case_statements(self): + """Test CASE statements for conditional logic.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_small) + + result = ctx.sql( + """ + SELECT + lat, lon, + temperature, + CASE + WHEN temperature < 0 THEN 'Freezing' + WHEN temperature < 10 THEN 'Cold' + WHEN temperature < 20 THEN 'Cool' + WHEN temperature < 30 THEN 'Warm' + ELSE 'Hot' + END as temp_category, + precipitation, + CASE + WHEN precipitation < 5 THEN 'Dry' + WHEN precipitation < 20 THEN 'Light' + WHEN precipitation < 50 THEN 'Moderate' + ELSE 'Heavy' + END as precip_category + FROM weather + LIMIT 50 + """ + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertIn('temp_category', result.columns) + self.assertIn('precip_category', result.columns) + + # Verify categories are assigned correctly + for _, row in result.iterrows(): + temp = row['temperature'] + temp_cat = row['temp_category'] + + if temp < 0: + self.assertEqual(temp_cat, 'Freezing') + elif temp < 10: + self.assertEqual(temp_cat, 'Cold') + elif temp < 20: + self.assertEqual(temp_cat, 'Cool') + elif temp < 30: + self.assertEqual(temp_cat, 'Warm') + else: + self.assertEqual(temp_cat, 'Hot') + + def test_mathematical_functions(self): + """Test mathematical functions in SQL.""" + ctx = XarrayContext() + ctx.from_dataset('weather', self.weather_small) + + result = ctx.sql( + """ + SELECT + lat, lon, + temperature, + ROUND(temperature, 1) as temp_rounded, + ABS(temperature - 20) as temp_diff_from_20, + SQRT(ABS(temperature)) as temp_sqrt, + precipitation, + LOG(precipitation + 1) as log_precip + FROM weather + WHERE precipitation > 0 + LIMIT 30 + """ + ).to_pandas() + + if len(result) > 0: + # Verify mathematical operations + for _, row in result.iterrows(): + temp = row['temperature'] + self.assertAlmostEqual( + row['temp_diff_from_20'], abs(temp - 20), places=5 + ) + if temp >= 0: + self.assertAlmostEqual(row['temp_sqrt'], np.sqrt(temp), places=5) + + def test_string_operations(self): + """Test string operations if applicable.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + # Test string operations on numeric data converted to string + result = ctx.sql( + """ + SELECT + lat, lon, + CAST(lat AS VARCHAR) as lat_str, + CONCAT('Lat: ', CAST(lat AS VARCHAR), ', Lon: ', CAST(lon AS VARCHAR)) as coordinates + FROM air + LIMIT 10 + """ + ).to_pandas() + + self.assertGreater(len(result), 0) + self.assertIn('coordinates', result.columns) + + # Verify concatenation worked + for _, row in result.iterrows(): + expected = f"Lat: {row['lat']}, Lon: {row['lon']}" + # Allow for slight formatting differences + self.assertIn('Lat:', row['coordinates']) + self.assertIn('Lon:', row['coordinates']) + + +class SqlErrorHandlingTestCase(unittest.TestCase): + """Test error handling and edge cases in SQL interface.""" + + def setUp(self): + """Set up dataset for error testing.""" + self.air = xr.tutorial.open_dataset('air_temperature') + self.air_small = self.air.isel( + time=slice(0, 10), lat=slice(0, 5), lon=slice(0, 8) + ).chunk({'time': 5}) + + def test_invalid_table_name(self): + """Test error handling for invalid table names.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + with self.assertRaises(Exception): + ctx.sql('SELECT * FROM nonexistent_table') + + def test_invalid_column_name(self): + """Test error handling for invalid column names.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + with self.assertRaises(Exception): + ctx.sql('SELECT nonexistent_column FROM air') + + def test_syntax_errors(self): + """Test handling of SQL syntax errors.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + # Invalid SQL syntax + with self.assertRaises(Exception): + ctx.sql('SELECT * FORM air') # Typo: FORM instead of FROM + + with self.assertRaises(Exception): + ctx.sql('SELECT * FROM air WHERE') # Incomplete WHERE clause + + def test_empty_dataset_handling(self): + """Test handling of empty datasets.""" + # Create empty dataset + empty_ds = xr.Dataset({'temp': (['x'], [])}, coords={'x': []}).chunk( + {'x': 1} + ) + + ctx = XarrayContext() + ctx.from_dataset('empty', empty_ds) + + # Should handle empty dataset gracefully + result = ctx.sql('SELECT * FROM empty').to_pandas() + self.assertEqual(len(result), 0) + + # Aggregations on empty dataset + result = ctx.sql('SELECT COUNT(*) as count FROM empty').to_pandas() + self.assertEqual(len(result), 1) + self.assertEqual(result['count'].iloc[0], 0) + + def test_large_limit_handling(self): + """Test handling of very large LIMIT values.""" + ctx = XarrayContext() + ctx.from_dataset('air', self.air_small) + + total_rows = ( + self.air_small.sizes['time'] + * self.air_small.sizes['lat'] + * self.air_small.sizes['lon'] + ) + + # Request more rows than exist + result = ctx.sql(f'SELECT * FROM air LIMIT {total_rows * 10}').to_pandas() + self.assertEqual(len(result), total_rows) + + if __name__ == '__main__': unittest.main()