terrakio-core 0.3.9__py3-none-any.whl → 0.4.0__py3-none-any.whl

terrakio_core/__init__.py

@@ -5,7 +5,7 @@ Terrakio Core
 Core components for Terrakio API clients.
 """
 
-__version__ = "0.3.9"
+__version__ = "0.4.0"
 
 from .async_client import AsyncClient
 from .sync_client import SyncClient as Client

terrakio_core/async_client.py

@@ -182,6 +182,11 @@ class AsyncClient(BaseClient):
             out_crs: str = "epsg:4326",
             resolution: int = -1,
             geom_fix: bool = False,
+            drop_nan: bool = False,
+            spatial_reduction: str = None,
+            temporal_reduction: str = None,
+            max_memory_mb: int = 500,
+            stream_to_disk: bool = False,
     ):
         """
         Compute zonal statistics for all geometries in a GeoDataFrame.
@@ -195,11 +200,20 @@ class AsyncClient(BaseClient):
             out_crs (str): Output coordinate reference system
             resolution (int): Resolution parameter
             geom_fix (bool): Whether to fix the geometry (default False)
+            drop_nan (bool): Whether to drop NaN values from the results (default False)
+            spatial_reduction (str): Reduction operation for spatial dimensions (x, y). 
+                                Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+            temporal_reduction (str): Reduction operation for temporal dimension (time).
+                                    Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+            max_memory_mb (int): Maximum memory threshold in MB (default 500MB)
+            stream_to_disk (bool): Whether to stream datasets to disk as NetCDF files (default False)
+
         Returns:
             geopandas.GeoDataFrame: GeoDataFrame with added columns for results, or None if inplace=True
+                                If stream_to_disk=True, large datasets are saved as NetCDF files with file paths stored.
 
         Raises:
