add catalog provider for review

Author: yc-606

catalog.py

@@ -0,0 +1,105 @@
+import xarray as xr
+import datafusion as dfn
+from xarray_sql.reader import read_xarray_table
+
+
+def group_vars_by_dims(ds):
+    """
+    Group variables in the dataset based on shared dims
+
+    ("time", "lat", "lon"):          ["temperature_2m", "wind_speed"],
+    ("time", "lat", "lon", "level"): ["pressure", "humidity"]
+    """
+    groups = {}
+
+    for var_name, var in ds.data_vars.items():
+        dims = var.dims
+
+        if dims not in groups:
+            groups[dims] = []
+
+        groups[dims].append(var_name)
+
+    return groups
+
+
+def dims_to_table_name(dims):
+    """
+    "time", "lat", "lon" -> "time_lat_lon"
+    """
+    return "_".join(dims)
+
+    
+class XarraySchemaProvider(dfn.catalog.SchemaProvider):
+    """
+    Custom datafusion schema that holds the tables
+    """
+
+    def __init__(self, ds, groups, chunks):
+        # dictionary to store the tables 
+        self.tables = {}
+
+        # create a table for for each group of vars
+        for dims, var_names in groups.items():
+            table_name = dims_to_table_name(dims)
+            subset = ds[var_names]
+            self.tables[table_name] = read_xarray_table(subset, chunks)
+
+    def table_names(self):
+        return set(self.tables.keys())
+
+    def table(self, name):
+        return self.tables.get(name)
+
+    def table_exist(self, name):
+        return name in self.tables
+
+    def register_table(self, name, table):
+        self.tables[name] = table
+
+    def deregister_table(self, name, cascade=True):
+        del self.tables[name]
+
+
+class XarrayCatalogProvider(dfn.catalog.CatalogProvider):
+    """
+    Custom datafusion catalog that holds the schemas
+    """
+    #Constructor
+    def __init__(self, ds, schema_name, chunks):
+        groups = group_vars_by_dims(ds)
+
+        # dictionary to store schemas using previous schema class
+        """
+        "data": {
+        "time_lat_lon":       [temperature_2m, wind_speed],
+        "time_lat_lon_level": [pressure, humidity]
+        }
+        """
+        self.schemas = {
+            schema_name: XarraySchemaProvider(ds, groups, chunks)
+        }
+
+    """
+    Other methods from test_catalog.py
+    """
+    def schema_names(self):
+        return set(self.schemas.keys())
+
+    def schema(self, name):
+        return self.schemas.get(name)
+
+    def register_schema(self, name, schema):
+        self.schemas[name] = schema
+
+    def deregister_schema(self, name, cascade=True):
+        del self.schemas[name]
+
+
+def register_catalog_from_dataset(ctx, ds, catalog_name="xarray", schema_name="data", chunks=None):
+    """
+    Main function. Takes an xarray dataset and registers it
+    with DataFusion so you can query it with SQL.
+    """
+    catalog = XarrayCatalogProvider(ds, schema_name, chunks)
+    ctx.register_catalog_provider(catalog_name, catalog)

catalog_test.py

@@ -0,0 +1,63 @@
+import numpy as np
+import xarray as xr
+from context import XarrayContext
+
+
+# create a fake era5 dataset for testing
+times = np.array(["2020-01-01", "2020-01-02", "2020-01-03"], dtype="datetime64")
+lats = np.array([0.0, 1.0, 2.0])
+lons = np.array([0.0, 1.0, 2.0])
+levels = np.array([500, 850])
+
+ds = xr.Dataset(
+    {
+        "temperature_2m": (["time", "lat", "lon"], np.random.rand(3, 3, 3)),
+        "wind_speed":     (["time", "lat", "lon"], np.random.rand(3, 3, 3)),
+        "pressure":       (["time", "lat", "lon", "level"], np.random.rand(3, 3, 3, 2)),
+        "humidity":       (["time", "lat", "lon", "level"], np.random.rand(3, 3, 3, 2)),
+    },
+    coords={
+        "time":  times,
+        "lat":   lats,
+        "lon":   lons,
+        "level": levels,
+    }
+).chunk({"time": 1})
+
+print("Variables:", list(ds.data_vars))
+print("Dimensions:", list(ds.dims))
+
+ctx = XarrayContext()
+ctx.register_catalog_from_dataset(ds)
+
+print("\nCatalogs:", ctx.catalog_names())
+print("Schemas:", ctx.catalog("xarray").schema_names())
+print("Tables:", ctx.catalog("xarray").schema("data").table_names())
+
+print("\n--- Surface variables (time, lat, lon) ---")
+result = ctx.sql("SELECT * FROM xarray.data.time_lat_lon LIMIT 5").collect()
+for batch in result:
+    print(batch.to_pandas())
+
+
+print("\n--- Atmospheric variables (time, lat, lon, level) ---")
+result = ctx.sql("SELECT * FROM xarray.data.time_lat_lon_level LIMIT 5").collect()
+for batch in result:
+    print(batch.to_pandas())
+
+print("\n--- Joined surface + atmospheric on shared dims ---")
+result = ctx.sql("""
+    SELECT
+        s.time, s.lat, s.lon,
+        s.temperature_2m,
+        a.level,
+        a.pressure
+    FROM xarray.data.time_lat_lon s
+    JOIN xarray.data.time_lat_lon_level a
+        ON s.time = a.time
+        AND s.lat = a.lat
+        AND s.lon = a.lon
+    LIMIT 10
+""").collect()
+for batch in result:
+    print(batch.to_pandas())

xarray_context.py

@@ -0,0 +1,13 @@
+import xarray as xr
+from datafusion import SessionContext
+from catalog import register_catalog_from_dataset
+
+
+class XarrayContext(SessionContext):
+    """
+    A regular DataFusion SessionContext but with an extra method
+    for registering xarray datasets.
+    """
+
+    def register_catalog_from_dataset(self, ds, catalog_name="xarray", schema_name="data", chunks=None):
+        register_catalog_from_dataset(self, ds, catalog_name, schema_name, chunks)