From 5540fe832fee5d0477d5b35476e3c0fdadba91f1 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 12 Jul 2025 17:30:16 -0700 Subject: [PATCH 01/10] Streaming experiment using record batches. --- xarray_sql/df.py | 26 +++++++++++++++++++++++--- xarray_sql/df_test.py | 1 + xarray_sql/sql.py | 1 + 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index af25c35a..a16fec9d 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,12 +1,14 @@ import itertools import typing as t +import datafusion import numpy as np import pandas as pd import pyarrow as pa import xarray as xr +from datafusion.context import ArrowStreamExportable -Block = t.Dict[str, slice] +Block = t.Dict[t.Hashable, slice] Chunks = t.Optional[t.Dict[str, int]] @@ -57,6 +59,20 @@ def _block_len(block: Block) -> int: return np.prod([v.stop - v.start for v in block.values()]) +def from_map_to_batches( + func: t.Callable[[...], pd.DataFrame], *iterables, args: t.Optional[t.Tuple] = None, schema: pa.Schema = None, **kwargs +) -> pa.RecordBatchReader: + if args is None: + args = () + + def apply_map(): + for items in zip(*iterables): + df = func(*items, *args, **kwargs) + yield pa.RecordBatch.from_pandas(df, schema=schema) + + return pa.RecordBatchReader.from_batches(schema, apply_map()) + + def from_map( func: t.Callable, *iterables, args: t.Optional[t.Tuple] = None, **kwargs ) -> pa.Table: @@ -109,7 +125,7 @@ def from_map( return pa.concat_tables(results) -def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.Table: +def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> ArrowStreamExportable: """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. Args: @@ -131,4 +147,8 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.Table: def pivot(b: Block) -> pd.DataFrame: return ds.isel(b).to_dataframe().reset_index() - return from_map(pivot, blocks) + schema = pa.Schema.from_pandas(pivot(blocks[0])) + last_schema = pa.Schema.from_pandas(pivot(blocks[-1])) + assert schema == last_schema, 'Schemas must be consistent across blocks!' + + return from_map_to_batches(pivot, blocks, schema=schema) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index fedef88c..448c9e08 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -81,6 +81,7 @@ def test_data_equal__one__last(self): self.assertEqual(self.air.isel(iselection), ds) +@unittest.skip("Chose non-pa.Table implementation.") class PyArrowTableTest(DaskTestCase): def test_sanity(self): diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 11553b99..1f82578b 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -15,3 +15,4 @@ def from_dataset( ): arrow_table = read_xarray(input_table, chunks) return self.from_arrow(arrow_table, table_name) + From 6e09fab854053fb6f570e38f0fe37fef7a4b74e2 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 12:13:40 -0700 Subject: [PATCH 02/10] Updated df docs and functionality. Added correctness tests for from_map_batched --- xarray_sql/df.py | 35 ++++++-- xarray_sql/df_test.py | 201 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 227 insertions(+), 9 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index a16fec9d..a8622151 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,7 +1,6 @@ import itertools import typing as t -import datafusion import numpy as np import pandas as pd import pyarrow as pa @@ -56,12 +55,28 @@ def explode(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[xr.Dataset]: def _block_len(block: Block) -> int: - return np.prod([v.stop - v.start for v in block.values()]) + return int(np.prod([v.stop - v.start for v in block.values()])) -def from_map_to_batches( +def from_map_batched( func: t.Callable[[...], pd.DataFrame], *iterables, args: t.Optional[t.Tuple] = None, schema: pa.Schema = None, **kwargs ) -> pa.RecordBatchReader: + """Create a PyArrow RecordBatchReader by mapping a function over iterables. + + This is equivalent to dask's from_map but returns a PyArrow + RecordBatchReader that can be used with DataFusion. It iterates over + RecordBatches which are created via the `func` one-at-a-time. + + Args: + func: Function to apply to each element of the iterables. + *iterables: Iterable objects to map the function over. + schema: Optional schema needed for the RecordBatchReader. + args: Additional positional arguments to pass to func. + **kwargs: Additional keyword arguments to pass to func. + + Returns: + A PyArrow RecordBatchReader containing the stream of RecordBatches. + """ if args is None: args = () @@ -125,6 +140,11 @@ def from_map( return pa.concat_tables(results) +def pivot(ds: xr.Dataset) -> pd.DataFrame: + """Converts an xarray Dataset to a pandas DataFrame.""" + return ds.to_dataframe().reset_index() + + def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> ArrowStreamExportable: """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. @@ -144,11 +164,10 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> ArrowStreamExportable: blocks = list(block_slices(ds, chunks)) - def pivot(b: Block) -> pd.DataFrame: - return ds.isel(b).to_dataframe().reset_index() + def pivot_block(b: Block): return pivot(ds.isel(b)) - schema = pa.Schema.from_pandas(pivot(blocks[0])) - last_schema = pa.Schema.from_pandas(pivot(blocks[-1])) + schema = pa.Schema.from_pandas(pivot_block(blocks[0])) + last_schema = pa.Schema.from_pandas(pivot_block(blocks[-1])) assert schema == last_schema, 'Schemas must be consistent across blocks!' - return from_map_to_batches(pivot, blocks, schema=schema) + return from_map_batched(pivot_block, blocks, schema=schema) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index 448c9e08..d920bdfb 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -6,7 +6,7 @@ import pyarrow as pa import xarray as xr -from .df import explode, read_xarray, block_slices, from_map +from .df import explode, read_xarray, block_slices, from_map, pivot, from_map_batched def rand_wx(start: str, end: str) -> xr.Dataset: @@ -34,6 +34,39 @@ def rand_wx(start: str, end: str) -> xr.Dataset: ) +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) + + # Create coordinates + 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) + + # Create large data arrays + 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 + + class DaskTestCase(unittest.TestCase): def setUp(self) -> None: @@ -179,5 +212,171 @@ def make_arrow_table(x): self.assertEqual(len(result), 3) +class TestFromMapBatchedCorrectness(DaskTestCase): + """Test correctness of from_map_batched function.""" + + def test_basic_functionality(self): + """Test that from_map_batched produces correct RecordBatchReader.""" + blocks = list(block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) + + # Get expected schema + first_block_df = pivot(self.air_small.isel(blocks[0])) + expected_schema = pa.Schema.from_pandas(first_block_df) + + # Create RecordBatchReader + reader = from_map_batched( + pivot, + [self.air_small.isel(block) for block in blocks], + schema=expected_schema + ) + + # Verify it's a RecordBatchReader + self.assertIsInstance(reader, pa.RecordBatchReader) + + # Verify schema + self.assertEqual(reader.schema, expected_schema) + + # Read all batches and verify content + batches = list(reader) + self.assertGreater(len(batches), 0) + + # Verify each batch has the correct schema + for batch in batches: + self.assertEqual(batch.schema, expected_schema) + self.assertGreater(len(batch), 0) + + def test_multiple_iterables(self): + """Test from_map_batched with 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 + ) + + # Read all data + table = reader.read_all() + df = table.to_pandas() + + # Verify results + 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_with_args_and_kwargs(self): + """Test from_map_batched with additional args and kwargs.""" + def multiply_and_add(x, multiplier, offset=0): + result = pd.DataFrame({ + 'x': [x], + 'result': [x * multiplier + offset] + }) + return result + + values = [1, 2, 3] + expected_schema = pa.schema([ + ('x', pa.int64()), + ('result', pa.int64()) + ]) + + reader = from_map_batched( + multiply_and_add, + values, + args=(2,), # multiplier = 2 + offset=5, # offset = 5 + schema=expected_schema + ) + + table = reader.read_all() + df = table.to_pandas() + + # Verify results: x * 2 + 5 + expected_df = pd.DataFrame({ + 'x': [1, 2, 3], + 'result': [7, 9, 11] # (1*2+5, 2*2+5, 3*2+5) + }) + + pd.testing.assert_frame_equal(df, expected_df) + + def test_empty_iterables(self): + """Test from_map_batched with 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) + self.assertEqual(len(batches), 0) + + def test_consistency_with_regular_map(self): + """Test that results are consistent with regular mapping.""" + blocks = list(block_slices(self.air_small, chunks={'time': 4, 'lat': 3})) + datasets = [self.air_small.isel(block) for block in blocks] + + # Get schema from first block + first_df = pivot(datasets[0]) + schema = pa.Schema.from_pandas(first_df) + + # Use from_map_batched + reader = from_map_batched(pivot, datasets, schema=schema) + batched_table = reader.read_all() + + # Regular map approach + regular_dfs = [pivot(ds) for ds in datasets] + regular_table = pa.Table.from_pandas(pd.concat(regular_dfs, ignore_index=True)) + + # Results should be identical + self.assertEqual(batched_table.schema, regular_table.schema) + self.assertEqual(len(batched_table), len(regular_table)) + + # Compare data (allowing for potential column order differences) + 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_integration_with_datafusion_via_read_xarray(self): + """Test integration with the read_xarray function that uses from_map_batched.""" + + # Use a small dataset + 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}) + + # read_xarray uses from_map_batched internally + arrow_stream = read_xarray(air_chunked, chunks={'time': 25, 'lat': 5, 'lon': 8}) + + # Verify it returns a proper ArrowStreamExportable (RecordBatchReader) + self.assertTrue(hasattr(arrow_stream, 'schema')) + self.assertTrue(hasattr(arrow_stream, '__iter__')) + + # Should be able to read all data + table = arrow_stream.read_all() + self.assertGreater(len(table), 0) + + # Should have the expected columns + expected_columns = {'time', 'lat', 'lon', 'air'} + actual_columns = set(table.column_names) + self.assertTrue(expected_columns.issubset(actual_columns)) + + + + + if __name__ == '__main__': unittest.main() From 75a3d6936206b36b176627dff86055c1de042e77 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 12:58:39 -0700 Subject: [PATCH 03/10] Added a test that provides memory bounds for read_xarray. --- xarray_sql/df_test.py | 47 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index d920bdfb..9eec5730 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -1,4 +1,5 @@ import itertools +import tracemalloc import unittest import numpy as np @@ -375,6 +376,52 @@ def test_integration_with_datafusion_via_read_xarray(self): self.assertTrue(expected_columns.issubset(actual_columns)) +class ReadXarrayStreamingTest(unittest.TestCase): + def setUp(self): + self.large_ds = create_large_dataset().chunk({'time': 25}) + + def test_read_xarray__loads_one_chunk_at_a_time(self): + tracemalloc.start() + iterable = read_xarray(self.large_ds) + first_size, first_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + + sizes, peaks = [], [] + + first_chunk = self.large_ds.isel(next(block_slices(self.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 + self.assertLess(first_size, chunk_size * 3) + self.assertLess(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) + + mean_size = np.mean(sizes) + mean_peak = np.mean(peaks) + + # Provide a bound for each size and peak + for size in sizes: + self.assertGreater(mean_size*1.1, size) + self.assertGreater(chunk_size * 3, size) + # malloc size is about 2.66x chunk_size on average + self.assertLess(chunk_size * 2, size) + + for peak in peaks: + self.assertGreater(mean_peak*1.1, peak) + self.assertGreater(chunk_size * 7, peak) + # malloc peak is about 6.89x chunk_size on average + self.assertLess(chunk_size * 4, peak) + + tracemalloc.stop() + + + From 1e3ff80e38842d0fc39ed68d6d6b73e6c4b49f24 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:00:40 -0700 Subject: [PATCH 04/10] Added a peak memory assert. --- xarray_sql/df_test.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index 9eec5730..ace628ed 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -418,6 +418,9 @@ def test_read_xarray__loads_one_chunk_at_a_time(self): # malloc peak is about 6.89x chunk_size on average self.assertLess(chunk_size * 4, peak) + # The peak malloc should never be more than the original dataset! + self.assertLess(max(peaks), self.large_ds.nbytes) + tracemalloc.stop() From dce91a1fe31091550c3905a6e8a842d2bb5a875f Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:00:59 -0700 Subject: [PATCH 05/10] reformat --- xarray_sql/df_test.py | 288 +++++++++++++++++++++--------------------- 1 file changed, 144 insertions(+), 144 deletions(-) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index ace628ed..fd2d4fbe 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -11,28 +11,28 @@ 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.'), - ) + 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): @@ -70,147 +70,147 @@ def adding_function(x, y): class DaskTestCase(unittest.TestCase): - def setUp(self) -> None: - self.air = xr.tutorial.open_dataset('air_temperature') - self.chunks = {'time': 240} - self.air = self.air.chunk(self.chunks) + def setUp(self) -> None: + self.air = xr.tutorial.open_dataset('air_temperature') + self.chunks = {'time': 240} + self.air = self.air.chunk(self.chunks) - self.air_small = self.air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(self.chunks) - self.randwx = rand_wx('1995-01-13T00', '1995-01-13T01') + self.air_small = self.air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(self.chunks) + self.randwx = rand_wx('1995-01-13T00', '1995-01-13T01') class ExplodeTest(DaskTestCase): - def test_cardinality(self): - dss = explode(self.air) - self.assertEqual( - len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) - ) + def test_cardinality(self): + dss = explode(self.air) + self.assertEqual( + len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) + ) - def test_dim_sizes__one(self): - ds = next(iter(explode(self.air))) - for k, v in self.chunks.items(): - self.assertIn(k, ds.dims) - self.assertEqual(v, ds.sizes[k]) - - def skip_test_dim_sizes__all(self): - # TODO(alxmrs): Why is this test slow? - dss = explode(self.air) - self.assertEqual( - [tuple(ds.dims.values()) for ds in dss], - list(itertools.product(*self.air.chunksizes.values())), - ) + def test_dim_sizes__one(self): + ds = next(iter(explode(self.air))) + for k, v in self.chunks.items(): + self.assertIn(k, ds.dims) + self.assertEqual(v, ds.sizes[k]) + + def skip_test_dim_sizes__all(self): + # TODO(alxmrs): Why is this test slow? + dss = explode(self.air) + self.assertEqual( + [tuple(ds.dims.values()) for ds in dss], + list(itertools.product(*self.air.chunksizes.values())), + ) - def test_data_equal__one__first(self): - ds = next(iter(explode(self.air))) - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - self.assertEqual(self.air.isel(iselection), ds) + def test_data_equal__one__first(self): + ds = next(iter(explode(self.air))) + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + self.assertEqual(self.air.isel(iselection), ds) - def test_data_equal__one__last(self): - dss = list(explode(self.air)) - ds = dss[-1] - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - self.assertEqual(self.air.isel(iselection), ds) + def test_data_equal__one__last(self): + dss = list(explode(self.air)) + ds = dss[-1] + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + self.assertEqual(self.air.isel(iselection), ds) @unittest.skip("Chose non-pa.Table implementation.") class PyArrowTableTest(DaskTestCase): - def test_sanity(self): - table = read_xarray(self.air_small) - self.assertIsNotNone(table) - self.assertIsInstance(table, pa.Table) - self.assertEqual(len(table), np.prod(list(self.air_small.sizes.values()))) - - def test_columns(self): - table = read_xarray(self.air_small) - cols = table.column_names - self.assertEqual(cols, ['lat', 'time', 'lon', 'air']) - - def test_dtypes(self): - table = read_xarray(self.air_small) - # Convert to pandas to check dtypes - df = table.to_pandas() - types = list(df.dtypes) - self.assertEqual([self.air_small[c].dtype for c in df.columns], types) - - def test_different_chunk_sizes(self): - default_table = read_xarray(self.air_small) - chunked_table = read_xarray(self.air_small, dict(time=5)) - - # Both should produce valid tables - self.assertIsInstance(default_table, pa.Table) - self.assertIsInstance(chunked_table, pa.Table) - # Should have same number of rows - self.assertEqual(len(default_table), len(chunked_table)) - - def test_chunk_perf(self): - table = read_xarray(self.air, chunks=dict(time=6)) - self.assertIsNotNone(table) - self.assertEqual(len(table), np.prod(list(self.air.sizes.values()))) - - def test_column_metadata_preserved(self): - try: - table = read_xarray(self.randwx, chunks=dict(time=24)) - self.assertIsInstance(table, pa.Table) - except Exception as e: - self.fail(f'Unexpected error: {e}') + def test_sanity(self): + table = read_xarray(self.air_small) + self.assertIsNotNone(table) + self.assertIsInstance(table, pa.Table) + self.assertEqual(len(table), np.prod(list(self.air_small.sizes.values()))) + + def test_columns(self): + table = read_xarray(self.air_small) + cols = table.column_names + self.assertEqual(cols, ['lat', 'time', 'lon', 'air']) + + def test_dtypes(self): + table = read_xarray(self.air_small) + # Convert to pandas to check dtypes + df = table.to_pandas() + types = list(df.dtypes) + self.assertEqual([self.air_small[c].dtype for c in df.columns], types) + + def test_different_chunk_sizes(self): + default_table = read_xarray(self.air_small) + chunked_table = read_xarray(self.air_small, dict(time=5)) + + # Both should produce valid tables + self.assertIsInstance(default_table, pa.Table) + self.assertIsInstance(chunked_table, pa.Table) + # Should have same number of rows + self.assertEqual(len(default_table), len(chunked_table)) + + def test_chunk_perf(self): + table = read_xarray(self.air, chunks=dict(time=6)) + self.assertIsNotNone(table) + self.assertEqual(len(table), np.prod(list(self.air.sizes.values()))) + + def test_column_metadata_preserved(self): + try: + table = read_xarray(self.randwx, chunks=dict(time=24)) + self.assertIsInstance(table, pa.Table) + except Exception as e: + self.fail(f'Unexpected error: {e}') class FromMapTest(unittest.TestCase): - def test_basic_from_map(self): - """Test basic from_map functionality with pandas DataFrames.""" + def test_basic_from_map(self): + """Test basic from_map functionality with pandas DataFrames.""" - def make_df(x): - return pd.DataFrame({'value': [x, x * 2], 'index': [0, 1]}) + def make_df(x): + return pd.DataFrame({'value': [x, x * 2], 'index': [0, 1]}) - result = from_map(make_df, [1, 2, 3]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 6) # 3 inputs * 2 rows each - self.assertEqual(result.column_names, ['value', 'index']) + result = from_map(make_df, [1, 2, 3]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 6) # 3 inputs * 2 rows each + self.assertEqual(result.column_names, ['value', 'index']) - def test_from_map_with_multiple_iterables(self): - """Test from_map with multiple iterables.""" + def test_from_map_with_multiple_iterables(self): + """Test from_map with multiple iterables.""" - def add_values(x, y): - return pd.DataFrame({'sum': [x + y], 'x': [x], 'y': [y]}) + def add_values(x, y): + return pd.DataFrame({'sum': [x + y], 'x': [x], 'y': [y]}) - result = from_map(add_values, [1, 2], [10, 20]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 2) + result = from_map(add_values, [1, 2], [10, 20]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 2) - # Convert to pandas to check values - df = result.to_pandas() - self.assertEqual(list(df['sum']), [11, 22]) + # Convert to pandas to check values + df = result.to_pandas() + self.assertEqual(list(df['sum']), [11, 22]) - def test_from_map_with_args(self): - """Test from_map with additional arguments.""" + def test_from_map_with_args(self): + """Test from_map with additional arguments.""" - def multiply_and_add(x, multiplier, add_value): - return pd.DataFrame({'result': [x * multiplier + add_value]}) + 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)) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 3) + result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 3) - df = result.to_pandas() - self.assertEqual( - list(df['result']), [12, 14, 16] - ) # (1*2+10, 2*2+10, 3*2+10) + df = result.to_pandas() + self.assertEqual( + list(df['result']), [12, 14, 16] + ) # (1*2+10, 2*2+10, 3*2+10) - def test_from_map_with_pyarrow_tables(self): - """Test from_map when function returns PyArrow tables.""" + def test_from_map_with_pyarrow_tables(self): + """Test from_map when function returns PyArrow tables.""" - def make_arrow_table(x): - df = pd.DataFrame({'value': [x]}) - return pa.Table.from_pandas(df) + 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]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 3) + result = from_map(make_arrow_table, [1, 2, 3]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 3) class TestFromMapBatchedCorrectness(DaskTestCase): @@ -218,7 +218,8 @@ class TestFromMapBatchedCorrectness(DaskTestCase): def test_basic_functionality(self): """Test that from_map_batched produces correct RecordBatchReader.""" - blocks = list(block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) + blocks = list( + block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) # Get expected schema first_block_df = pivot(self.air_small.isel(blocks[0])) @@ -279,6 +280,7 @@ def test_multiple_iterables(self): def test_with_args_and_kwargs(self): """Test from_map_batched with additional args and kwargs.""" + def multiply_and_add(x, multiplier, offset=0): result = pd.DataFrame({ 'x': [x], @@ -296,7 +298,7 @@ def multiply_and_add(x, multiplier, offset=0): multiply_and_add, values, args=(2,), # multiplier = 2 - offset=5, # offset = 5 + offset=5, # offset = 5 schema=expected_schema ) @@ -346,8 +348,10 @@ def test_consistency_with_regular_map(self): self.assertEqual(len(batched_table), len(regular_table)) # Compare data (allowing for potential column order differences) - 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) + 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) @@ -407,13 +411,13 @@ def test_read_xarray__loads_one_chunk_at_a_time(self): # Provide a bound for each size and peak for size in sizes: - self.assertGreater(mean_size*1.1, size) + self.assertGreater(mean_size * 1.1, size) self.assertGreater(chunk_size * 3, size) # malloc size is about 2.66x chunk_size on average self.assertLess(chunk_size * 2, size) for peak in peaks: - self.assertGreater(mean_peak*1.1, peak) + self.assertGreater(mean_peak * 1.1, peak) self.assertGreater(chunk_size * 7, peak) # malloc peak is about 6.89x chunk_size on average self.assertLess(chunk_size * 4, peak) @@ -424,9 +428,5 @@ def test_read_xarray__loads_one_chunk_at_a_time(self): tracemalloc.stop() - - - - if __name__ == '__main__': - unittest.main() + unittest.main() From 586c25309c02cec7675a0356234134133e6b18ee Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:20:04 -0700 Subject: [PATCH 06/10] Apply consistent indentation style --- xarray_sql/core.py | 2 +- xarray_sql/df.py | 19 +- xarray_sql/df_integrationtest.py | 24 +- xarray_sql/df_test.py | 724 +++++++++++++++---------------- xarray_sql/sql.py | 1 - xarray_sql/sql_test.py | 18 +- 6 files changed, 394 insertions(+), 394 deletions(-) diff --git a/xarray_sql/core.py b/xarray_sql/core.py index afbd2eca..f0a777a4 100644 --- a/xarray_sql/core.py +++ b/xarray_sql/core.py @@ -39,7 +39,7 @@ def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: prod_vals = (ds.coords[k].values for k in dim_keys) coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( - -1, len(dim_keys) + -1, len(dim_keys) ) for i, d in enumerate(dim_keys): diff --git a/xarray_sql/df.py b/xarray_sql/df.py index a8622151..301be9db 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -18,7 +18,7 @@ def _get_chunk_slicer( if dim in chunk_index: which_chunk = chunk_index[dim] return slice( - chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] + chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] ) return slice(None) @@ -40,11 +40,11 @@ def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]: ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) blocks = ( - { - dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) - for dim in ds.dims - } - for chunk_index in chunk_idxs + { + dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) + for dim in ds.dims + } + for chunk_index in chunk_idxs ) yield from blocks @@ -59,7 +59,8 @@ def _block_len(block: Block) -> int: def from_map_batched( - func: t.Callable[[...], pd.DataFrame], *iterables, args: t.Optional[t.Tuple] = None, schema: pa.Schema = None, **kwargs + func: t.Callable[[...], pd.DataFrame], *iterables, args: t.Optional[t.Tuple] = None, + schema: pa.Schema = None, **kwargs ) -> pa.RecordBatchReader: """Create a PyArrow RecordBatchReader by mapping a function over iterables. @@ -128,7 +129,7 @@ def from_map( pa_table = pa.Table.from_pandas(df) except Exception as e: raise ValueError( - f"Cannot convert function result to PyArrow Table: {e}" + f"Cannot convert function result to PyArrow Table: {e}" ) results.append(pa_table) @@ -159,7 +160,7 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> ArrowStreamExportable: """ fst = next(iter(ds.values())).dims assert all( - da.dims == fst for da in ds.values() + da.dims == fst for da in ds.values() ), "All dimensions must be equal. Please filter data_vars in the Dataset." blocks = list(block_slices(ds, chunks)) diff --git a/xarray_sql/df_integrationtest.py b/xarray_sql/df_integrationtest.py index 6f0eae4d..fe438a86 100644 --- a/xarray_sql/df_integrationtest.py +++ b/xarray_sql/df_integrationtest.py @@ -9,21 +9,21 @@ class Era5TestCast(unittest.TestCase): def test_open_era5(self): era5_ds = xr.open_zarr( - 'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2', - chunks={'time': 240, 'level': 1}, + 'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2', + chunks={'time': 240, 'level': 1}, ) era5_wind_df = read_xarray( - era5_ds[['u_component_of_wind', 'v_component_of_wind']] + era5_ds[['u_component_of_wind', 'v_component_of_wind']] ) self.assertEqual( - list(era5_wind_df.columns), - [ - 'time', - 'level', - 'latitude', - 'longitude', - 'u_component_of_wind', - 'v_component_of_wind', - ], + list(era5_wind_df.columns), + [ + 'time', + 'level', + 'latitude', + 'longitude', + 'u_component_of_wind', + 'v_component_of_wind', + ], ) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index fd2d4fbe..f825e90b 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -11,422 +11,422 @@ 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.'), - ) + 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) - - # Create coordinates - 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) - - # Create large data arrays - 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, - }) + """Create a large xarray dataset for memory testing.""" + np.random.seed(42) + + # Create coordinates + 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) + + # Create large data arrays + 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 + """Simple function that adds two values and returns a DataFrame.""" + result = pd.DataFrame({ + 'x': [x], + 'y': [y], + 'sum': [x + y] + }) + return result class DaskTestCase(unittest.TestCase): - def setUp(self) -> None: - self.air = xr.tutorial.open_dataset('air_temperature') - self.chunks = {'time': 240} - self.air = self.air.chunk(self.chunks) + def setUp(self) -> None: + self.air = xr.tutorial.open_dataset('air_temperature') + self.chunks = {'time': 240} + self.air = self.air.chunk(self.chunks) - self.air_small = self.air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(self.chunks) - self.randwx = rand_wx('1995-01-13T00', '1995-01-13T01') + self.air_small = self.air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(self.chunks) + self.randwx = rand_wx('1995-01-13T00', '1995-01-13T01') class ExplodeTest(DaskTestCase): - def test_cardinality(self): - dss = explode(self.air) - self.assertEqual( - len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) - ) - - def test_dim_sizes__one(self): - ds = next(iter(explode(self.air))) - for k, v in self.chunks.items(): - self.assertIn(k, ds.dims) - self.assertEqual(v, ds.sizes[k]) - - def skip_test_dim_sizes__all(self): - # TODO(alxmrs): Why is this test slow? - dss = explode(self.air) - self.assertEqual( - [tuple(ds.dims.values()) for ds in dss], - list(itertools.product(*self.air.chunksizes.values())), - ) - - def test_data_equal__one__first(self): - ds = next(iter(explode(self.air))) - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - self.assertEqual(self.air.isel(iselection), ds) - - def test_data_equal__one__last(self): - dss = list(explode(self.air)) - ds = dss[-1] - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - self.assertEqual(self.air.isel(iselection), ds) + def test_cardinality(self): + dss = explode(self.air) + self.assertEqual( + len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) + ) + + def test_dim_sizes__one(self): + ds = next(iter(explode(self.air))) + for k, v in self.chunks.items(): + self.assertIn(k, ds.dims) + self.assertEqual(v, ds.sizes[k]) + + def skip_test_dim_sizes__all(self): + # TODO(alxmrs): Why is this test slow? + dss = explode(self.air) + self.assertEqual( + [tuple(ds.dims.values()) for ds in dss], + list(itertools.product(*self.air.chunksizes.values())), + ) + + def test_data_equal__one__first(self): + ds = next(iter(explode(self.air))) + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + self.assertEqual(self.air.isel(iselection), ds) + + def test_data_equal__one__last(self): + dss = list(explode(self.air)) + ds = dss[-1] + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + self.assertEqual(self.air.isel(iselection), ds) @unittest.skip("Chose non-pa.Table implementation.") class PyArrowTableTest(DaskTestCase): - def test_sanity(self): - table = read_xarray(self.air_small) - self.assertIsNotNone(table) - self.assertIsInstance(table, pa.Table) - self.assertEqual(len(table), np.prod(list(self.air_small.sizes.values()))) - - def test_columns(self): - table = read_xarray(self.air_small) - cols = table.column_names - self.assertEqual(cols, ['lat', 'time', 'lon', 'air']) - - def test_dtypes(self): - table = read_xarray(self.air_small) - # Convert to pandas to check dtypes - df = table.to_pandas() - types = list(df.dtypes) - self.assertEqual([self.air_small[c].dtype for c in df.columns], types) - - def test_different_chunk_sizes(self): - default_table = read_xarray(self.air_small) - chunked_table = read_xarray(self.air_small, dict(time=5)) - - # Both should produce valid tables - self.assertIsInstance(default_table, pa.Table) - self.assertIsInstance(chunked_table, pa.Table) - # Should have same number of rows - self.assertEqual(len(default_table), len(chunked_table)) - - def test_chunk_perf(self): - table = read_xarray(self.air, chunks=dict(time=6)) - self.assertIsNotNone(table) - self.assertEqual(len(table), np.prod(list(self.air.sizes.values()))) - - def test_column_metadata_preserved(self): - try: - table = read_xarray(self.randwx, chunks=dict(time=24)) - self.assertIsInstance(table, pa.Table) - except Exception as e: - self.fail(f'Unexpected error: {e}') + def test_sanity(self): + table = read_xarray(self.air_small) + self.assertIsNotNone(table) + self.assertIsInstance(table, pa.Table) + self.assertEqual(len(table), np.prod(list(self.air_small.sizes.values()))) + + def test_columns(self): + table = read_xarray(self.air_small) + cols = table.column_names + self.assertEqual(cols, ['lat', 'time', 'lon', 'air']) + + def test_dtypes(self): + table = read_xarray(self.air_small) + # Convert to pandas to check dtypes + df = table.to_pandas() + types = list(df.dtypes) + self.assertEqual([self.air_small[c].dtype for c in df.columns], types) + + def test_different_chunk_sizes(self): + default_table = read_xarray(self.air_small) + chunked_table = read_xarray(self.air_small, dict(time=5)) + + # Both should produce valid tables + self.assertIsInstance(default_table, pa.Table) + self.assertIsInstance(chunked_table, pa.Table) + # Should have same number of rows + self.assertEqual(len(default_table), len(chunked_table)) + + def test_chunk_perf(self): + table = read_xarray(self.air, chunks=dict(time=6)) + self.assertIsNotNone(table) + self.assertEqual(len(table), np.prod(list(self.air.sizes.values()))) + + def test_column_metadata_preserved(self): + try: + table = read_xarray(self.randwx, chunks=dict(time=24)) + self.assertIsInstance(table, pa.Table) + except Exception as e: + self.fail(f'Unexpected error: {e}') class FromMapTest(unittest.TestCase): - def test_basic_from_map(self): - """Test basic from_map functionality with pandas DataFrames.""" + def test_basic_from_map(self): + """Test basic from_map functionality with pandas DataFrames.""" - def make_df(x): - return pd.DataFrame({'value': [x, x * 2], 'index': [0, 1]}) + def make_df(x): + return pd.DataFrame({'value': [x, x * 2], 'index': [0, 1]}) - result = from_map(make_df, [1, 2, 3]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 6) # 3 inputs * 2 rows each - self.assertEqual(result.column_names, ['value', 'index']) + result = from_map(make_df, [1, 2, 3]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 6) # 3 inputs * 2 rows each + self.assertEqual(result.column_names, ['value', 'index']) - def test_from_map_with_multiple_iterables(self): - """Test from_map with multiple iterables.""" + def test_from_map_with_multiple_iterables(self): + """Test from_map with multiple iterables.""" - def add_values(x, y): - return pd.DataFrame({'sum': [x + y], 'x': [x], 'y': [y]}) + def add_values(x, y): + return pd.DataFrame({'sum': [x + y], 'x': [x], 'y': [y]}) - result = from_map(add_values, [1, 2], [10, 20]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 2) + result = from_map(add_values, [1, 2], [10, 20]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 2) - # Convert to pandas to check values - df = result.to_pandas() - self.assertEqual(list(df['sum']), [11, 22]) + # Convert to pandas to check values + df = result.to_pandas() + self.assertEqual(list(df['sum']), [11, 22]) - def test_from_map_with_args(self): - """Test from_map with additional arguments.""" + def test_from_map_with_args(self): + """Test from_map with additional arguments.""" - def multiply_and_add(x, multiplier, add_value): - return pd.DataFrame({'result': [x * multiplier + add_value]}) + 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)) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 3) + result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 3) - df = result.to_pandas() - self.assertEqual( - list(df['result']), [12, 14, 16] - ) # (1*2+10, 2*2+10, 3*2+10) + df = result.to_pandas() + self.assertEqual( + list(df['result']), [12, 14, 16] + ) # (1*2+10, 2*2+10, 3*2+10) - def test_from_map_with_pyarrow_tables(self): - """Test from_map when function returns PyArrow tables.""" + def test_from_map_with_pyarrow_tables(self): + """Test from_map when function returns PyArrow tables.""" - def make_arrow_table(x): - df = pd.DataFrame({'value': [x]}) - return pa.Table.from_pandas(df) + 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]) - self.assertIsInstance(result, pa.Table) - self.assertEqual(len(result), 3) + result = from_map(make_arrow_table, [1, 2, 3]) + self.assertIsInstance(result, pa.Table) + self.assertEqual(len(result), 3) class TestFromMapBatchedCorrectness(DaskTestCase): - """Test correctness of from_map_batched function.""" - - def test_basic_functionality(self): - """Test that from_map_batched produces correct RecordBatchReader.""" - blocks = list( - block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) - - # Get expected schema - first_block_df = pivot(self.air_small.isel(blocks[0])) - expected_schema = pa.Schema.from_pandas(first_block_df) - - # Create RecordBatchReader - reader = from_map_batched( - pivot, - [self.air_small.isel(block) for block in blocks], - schema=expected_schema - ) - - # Verify it's a RecordBatchReader - self.assertIsInstance(reader, pa.RecordBatchReader) - - # Verify schema - self.assertEqual(reader.schema, expected_schema) - - # Read all batches and verify content - batches = list(reader) - self.assertGreater(len(batches), 0) - - # Verify each batch has the correct schema - for batch in batches: - self.assertEqual(batch.schema, expected_schema) - self.assertGreater(len(batch), 0) - - def test_multiple_iterables(self): - """Test from_map_batched with 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 - ) - - # Read all data - table = reader.read_all() - df = table.to_pandas() - - # Verify results - 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_with_args_and_kwargs(self): - """Test from_map_batched with additional args and kwargs.""" - - def multiply_and_add(x, multiplier, offset=0): - result = pd.DataFrame({ - 'x': [x], - 'result': [x * multiplier + offset] - }) - return result - - values = [1, 2, 3] - expected_schema = pa.schema([ - ('x', pa.int64()), - ('result', pa.int64()) - ]) - - reader = from_map_batched( - multiply_and_add, - values, - args=(2,), # multiplier = 2 - offset=5, # offset = 5 - schema=expected_schema - ) - - table = reader.read_all() - df = table.to_pandas() - - # Verify results: x * 2 + 5 - expected_df = pd.DataFrame({ - 'x': [1, 2, 3], - 'result': [7, 9, 11] # (1*2+5, 2*2+5, 3*2+5) - }) - - pd.testing.assert_frame_equal(df, expected_df) - - def test_empty_iterables(self): - """Test from_map_batched with 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) - self.assertEqual(len(batches), 0) - - def test_consistency_with_regular_map(self): - """Test that results are consistent with regular mapping.""" - blocks = list(block_slices(self.air_small, chunks={'time': 4, 'lat': 3})) - datasets = [self.air_small.isel(block) for block in blocks] - - # Get schema from first block - first_df = pivot(datasets[0]) - schema = pa.Schema.from_pandas(first_df) - - # Use from_map_batched - reader = from_map_batched(pivot, datasets, schema=schema) - batched_table = reader.read_all() - - # Regular map approach - regular_dfs = [pivot(ds) for ds in datasets] - regular_table = pa.Table.from_pandas(pd.concat(regular_dfs, ignore_index=True)) - - # Results should be identical - self.assertEqual(batched_table.schema, regular_table.schema) - self.assertEqual(len(batched_table), len(regular_table)) - - # Compare data (allowing for potential column order differences) - 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_integration_with_datafusion_via_read_xarray(self): - """Test integration with the read_xarray function that uses from_map_batched.""" - - # Use a small dataset - 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}) - - # read_xarray uses from_map_batched internally - arrow_stream = read_xarray(air_chunked, chunks={'time': 25, 'lat': 5, 'lon': 8}) - - # Verify it returns a proper ArrowStreamExportable (RecordBatchReader) - self.assertTrue(hasattr(arrow_stream, 'schema')) - self.assertTrue(hasattr(arrow_stream, '__iter__')) - - # Should be able to read all data - table = arrow_stream.read_all() - self.assertGreater(len(table), 0) + """Test correctness of from_map_batched function.""" + + def test_basic_functionality(self): + """Test that from_map_batched produces correct RecordBatchReader.""" + blocks = list( + block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) + + # Get expected schema + first_block_df = pivot(self.air_small.isel(blocks[0])) + expected_schema = pa.Schema.from_pandas(first_block_df) + + # Create RecordBatchReader + reader = from_map_batched( + pivot, + [self.air_small.isel(block) for block in blocks], + schema=expected_schema + ) + + # Verify it's a RecordBatchReader + self.assertIsInstance(reader, pa.RecordBatchReader) + + # Verify schema + self.assertEqual(reader.schema, expected_schema) + + # Read all batches and verify content + batches = list(reader) + self.assertGreater(len(batches), 0) + + # Verify each batch has the correct schema + for batch in batches: + self.assertEqual(batch.schema, expected_schema) + self.assertGreater(len(batch), 0) + + def test_multiple_iterables(self): + """Test from_map_batched with 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 + ) + + # Read all data + table = reader.read_all() + df = table.to_pandas() + + # Verify results + 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_with_args_and_kwargs(self): + """Test from_map_batched with additional args and kwargs.""" + + def multiply_and_add(x, multiplier, offset=0): + result = pd.DataFrame({ + 'x': [x], + 'result': [x * multiplier + offset] + }) + return result + + values = [1, 2, 3] + expected_schema = pa.schema([ + ('x', pa.int64()), + ('result', pa.int64()) + ]) + + reader = from_map_batched( + multiply_and_add, + values, + args=(2,), # multiplier = 2 + offset=5, # offset = 5 + schema=expected_schema + ) + + table = reader.read_all() + df = table.to_pandas() + + # Verify results: x * 2 + 5 + expected_df = pd.DataFrame({ + 'x': [1, 2, 3], + 'result': [7, 9, 11] # (1*2+5, 2*2+5, 3*2+5) + }) + + pd.testing.assert_frame_equal(df, expected_df) + + def test_empty_iterables(self): + """Test from_map_batched with 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) + self.assertEqual(len(batches), 0) + + def test_consistency_with_regular_map(self): + """Test that results are consistent with regular mapping.""" + blocks = list(block_slices(self.air_small, chunks={'time': 4, 'lat': 3})) + datasets = [self.air_small.isel(block) for block in blocks] + + # Get schema from first block + first_df = pivot(datasets[0]) + schema = pa.Schema.from_pandas(first_df) + + # Use from_map_batched + reader = from_map_batched(pivot, datasets, schema=schema) + batched_table = reader.read_all() + + # Regular map approach + regular_dfs = [pivot(ds) for ds in datasets] + regular_table = pa.Table.from_pandas(pd.concat(regular_dfs, ignore_index=True)) + + # Results should be identical + self.assertEqual(batched_table.schema, regular_table.schema) + self.assertEqual(len(batched_table), len(regular_table)) + + # Compare data (allowing for potential column order differences) + 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_integration_with_datafusion_via_read_xarray(self): + """Test integration with the read_xarray function that uses from_map_batched.""" + + # Use a small dataset + 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}) + + # read_xarray uses from_map_batched internally + arrow_stream = read_xarray(air_chunked, chunks={'time': 25, 'lat': 5, 'lon': 8}) + + # Verify it returns a proper ArrowStreamExportable (RecordBatchReader) + self.assertTrue(hasattr(arrow_stream, 'schema')) + self.assertTrue(hasattr(arrow_stream, '__iter__')) + + # Should be able to read all data + table = arrow_stream.read_all() + self.assertGreater(len(table), 0) - # Should have the expected columns - expected_columns = {'time', 'lat', 'lon', 'air'} - actual_columns = set(table.column_names) - self.assertTrue(expected_columns.issubset(actual_columns)) + # Should have the expected columns + expected_columns = {'time', 'lat', 'lon', 'air'} + actual_columns = set(table.column_names) + self.assertTrue(expected_columns.issubset(actual_columns)) class ReadXarrayStreamingTest(unittest.TestCase): - def setUp(self): - self.large_ds = create_large_dataset().chunk({'time': 25}) + def setUp(self): + self.large_ds = create_large_dataset().chunk({'time': 25}) - def test_read_xarray__loads_one_chunk_at_a_time(self): - tracemalloc.start() - iterable = read_xarray(self.large_ds) - first_size, first_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() + def test_read_xarray__loads_one_chunk_at_a_time(self): + tracemalloc.start() + iterable = read_xarray(self.large_ds) + first_size, first_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() - sizes, peaks = [], [] + sizes, peaks = [], [] - first_chunk = self.large_ds.isel(next(block_slices(self.large_ds))) - chunk_size = first_chunk.nbytes + first_chunk = self.large_ds.isel(next(block_slices(self.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 - self.assertLess(first_size, chunk_size * 3) - self.assertLess(first_peak, chunk_size * 6) + # Creating the iterator should be inexpensive -- less than one chunk. + # We multiply by constant factors because chunks have additional overhead + self.assertLess(first_size, chunk_size * 3) + self.assertLess(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 it in iterable: + _ = it + cur_size, cur_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + sizes.append(cur_size), peaks.append(cur_peak) - mean_size = np.mean(sizes) - mean_peak = np.mean(peaks) + mean_size = np.mean(sizes) + mean_peak = np.mean(peaks) - # Provide a bound for each size and peak - for size in sizes: - self.assertGreater(mean_size * 1.1, size) - self.assertGreater(chunk_size * 3, size) - # malloc size is about 2.66x chunk_size on average - self.assertLess(chunk_size * 2, size) + # Provide a bound for each size and peak + for size in sizes: + self.assertGreater(mean_size * 1.1, size) + self.assertGreater(chunk_size * 3, size) + # malloc size is about 2.66x chunk_size on average + self.assertLess(chunk_size * 2, size) - for peak in peaks: - self.assertGreater(mean_peak * 1.1, peak) - self.assertGreater(chunk_size * 7, peak) - # malloc peak is about 6.89x chunk_size on average - self.assertLess(chunk_size * 4, peak) + for peak in peaks: + self.assertGreater(mean_peak * 1.1, peak) + self.assertGreater(chunk_size * 7, peak) + # malloc peak is about 6.89x chunk_size on average + self.assertLess(chunk_size * 4, peak) - # The peak malloc should never be more than the original dataset! - self.assertLess(max(peaks), self.large_ds.nbytes) + # The peak malloc should never be more than the original dataset! + self.assertLess(max(peaks), self.large_ds.nbytes) - tracemalloc.stop() + tracemalloc.stop() if __name__ == '__main__': - unittest.main() + unittest.main() diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 1f82578b..11553b99 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -15,4 +15,3 @@ def from_dataset( ): arrow_table = read_xarray(input_table, chunks) return self.from_arrow(arrow_table, table_name) - diff --git a/xarray_sql/sql_test.py b/xarray_sql/sql_test.py index f550f290..f9a68e6f 100644 --- a/xarray_sql/sql_test.py +++ b/xarray_sql/sql_test.py @@ -22,14 +22,14 @@ def test_agg_small(self): c.from_dataset('air', self.air_small) query = c.sql( - """ - SELECT - "lat", "lon", SUM("air") as air_total - FROM - "air" - GROUP BY - "lat", "lon" - """ + """ + SELECT + "lat", "lon", SUM("air") as air_total + FROM + "air" + GROUP BY + "lat", "lon" + """ ) result = query.to_pandas() @@ -43,7 +43,7 @@ def test_agg_regular(self): c.from_dataset('air', self.air) query = c.sql( - """ + """ SELECT "lat", "lon", AVG("air") as air_total FROM From 627e7476bd98ede8c992e00f7603defe79742389 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:20:51 -0700 Subject: [PATCH 07/10] Better name for private function. --- xarray_sql/df.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 301be9db..0c59d32f 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -81,12 +81,12 @@ def from_map_batched( if args is None: args = () - def apply_map(): + def map_batches(): for items in zip(*iterables): df = func(*items, *args, **kwargs) yield pa.RecordBatch.from_pandas(df, schema=schema) - return pa.RecordBatchReader.from_batches(schema, apply_map()) + return pa.RecordBatchReader.from_batches(schema, map_batches()) def from_map( From 5b79cbeebc5488acf1663a91dce0c42a3e3cf043 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:24:23 -0700 Subject: [PATCH 08/10] Applied pyink formatting. --- xarray_sql/core.py | 2 +- xarray_sql/df.py | 28 +++--- xarray_sql/df_integrationtest.py | 24 ++--- xarray_sql/df_test.py | 149 +++++++++++++++---------------- xarray_sql/sql_test.py | 4 +- 5 files changed, 105 insertions(+), 102 deletions(-) diff --git a/xarray_sql/core.py b/xarray_sql/core.py index f0a777a4..afbd2eca 100644 --- a/xarray_sql/core.py +++ b/xarray_sql/core.py @@ -39,7 +39,7 @@ def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: prod_vals = (ds.coords[k].values for k in dim_keys) coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( - -1, len(dim_keys) + -1, len(dim_keys) ) for i, d in enumerate(dim_keys): diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 0c59d32f..1ecc27ab 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -18,7 +18,7 @@ def _get_chunk_slicer( if dim in chunk_index: which_chunk = chunk_index[dim] return slice( - chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] + chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] ) return slice(None) @@ -40,11 +40,11 @@ def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]: ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) blocks = ( - { - dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) - for dim in ds.dims - } - for chunk_index in chunk_idxs + { + dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) + for dim in ds.dims + } + for chunk_index in chunk_idxs ) yield from blocks @@ -59,8 +59,11 @@ def _block_len(block: Block) -> int: def from_map_batched( - func: t.Callable[[...], pd.DataFrame], *iterables, args: t.Optional[t.Tuple] = None, - schema: pa.Schema = None, **kwargs + func: t.Callable[[...], pd.DataFrame], + *iterables, + args: t.Optional[t.Tuple] = None, + schema: pa.Schema = None, + **kwargs, ) -> pa.RecordBatchReader: """Create a PyArrow RecordBatchReader by mapping a function over iterables. @@ -129,7 +132,7 @@ def from_map( pa_table = pa.Table.from_pandas(df) except Exception as e: raise ValueError( - f"Cannot convert function result to PyArrow Table: {e}" + f"Cannot convert function result to PyArrow Table: {e}" ) results.append(pa_table) @@ -160,15 +163,16 @@ def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> ArrowStreamExportable: """ fst = next(iter(ds.values())).dims assert all( - da.dims == fst for da in ds.values() + da.dims == fst for da in ds.values() ), "All dimensions must be equal. Please filter data_vars in the Dataset." blocks = list(block_slices(ds, chunks)) - def pivot_block(b: Block): return pivot(ds.isel(b)) + def pivot_block(b: Block): + return pivot(ds.isel(b)) schema = pa.Schema.from_pandas(pivot_block(blocks[0])) last_schema = pa.Schema.from_pandas(pivot_block(blocks[-1])) - assert schema == last_schema, 'Schemas must be consistent across blocks!' + assert schema == last_schema, "Schemas must be consistent across blocks!" return from_map_batched(pivot_block, blocks, schema=schema) diff --git a/xarray_sql/df_integrationtest.py b/xarray_sql/df_integrationtest.py index fe438a86..6f0eae4d 100644 --- a/xarray_sql/df_integrationtest.py +++ b/xarray_sql/df_integrationtest.py @@ -9,21 +9,21 @@ class Era5TestCast(unittest.TestCase): def test_open_era5(self): era5_ds = xr.open_zarr( - 'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2', - chunks={'time': 240, 'level': 1}, + 'gs://gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2', + chunks={'time': 240, 'level': 1}, ) era5_wind_df = read_xarray( - era5_ds[['u_component_of_wind', 'v_component_of_wind']] + era5_ds[['u_component_of_wind', 'v_component_of_wind']] ) self.assertEqual( - list(era5_wind_df.columns), - [ - 'time', - 'level', - 'latitude', - 'longitude', - 'u_component_of_wind', - 'v_component_of_wind', - ], + list(era5_wind_df.columns), + [ + 'time', + 'level', + 'latitude', + 'longitude', + 'u_component_of_wind', + 'v_component_of_wind', + ], ) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index f825e90b..5193de84 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -20,18 +20,18 @@ def rand_wx(start: str, end: str) -> xr.Dataset: 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.'), + 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.'), ) @@ -48,23 +48,22 @@ def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100): 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, - }) + 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] - }) + result = pd.DataFrame({'x': [x], 'y': [y], 'sum': [x + y]}) return result @@ -76,7 +75,7 @@ def setUp(self) -> None: self.air = self.air.chunk(self.chunks) self.air_small = self.air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) ).chunk(self.chunks) self.randwx = rand_wx('1995-01-13T00', '1995-01-13T01') @@ -86,7 +85,7 @@ class ExplodeTest(DaskTestCase): def test_cardinality(self): dss = explode(self.air) self.assertEqual( - len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) + len(list(dss)), np.prod([len(c) for c in self.air.chunks.values()]) ) def test_dim_sizes__one(self): @@ -99,8 +98,8 @@ def skip_test_dim_sizes__all(self): # TODO(alxmrs): Why is this test slow? dss = explode(self.air) self.assertEqual( - [tuple(ds.dims.values()) for ds in dss], - list(itertools.product(*self.air.chunksizes.values())), + [tuple(ds.dims.values()) for ds in dss], + list(itertools.product(*self.air.chunksizes.values())), ) def test_data_equal__one__first(self): @@ -115,7 +114,7 @@ def test_data_equal__one__last(self): self.assertEqual(self.air.isel(iselection), ds) -@unittest.skip("Chose non-pa.Table implementation.") +@unittest.skip('Chose non-pa.Table implementation.') class PyArrowTableTest(DaskTestCase): def test_sanity(self): @@ -198,7 +197,7 @@ def multiply_and_add(x, multiplier, add_value): df = result.to_pandas() self.assertEqual( - list(df['result']), [12, 14, 16] + list(df['result']), [12, 14, 16] ) # (1*2+10, 2*2+10, 3*2+10) def test_from_map_with_pyarrow_tables(self): @@ -219,7 +218,8 @@ class TestFromMapBatchedCorrectness(DaskTestCase): def test_basic_functionality(self): """Test that from_map_batched produces correct RecordBatchReader.""" blocks = list( - block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4})) + block_slices(self.air_small, chunks={'time': 4, 'lat': 3, 'lon': 4}) + ) # Get expected schema first_block_df = pivot(self.air_small.isel(blocks[0])) @@ -227,9 +227,9 @@ def test_basic_functionality(self): # Create RecordBatchReader reader = from_map_batched( - pivot, - [self.air_small.isel(block) for block in blocks], - schema=expected_schema + pivot, + [self.air_small.isel(block) for block in blocks], + schema=expected_schema, ) # Verify it's a RecordBatchReader @@ -252,17 +252,12 @@ def test_multiple_iterables(self): 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()) - ]) + 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 + adding_function, x_values, y_values, schema=expected_schema ) # Read all data @@ -270,11 +265,13 @@ def test_multiple_iterables(self): df = table.to_pandas() # Verify results - expected_df = pd.DataFrame({ - 'x': x_values, - 'y': y_values, - 'sum': [x + y for x, y in zip(x_values, y_values)] - }) + 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) @@ -282,34 +279,27 @@ def test_with_args_and_kwargs(self): """Test from_map_batched with additional args and kwargs.""" def multiply_and_add(x, multiplier, offset=0): - result = pd.DataFrame({ - 'x': [x], - 'result': [x * multiplier + offset] - }) + result = pd.DataFrame({'x': [x], 'result': [x * multiplier + offset]}) return result values = [1, 2, 3] - expected_schema = pa.schema([ - ('x', pa.int64()), - ('result', pa.int64()) - ]) + expected_schema = pa.schema([('x', pa.int64()), ('result', pa.int64())]) reader = from_map_batched( - multiply_and_add, - values, - args=(2,), # multiplier = 2 - offset=5, # offset = 5 - schema=expected_schema + multiply_and_add, + values, + args=(2,), # multiplier = 2 + offset=5, # offset = 5 + schema=expected_schema, ) table = reader.read_all() df = table.to_pandas() # Verify results: x * 2 + 5 - expected_df = pd.DataFrame({ - 'x': [1, 2, 3], - 'result': [7, 9, 11] # (1*2+5, 2*2+5, 3*2+5) - }) + expected_df = pd.DataFrame( + {'x': [1, 2, 3], 'result': [7, 9, 11]} # (1*2+5, 2*2+5, 3*2+5) + ) pd.testing.assert_frame_equal(df, expected_df) @@ -318,9 +308,7 @@ def test_empty_iterables(self): empty_schema = pa.schema([('value', pa.int64())]) reader = from_map_batched( - lambda x: pd.DataFrame({'value': [x]}), - [], - schema=empty_schema + lambda x: pd.DataFrame({'value': [x]}), [], schema=empty_schema ) batches = list(reader) @@ -341,17 +329,25 @@ def test_consistency_with_regular_map(self): # Regular map approach regular_dfs = [pivot(ds) for ds in datasets] - regular_table = pa.Table.from_pandas(pd.concat(regular_dfs, ignore_index=True)) + regular_table = pa.Table.from_pandas( + pd.concat(regular_dfs, ignore_index=True) + ) # Results should be identical self.assertEqual(batched_table.schema, regular_table.schema) self.assertEqual(len(batched_table), len(regular_table)) # Compare data (allowing for potential column order differences) - 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) + 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) @@ -364,7 +360,9 @@ def test_integration_with_datafusion_via_read_xarray(self): air_chunked = air_small.chunk({'time': 25, 'lat': 5, 'lon': 8}) # read_xarray uses from_map_batched internally - arrow_stream = read_xarray(air_chunked, chunks={'time': 25, 'lat': 5, 'lon': 8}) + arrow_stream = read_xarray( + air_chunked, chunks={'time': 25, 'lat': 5, 'lon': 8} + ) # Verify it returns a proper ArrowStreamExportable (RecordBatchReader) self.assertTrue(hasattr(arrow_stream, 'schema')) @@ -381,6 +379,7 @@ def test_integration_with_datafusion_via_read_xarray(self): class ReadXarrayStreamingTest(unittest.TestCase): + def setUp(self): self.large_ds = create_large_dataset().chunk({'time': 25}) diff --git a/xarray_sql/sql_test.py b/xarray_sql/sql_test.py index f9a68e6f..041f1320 100644 --- a/xarray_sql/sql_test.py +++ b/xarray_sql/sql_test.py @@ -22,7 +22,7 @@ def test_agg_small(self): c.from_dataset('air', self.air_small) query = c.sql( - """ + """ SELECT "lat", "lon", SUM("air") as air_total FROM @@ -43,7 +43,7 @@ def test_agg_regular(self): c.from_dataset('air', self.air) query = c.sql( - """ + """ SELECT "lat", "lon", AVG("air") as air_total FROM From 425e3176014b8134a0c954bb67e6184e6f8aa617 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:28:52 -0700 Subject: [PATCH 09/10] More accurate comment. --- xarray_sql/df.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 1ecc27ab..2e941b88 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -72,7 +72,8 @@ def from_map_batched( RecordBatches which are created via the `func` one-at-a-time. Args: - func: Function to apply to each element of the iterables. + func: Function to apply to each element of the iterables. Currently, the function + must return a Pandas DataFrame. *iterables: Iterable objects to map the function over. schema: Optional schema needed for the RecordBatchReader. args: Additional positional arguments to pass to func. From bd88b49fd75acdc02cc76fb2abfd3fd35336cf21 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 13 Jul 2025 13:29:53 -0700 Subject: [PATCH 10/10] Delete stale tests. --- xarray_sql/df_test.py | 44 ------------------------------------------- 1 file changed, 44 deletions(-) diff --git a/xarray_sql/df_test.py b/xarray_sql/df_test.py index 5193de84..403f8155 100644 --- a/xarray_sql/df_test.py +++ b/xarray_sql/df_test.py @@ -114,50 +114,6 @@ def test_data_equal__one__last(self): self.assertEqual(self.air.isel(iselection), ds) -@unittest.skip('Chose non-pa.Table implementation.') -class PyArrowTableTest(DaskTestCase): - - def test_sanity(self): - table = read_xarray(self.air_small) - self.assertIsNotNone(table) - self.assertIsInstance(table, pa.Table) - self.assertEqual(len(table), np.prod(list(self.air_small.sizes.values()))) - - def test_columns(self): - table = read_xarray(self.air_small) - cols = table.column_names - self.assertEqual(cols, ['lat', 'time', 'lon', 'air']) - - def test_dtypes(self): - table = read_xarray(self.air_small) - # Convert to pandas to check dtypes - df = table.to_pandas() - types = list(df.dtypes) - self.assertEqual([self.air_small[c].dtype for c in df.columns], types) - - def test_different_chunk_sizes(self): - default_table = read_xarray(self.air_small) - chunked_table = read_xarray(self.air_small, dict(time=5)) - - # Both should produce valid tables - self.assertIsInstance(default_table, pa.Table) - self.assertIsInstance(chunked_table, pa.Table) - # Should have same number of rows - self.assertEqual(len(default_table), len(chunked_table)) - - def test_chunk_perf(self): - table = read_xarray(self.air, chunks=dict(time=6)) - self.assertIsNotNone(table) - self.assertEqual(len(table), np.prod(list(self.air.sizes.values()))) - - def test_column_metadata_preserved(self): - try: - table = read_xarray(self.randwx, chunks=dict(time=24)) - self.assertIsInstance(table, pa.Table) - except Exception as e: - self.fail(f'Unexpected error: {e}') - - class FromMapTest(unittest.TestCase): def test_basic_from_map(self):