-            ValueError: If concurrency is too high
+            ValueError: If concurrency is too high or if data exceeds memory limit without streaming
             APIError: If the API request fails
         """
         return await _zonal_stats(
@@ -211,7 +225,12 @@ class AsyncClient(BaseClient):
             in_crs=in_crs,
             out_crs=out_crs,
             resolution=resolution,
-            geom_fix=geom_fix
+            geom_fix=geom_fix,
+            drop_nan=drop_nan,
+            spatial_reduction=spatial_reduction,
+            temporal_reduction=temporal_reduction,
+            max_memory_mb=max_memory_mb,
+            stream_to_disk=stream_to_disk
         )
 
     async def create_dataset_file(

terrakio_core/client.py

@@ -1,11 +1,14 @@
 from typing import Optional
 import logging
+import warnings
 from terrakio_core.config import read_config_file, DEFAULT_CONFIG_FILE
 from abc import abstractmethod
+import xarray as xr
+
 
 class BaseClient():
     def __init__(self, url: Optional[str] = None, api_key: Optional[str] = None, verbose: bool = False):
-
+        
         self.verbose = verbose
         self.logger = logging.getLogger("terrakio")
         if verbose:
@@ -21,17 +24,110 @@ class BaseClient():
         self.url = url
         self.key = api_key
 
-        config = read_config_file( DEFAULT_CONFIG_FILE, logger = self.logger)
-        
+        config = read_config_file(DEFAULT_CONFIG_FILE, logger=self.logger)
+                
         if self.url is None:
             self.url = config.get('url')
-        
+                
         if self.key is None:
             self.key = config.get('key')
 
         self.token = config.get('token')
-
         
+        # Apply xarray printing fix to prevent crashes with GeoDataFrames
+        self._apply_xarray_fix()
+    
+    def _apply_xarray_fix(self):
+        """
+        Apply xarray printing fix to prevent crashes when GeoDataFrames contain xarray objects.
+        This fix is applied automatically when the client is initialized.
+        """
+        try:
+            
+            # Check if fix is already applied globally
+            if hasattr(xr.DataArray, '_terrakio_fix_applied'):
+                if self.verbose:
+                    self.logger.info("xarray printing fix already applied")
+                return
+            
+            # Store original methods for potential restoration
+            if not hasattr(xr.DataArray, '_original_iter'):
+                xr.DataArray._original_iter = xr.DataArray.__iter__
+                xr.Dataset._original_iter = xr.Dataset.__iter__
+            
+            # Define safe iteration methods that prevent pandas from iterating
+            # but leave __repr__ and __str__ untouched for normal xarray printing
+            def safe_dataarray_iter(self):
+                # Return infinite iterator that always yields the same safe value
+                name = getattr(self, 'name', None) or 'unnamed'
+                shape_str = 'x'.join(map(str, self.shape)) if hasattr(self, 'shape') else 'unknown'
+                placeholder = f"<DataArray '{name}' {shape_str}>"
+                while True:
+                    yield placeholder
+            
+            def safe_dataset_iter(self):
+                # Return infinite iterator that always yields the same safe value
+                num_vars = len(self.data_vars) if hasattr(self, 'data_vars') else 0
+                num_dims = len(self.dims) if hasattr(self, 'dims') else 0
+                placeholder = f"<Dataset: {num_vars} vars, {num_dims} dims>"
+                while True:
+                    yield placeholder
+            
+            # Apply only the iteration fix - leave __repr__ and __str__ untouched
+            xr.DataArray.__iter__ = safe_dataarray_iter
+            xr.Dataset.__iter__ = safe_dataset_iter
+            
+            # Mark as applied to avoid duplicate applications
+            xr.DataArray._terrakio_fix_applied = True
+            xr.Dataset._terrakio_fix_applied = True
+            
+            if self.verbose:
+                self.logger.info("xarray iteration fix applied - GeoDataFrames with xarray objects will print safely, direct xarray printing unchanged")
+                
+        except ImportError:
+            # xarray not installed, skip the fix
+            if self.verbose:
+                self.logger.info("xarray not installed, skipping printing fix")
+        except Exception as e:
+            # Log warning but don't fail initialization
+            warning_msg = f"Failed to apply xarray printing fix: {e}"
+            warnings.warn(warning_msg)
+            if self.verbose:
+                self.logger.warning(warning_msg)
+    
+    def restore_xarray_printing(self):
+        """
+        Restore original xarray printing behavior.
+        Call this method if you want to see full xarray representations again.
+        """
+        try:
+            import xarray as xr
+            
+            if hasattr(xr.DataArray, '_original_iter'):
+                xr.DataArray.__iter__ = xr.DataArray._original_iter
+                xr.Dataset.__iter__ = xr.Dataset._original_iter
+                
+                # Remove the fix markers
+                if hasattr(xr.DataArray, '_terrakio_fix_applied'):
+                    delattr(xr.DataArray, '_terrakio_fix_applied')
+                if hasattr(xr.Dataset, '_terrakio_fix_applied'):
+                    delattr(xr.Dataset, '_terrakio_fix_applied')
+                
+                if self.verbose:
+                    self.logger.info("Original xarray iteration behavior restored")
+            else:
+                if self.verbose:
+                    self.logger.info("No xarray fix to restore")
+                    
+        except ImportError:
+            if self.verbose:
+                self.logger.info("xarray not available")
+        except Exception as e:
+            warning_msg = f"Failed to restore xarray printing: {e}"
+            warnings.warn(warning_msg)
+            if self.verbose:
+                self.logger.warning(warning_msg)
+    
     @abstractmethod
     def _setup_session(self):
         """Initialize the HTTP session - implemented by sync/async clients"""

terrakio_core/convenience_functions/convenience_functions.py

@@ -11,37 +11,39 @@ from ..helper.tiles import tiles
 import uuid
 import xarray as xr
 
-
-async def zonal_stats(
+async def request_data(
         client,
         gdf: GeoDataFrame,
         expr: str,
         conc: int = 20,
-        inplace: bool = False,
         in_crs: str = "epsg:4326",
         out_crs: str = "epsg:4326",
         resolution: int = -1,
         geom_fix: bool = False,
+        max_memory_mb: int = 500,
+        stream_to_disk: bool = None,
 ):
     """
-    Compute zonal statistics for all geometries in a GeoDataFrame.
+    Request xarray datasets for all geometries in a GeoDataFrame.
 
     Args:
         client: The AsyncClient instance
         gdf (GeoDataFrame): GeoDataFrame containing geometries
         expr (str): Terrakio expression to evaluate, can include spatial aggregations
         conc (int): Number of concurrent requests to make
-        inplace (bool): Whether to modify the input GeoDataFrame in place
         in_crs (str): Input coordinate reference system
         out_crs (str): Output coordinate reference system
         resolution (int): Resolution parameter
         geom_fix (bool): Whether to fix the geometry (default False)
+        max_memory_mb (int): Maximum memory threshold in MB (default 500MB)
+        stream_to_disk (bool): Whether to stream large datasets to disk. If None, will be determined automatically.
+
     Returns:
