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

terrakio_core/convenience_functions/convenience_functions.py

@@ -9,36 +9,41 @@ from ..exceptions import APIError, ConfigurationError
 from ..helper.bounded_taskgroup import BoundedTaskGroup
 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 added columns for results, or None if inplace=True
+        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:
@@ -46,32 +51,74 @@ 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, index):
-        """Process a single geometry"""
+    async def process_geometry(geom):
+        """Process a single geometry and return the dataset"""
         nonlocal completed_count
         
         try:
             feature = {
                 "type": "Feature",
                 "geometry": mapping(geom),
-                "properties": {"index": index}
+                "properties": {}
             }
-            result = await client.geoquery(expr=expr, feature=feature, output="csv",
+            # Request xarray dataset
+            result = await client.geoquery(expr=expr, feature=feature,
                                         in_crs=in_crs, out_crs=out_crs, resolution=resolution, geom_fix=geom_fix)
-            
             if isinstance(result, dict) and result.get("error"):
-                error_msg = f"Request {index} failed: {result.get('error_message', 'Unknown error')}"
+                error_msg = f"Request failed: {result.get('error_message', 'Unknown error')}"
                 if result.get('status_code'):
-                    error_msg = f"Request {index} failed with status {result['status_code']}: {result.get('error_message', 'Unknown error')}"
+                    error_msg = f"Request failed with status {result['status_code']}: {result.get('error_message', 'Unknown error')}"
                 raise APIError(error_msg)
             
-            if isinstance(result, pd.DataFrame):
-                result['_geometry_index'] = index
+            # Ensure we got an xarray Dataset
+            if not isinstance(result, xr.Dataset):
+                raise ValueError(f"Expected xarray Dataset, got {type(result)}")
             
             async with lock:
                 completed_count += 1
@@ -88,85 +135,271 @@ async def zonal_stats(
     try:
         async with BoundedTaskGroup(max_concurrency=conc) as tg:
             tasks = [
-                tg.create_task(process_geometry(gdf.geometry.iloc[idx], idx))
+                tg.create_task(process_geometry(gdf.geometry.iloc[idx]))
                 for idx in range(len(gdf))
             ]
         all_results = [task.result() for task in tasks]
-        
+
     except* Exception as eg:
         for e in eg.exceptions:
             if hasattr(e, 'response'):
                 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
         
-    combined_df = pd.concat(all_results, ignore_index=True)
+    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
     
-    has_time = 'time' in combined_df.columns
+    # 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
     
-    if has_time:
-        if '_geometry_index' not in combined_df.columns:
-            raise ValueError("Missing geometry index in results")
+    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
         
-        combined_df.set_index(['_geometry_index', 'time'], inplace=True)
+    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 = []
+    
+    # Process each row (geometry + dataset)
+    for i, row in gdf_with_datasets.iterrows():
+        dataset = row['dataset']
         
-        result_cols = combined_df.columns
+        # Create new row for this geometry
+        new_row = {}
         
-        result_rows = []
-        geometries = []
+        # 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]
         
-        for (geom_idx, time_val), row in combined_df.iterrows():
-            new_row = {}
+        # Process each variable in the dataset
+        data_vars = list(dataset.data_vars.keys())
+        for var_name in data_vars:
+            var_data = dataset[var_name]
             
-            for col in gdf.columns:
-                if col != 'geometry':
-                    new_row[col] = gdf.loc[geom_idx, col]
+            # 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')
             
-            for col in result_cols:
-                new_row[col] = row[col]
+            # 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
             
-            result_rows.append(new_row)
-            geometries.append(gdf.geometry.iloc[geom_idx])
-        
-        multi_index = pd.MultiIndex.from_tuples(
-            combined_df.index.tolist(),
-            names=['geometry_index', 'time']
-        )
-        
-        result_gdf = GeoDataFrame(
-            result_rows, 
-            geometry=geometries,
-            index=multi_index
-        )
+            # 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
         
-        if inplace:
-            return result_gdf
-        else:
-            return result_gdf
-    else:
-        result_gdf = gdf.copy() if not inplace else gdf
+        result_rows.append(new_row)
+        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_with_datasets.drop(gdf_with_datasets.index, inplace=True)
+        gdf_with_datasets.drop(gdf_with_datasets.columns, axis=1, inplace=True)
         
-        result_cols = [col for col in combined_df.columns if col not in ['_geometry_index']]
+        # Copy all data from result_gdf to gdf
+        for col in result_gdf.columns:
+            gdf_with_datasets[col] = result_gdf[col].values
         
-        geom_idx_to_row = {}
-        for idx, row in combined_df.iterrows():
-            geom_idx = int(row['_geometry_index'])
-            geom_idx_to_row[geom_idx] = row
+        # Ensure it remains a GeoDataFrame with correct geometry
+        gdf_with_datasets.geometry = result_gdf.geometry
         
-        for col in result_cols:
-            if col not in result_gdf.columns:
-                result_gdf[col] = None
-            
-            for geom_idx, row in geom_idx_to_row.items():
-                result_gdf.loc[geom_idx, col] = row[col]
-        if inplace:
-            return None
-        else:
-            return result_gdf
+        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,

terrakio_core/endpoints/mass_stats.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from urllib.parse import urlparse
 from ..helper.decorators import require_token, require_api_key, require_auth
 import aiohttp
-
+from typing import Dict, Any, Optional, List, Union
 class MassStats:
     def __init__(self, client):
         self._client = client
@@ -218,7 +218,7 @@ class MassStats:
             params["output"] = output
             
         return self._client._terrakio_request("GET", "mass_stats/download", params=params)
-    
+
     @require_api_key
     async def _upload_file(self, file_path: str, url: str, use_gzip: bool = False):
         """
@@ -237,6 +237,18 @@ class MassStats:
         except json.JSONDecodeError as e:
             raise ValueError(f"Invalid JSON in file {file_path}: {e}")
         
+        return await self._upload_json_data(json_data, url, use_gzip)
+
+    @require_api_key
+    async def _upload_json_data(self, json_data: Union[Dict, List], url: str, use_gzip: bool = False):
+        """
+        Helper method to upload JSON data directly to a signed URL.
+        
+        Args:
+            json_data: JSON data (dict or list) to upload
+            url: Signed URL to upload to
+            use_gzip: Whether to compress the data with gzip
+        """
         if hasattr(json, 'dumps') and 'ignore_nan' in json.dumps.__code__.co_varnames:
             dumps_kwargs = {'ignore_nan': True}
         else:
@@ -253,6 +265,7 @@ class MassStats:
             headers = {
                 'Content-Type': 'application/json'
             }
+        
         response = await self._client._regular_request("PUT", url, data=body, headers=headers)
         return response
     
@@ -386,36 +399,55 @@ class MassStats:
         region: str, 
         output: str, 
         config: Dict[str, Any],
-        request_json: Dict[str, Any], 
-        manifest_json: Dict[str, Any],
+        request_json: str,  # Path to request JSON file
         overwrite: bool = False, 
         skip_existing: bool = False,  
         location: str = None, 
         force_loc: bool = None, 
-        server: str = "dev-au.terrak.io"
+        server: str = None
     ) -> Dict[str, Any]:
         """
         Execute a mass stats job.
-
+        
         Args:
             name: The name of the job
             region: The region of the job
             output: The output of the job
             config: The config of the job
-            request_json: The request JSON
-            manifest_json: The manifest JSON
+            request_json: Path to the request JSON file
             overwrite: Whether to overwrite the job
             skip_existing: Whether to skip existing jobs
             location: The location of the job
             force_loc: Whether to force the location
             server: The server to use
-
+        
         Returns:
             API response as a dictionary
-
+        
         Raises:
             APIError: If the API request fails
         """
+        
+        def extract_manifest_from_request(request_data: List[Dict[str, Any]]) -> List[str]:
+            """Extract unique group names from request data to create manifest list."""
+            groups = []
+            seen_groups = set()
+            
+            for item in request_data:
+                if not isinstance(item, dict):
+                    raise ValueError("Each item in request JSON should be a dictionary")
+                    
+                if 'group' not in item:
+                    raise ValueError("Each item should have a 'group' field")
+                    
+                group = item['group']
+                if group not in seen_groups:
+                    groups.append(group)
+                    seen_groups.add(group)
+                    
+            return groups
+        
+        # Load and validate request JSON
         try:
             with open(request_json, 'r') as file:
                 request_data = json.load(file)
@@ -427,14 +459,35 @@ class MassStats:
             return e
         except json.JSONDecodeError as e:
             return e
-        upload_result = await self._upload_request(name = name, size = size, region = region, output = output, config = config, location = location, force_loc = force_loc, overwrite = overwrite, server = server, skip_existing = skip_existing)
+        
+        # Generate manifest from request data (kept in memory)
+        try:
+            manifest_groups = extract_manifest_from_request(request_data)
+        except Exception as e:
+            raise ValueError(f"Error extracting manifest from request JSON: {e}")
+        
+        # Get upload URLs
+        upload_result = await self._upload_request(
+            name=name, 
+            size=size, 
+            region=region, 
+            output=output, 
+            config=config, 
+            location=location, 
+            force_loc=force_loc, 
+            overwrite=overwrite, 
+            server=server, 
+            skip_existing=skip_existing
+        )
+        
         requests_url = upload_result.get('requests_url')
         manifest_url = upload_result.get('manifest_url')
+        
         if not requests_url:
             raise ValueError("No requests_url returned from server for request JSON upload")
         
+        # Upload request JSON file
         try:
-            # in this place we are uploading the request json file, we need to check whether the json is in the correct format or not
             self.validate_request(request_json)
             requests_response = await self._upload_file(request_json, requests_url, use_gzip=True)
             if requests_response.status not in [200, 201, 204]:
@@ -446,15 +499,17 @@ class MassStats:
         if not manifest_url:
             raise ValueError("No manifest_url returned from server for manifest JSON upload")
         
+        # Upload manifest JSON data directly (no temporary file needed)
         try:
-            manifest_response = await self._upload_file(manifest_json, manifest_url, use_gzip=False)
+            manifest_response = await self._upload_json_data(manifest_groups, manifest_url, use_gzip=False)
             if manifest_response.status not in [200, 201, 204]:
                 self._client.logger.error(f"Manifest upload error: {manifest_response.text()}")
                 raise Exception(f"Failed to upload manifest JSON: {manifest_response.text()}")
         except Exception as e:
-            raise Exception(f"Error uploading manifest JSON file {manifest_json}: {e}")
-
-        start_job_task_id =await self.start_job(upload_result.get("id"))
+            raise Exception(f"Error uploading manifest JSON: {e}")
+        
+        # Start the job
+        start_job_task_id = await self.start_job(upload_result.get("id"))
         return start_job_task_id
 
     @require_api_key