-        geopandas.GeoDataFrame: GeoDataFrame with variable dataarrays in separate columns.
-                               Each row represents one geometry with full time-dimensional dataarrays.
+        geopandas.GeoDataFrame: Copy of input GeoDataFrame with additional 'dataset' column
+                               containing the xarray Dataset for each geometry.
 
     Raises:
-        ValueError: If concurrency is too high
+        ValueError: If concurrency is too high or if data exceeds memory limit without streaming
         APIError: If the API request fails
     """
     if conc > 100:
@@ -49,13 +51,54 @@ async def zonal_stats(
     
     total_geometries = len(gdf)
     
+    # First, make a request with the first geometry to estimate total memory usage
+    client.logger.info("Estimating total memory usage...")
+    first_geom = gdf.geometry.iloc[0]
+    feature = {
+        "type": "Feature",
+        "geometry": mapping(first_geom),
+        "properties": {}
+    }
+    
+    try:
+        first_result = await client.geoquery(expr=expr, feature=feature,
+                                           in_crs=in_crs, out_crs=out_crs, resolution=resolution, geom_fix=geom_fix)
+        if isinstance(first_result, dict) and first_result.get("error"):
+            error_msg = f"Request failed: {first_result.get('error_message', 'Unknown error')}"
+            if first_result.get('status_code'):
+                error_msg = f"Request failed with status {first_result['status_code']}: {first_result.get('error_message', 'Unknown error')}"
+            raise APIError(error_msg)
+        
+        if not isinstance(first_result, xr.Dataset):
+            raise ValueError(f"Expected xarray Dataset, got {type(first_result)}")
+        
+        # Estimate total memory usage
+        single_dataset_size_bytes = estimate_dataset_size(first_result)
+        total_size_bytes = single_dataset_size_bytes * total_geometries
+        total_size_mb = total_size_bytes / (1024 * 1024)
+        
+        client.logger.info(f"Estimated total memory usage: {total_size_mb:.2f} MB for {total_geometries} geometries")
+        
+        # Check if we need to stream to disk
+        if stream_to_disk is None:
+            # Auto-determine based on memory usage
+            if total_size_mb > max_memory_mb:
+                client.logger.warning(f"The data you are requesting exceeds {max_memory_mb} MB, we recommend you to set the stream_to_disk parameter to True")
+                raise ValueError(f"The data you are requesting exceeds {max_memory_mb} MB, we recommend you to set the stream_to_disk parameter to True")
+        
+    except Exception as e:
+        if "recommend you to set the stream_to_disk parameter to True" in str(e):
+            raise
+        client.logger.error(f"Failed to estimate memory usage: {e}")
+        raise
+    
     client.logger.info(f"Processing {total_geometries} geometries with concurrency {conc}")
     
     completed_count = 0
     lock = asyncio.Lock()
     
     async def process_geometry(geom):
-        """Process a single geometry"""
+        """Process a single geometry and return the dataset"""
         nonlocal completed_count
         
         try:
@@ -103,53 +146,261 @@ async def zonal_stats(
                 raise APIError(f"API request failed: {e.response.text}")
         raise
     
-    client.logger.info("All requests completed! Processing results...")
+    client.logger.info("All requests completed!")
 
-    
     if not all_results:
         raise ValueError("No valid results were returned for any geometry")
     
+    # Create a copy of the input GeoDataFrame
+    result_gdf = gdf.copy()
+    
+    # Add the dataset column with the xarray datasets
+    result_gdf['dataset'] = all_results
+    
+    return result_gdf
+
+
+import os
+from pathlib import Path
+
+def estimate_dataset_size(dataset):
+    """
+    Estimate the memory size of an xarray dataset in bytes.
+    
+    Args:
+        dataset: xarray Dataset
+        
+    Returns:
+        int: Estimated size in bytes
+    """
+    total_size = 0
+    for var_name, var in dataset.data_vars.items():
+        # Get the dtype size in bytes
+        dtype_size = var.dtype.itemsize
+        # Get the total number of elements
+        total_elements = var.size
+        # Calculate total size for this variable
+        total_size += dtype_size * total_elements
+    
+    # Add coordinate sizes
+    for coord_name, coord in dataset.coords.items():
+        if coord_name not in dataset.dims:  # Don't double count dimension coordinates
+            dtype_size = coord.dtype.itemsize
+            total_elements = coord.size
+            total_size += dtype_size * total_elements
+    
+    return total_size
 
+def save_dataset_to_file(dataset, filepath):
+    """
+    Save dataset to NetCDF file.
+    
+    Args:
+        dataset: xarray Dataset
+        filepath: Path to save the file
+        
+    Returns:
+        str: Path to saved file
+    """
+    filepath = Path(filepath)
+    
+    if not str(filepath).endswith('.nc'):
+        filepath = filepath.with_suffix('.nc')
+    
+    dataset.to_netcdf(filepath)
+    return str(filepath)
+
+def post_processing(
+        gdf_with_datasets: GeoDataFrame,
+        spatial_reduction: str = None,
+        temporal_reduction: str = None,
+        drop_nan: bool = False,
+        inplace: bool = False,
+        stream_to_disk: bool = False,
+):
+    """
+    Post-process the GeoDataFrame with datasets to extract variables with optional reductions.
+
+    Args:
+        gdf_with_datasets (GeoDataFrame): GeoDataFrame with 'dataset' column containing xarray Datasets
+        spatial_reduction (str): Reduction operation for spatial dimensions (x, y). 
+                               Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+        temporal_reduction (str): Reduction operation for temporal dimension (time).
+                                Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+        drop_nan (bool): Whether to drop NaN values from the results (default False)
+        inplace (bool): Whether to modify the input GeoDataFrame in place
+        stream_to_disk (bool): Whether to stream datasets to disk as NetCDF files (default False)
+
+    Returns:
+        geopandas.GeoDataFrame: GeoDataFrame with variable dataarrays/values or file paths in separate columns.
+                               If stream_to_disk=True, large datasets are saved as NetCDF files with file paths stored.
+    """
+    if 'dataset' not in gdf_with_datasets.columns:
+        raise ValueError("Input GeoDataFrame must contain a 'dataset' column")
+    
+    # Validate reduction parameters
+    valid_reductions = ['mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count']
+    if spatial_reduction and spatial_reduction not in valid_reductions:
+        raise ValueError(f"spatial_reduction must be one of {valid_reductions}")
+    if temporal_reduction and temporal_reduction not in valid_reductions:
+        raise ValueError(f"temporal_reduction must be one of {valid_reductions}")
+    
     result_rows = []
     geometries = []
     
-    # Results should be in the same order as input geometries
-    for i, dataset in enumerate(all_results):
-        # Get the original geometry by index
-        original_geom = gdf.geometry.iloc[i]
-        # Create single row for this geometry
+    # Process each row (geometry + dataset)
+    for i, row in gdf_with_datasets.iterrows():
+        dataset = row['dataset']
+        
+        # Create new row for this geometry
         new_row = {}
-                 
-        # Copy original GeoDataFrame attributes
-        for col in gdf.columns:
-            if col != 'geometry':
-                new_row[col] = gdf.iloc[i][col]
-        # Store the full dataarray for each variable (with time dimension intact)
+        
+        # Copy original GeoDataFrame attributes (excluding dataset column)
+        for col in gdf_with_datasets.columns:
+            if col not in ['geometry', 'dataset']:
+                new_row[col] = row[col]
+        
+        # Process each variable in the dataset
         data_vars = list(dataset.data_vars.keys())
         for var_name in data_vars:
             var_data = dataset[var_name]
-            new_row[f'{var_name}_dataarray'] = var_data
+            
+            # Apply drop_nan if requested
+            if drop_nan:
+                # Drop spatial dimensions where all values are NaN
+                var_data = var_data.dropna(dim='x', how='all').dropna(dim='y', how='all')
+                
+                # Drop time dimensions where all values are NaN
+                if 'time' in var_data.dims:
+                    var_data = var_data.dropna(dim='time', how='all')
+            
+            # Check current dimensions to determine if aggregation is needed
+            current_dims = set(var_data.dims)
+            has_spatial_dims = bool(current_dims.intersection(['x', 'y']))
+            has_temporal_dim = 'time' in current_dims
+            
+            # Apply spatial reduction only if spatial dimensions exist and reduction is requested
+            if spatial_reduction and has_spatial_dims:
+                spatial_dims = [dim for dim in ['x', 'y'] if dim in var_data.dims]
+                if spatial_dims:
+                    if spatial_reduction == 'count':
+                        var_data = var_data.count(dim=spatial_dims)
+                    else:
+                        var_data = getattr(var_data, spatial_reduction)(dim=spatial_dims)
+            
+            # Apply temporal reduction only if time dimension exists and reduction is requested
+            if temporal_reduction and has_temporal_dim:
+                if temporal_reduction == 'count':
+                    var_data = var_data.count(dim='time')
+                else:
+                    var_data = getattr(var_data, temporal_reduction)(dim='time')
+            
+            # Handle streaming to disk if requested
+            if stream_to_disk:
+                # Create a single-variable dataset for saving
+                single_var_dataset = var_data.to_dataset(name=var_name)
+                
+                # Generate filename based on row index and variable name
+                filename = f"geometry_{i}_{var_name}.nc"
+                filepath = os.path.join(os.getcwd(), filename)
+                
+                # Save to disk and store file path
+                saved_path = save_dataset_to_file(single_var_dataset, filepath)
+                new_row[var_name] = f"file://{saved_path}"
+                
+                print(f"Dataset for geometry {i}, variable '{var_name}' saved to: {saved_path}")
+            else:
+                # Keep in memory
+                new_row[var_name] = var_data
+        
         result_rows.append(new_row)
-        geometries.append(original_geom)
+        geometries.append(row['geometry'])
+    
     # Create the result GeoDataFrame with default integer index
     result_gdf = GeoDataFrame(result_rows, geometry=geometries)
+    
     if inplace:
         # Clear original gdf and replace with result_gdf content
-        gdf.drop(gdf.index, inplace=True)
-        gdf.drop(gdf.columns, axis=1, inplace=True)
+        gdf_with_datasets.drop(gdf_with_datasets.index, inplace=True)
+        gdf_with_datasets.drop(gdf_with_datasets.columns, axis=1, inplace=True)
         
         # Copy all data from result_gdf to gdf
         for col in result_gdf.columns:
-            gdf[col] = result_gdf[col].values
+            gdf_with_datasets[col] = result_gdf[col].values
         
         # Ensure it remains a GeoDataFrame with correct geometry
-        gdf.geometry = result_gdf.geometry
+        gdf_with_datasets.geometry = result_gdf.geometry
         
         return None
     else:
-
         return result_gdf
 
+
+# Updated zonal_stats function that uses both parts
+async def zonal_stats(
+        client,
+        gdf: GeoDataFrame,
+        expr: str,
+        conc: int = 20,
+        inplace: bool = False,
+        in_crs: str = "epsg:4326",
+        out_crs: str = "epsg:4326",
+        resolution: int = -1,
+        geom_fix: bool = False,
+        drop_nan: bool = False,
+        spatial_reduction: str = None,
+        temporal_reduction: str = None,
+        max_memory_mb: int = 500,
+        stream_to_disk: bool = False,
+):
+    """
+    Compute zonal statistics for all geometries in a GeoDataFrame.
+    This is a convenience function that combines request_data and post_processing.
+    
+    Args:
+        client: The AsyncClient instance
+        gdf (GeoDataFrame): GeoDataFrame containing geometries
+        expr (str): Terrakio expression to evaluate, can include spatial aggregations
+        conc (int): Number of concurrent requests to make
+        inplace (bool): Whether to modify the input GeoDataFrame in place
+        in_crs (str): Input coordinate reference system
+        out_crs (str): Output coordinate reference system
+        resolution (int): Resolution parameter
+        geom_fix (bool): Whether to fix the geometry (default False)
+        drop_nan (bool): Whether to drop NaN values from the results (default False)
+        spatial_reduction (str): Reduction operation for spatial dimensions (x, y). 
+                               Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+        temporal_reduction (str): Reduction operation for temporal dimension (time).
+                                Options: 'mean', 'median', 'min', 'max', 'std', 'var', 'sum', 'count'
+        max_memory_mb (int): Maximum memory threshold in MB (default 500MB)
+        stream_to_disk (bool): Whether to stream datasets to disk as NetCDF files (default False)
+    """
+    # Step 1: Request data (with memory estimation)
+    gdf_with_datasets = await request_data(
+        client=client,
+        gdf=gdf,
+        expr=expr,
+        conc=conc,
+        in_crs=in_crs,
+        out_crs=out_crs,
+        resolution=resolution,
+        geom_fix=geom_fix,
+        max_memory_mb=max_memory_mb,
+        stream_to_disk=stream_to_disk
+    )
+    
+    # Step 2: Post-process with reductions and optional streaming
+    result = post_processing(
+        gdf_with_datasets=gdf_with_datasets,
+        spatial_reduction=spatial_reduction,
+        temporal_reduction=temporal_reduction,
+        drop_nan=drop_nan,
+        inplace=inplace,
+        stream_to_disk=stream_to_disk
+    )
+    
+    return result
+
 async def create_dataset_file(
     client,
     aoi: str,