Merge pull request #8 from terrafloww/feat/speed_up_indexer

Migrate codebase from httpx to aiohttp

Author: print-sid8

setup.py

@@ -8,7 +8,7 @@
 
 setup(
     name="rasteret",
-    version="0.1.17",
+    version="0.1.20",
     author="Sidharth Subramaniam",
     author_email="sid@terrafloww.com",
     description="Fast and efficient access to Cloud-Optimized GeoTIFFs (COGs)",

src/rasteret/fetch/cog.py

@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 import asyncio
-import httpx
+import aiohttp
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Tuple
 
@@ -50,32 +50,37 @@ class COGTileRequest:
 class COGReader:
     """Manages connection pooling and COG reading operations."""
 
-    def __init__(self, max_concurrent: int = 50):
+    def __init__(self, max_concurrent: int = 150):  # Increased from 50
         self.max_concurrent = max_concurrent
-        self.limits = httpx.Limits(
-            max_keepalive_connections=max_concurrent,
-            max_connections=max_concurrent,
-            keepalive_expiry=60.0,  # Shorter keepalive for HTTP/2
-        )
-        self.timeout = httpx.Timeout(30.0, connect=10.0)
+        # aiohttp connector settings
+        self.connector_limit = max_concurrent
+        self.connector_limit_per_host = max_concurrent
+        self.keepalive_timeout = 60.0
         self.client = None
         self.sem = None
-        self.batch_size = 12  # Reduced for better HTTP/2 multiplexing
+        self.batch_size = 20  # Increased for better throughput
 
     async def __aenter__(self):
-        self.client = httpx.AsyncClient(
-            timeout=self.timeout,
-            limits=self.limits,
-            http2=True,
-            verify=True,
-            trust_env=True,
+        # Create aiohttp connector with optimized settings
+        connector = aiohttp.TCPConnector(
+            limit=self.connector_limit,
+            limit_per_host=self.connector_limit_per_host,
+            keepalive_timeout=self.keepalive_timeout,
+            enable_cleanup_closed=True,
+        )
+
+        # Create aiohttp client session
+        timeout = aiohttp.ClientTimeout(total=30.0, connect=10.0)
+        self.client = aiohttp.ClientSession(
+            connector=connector,
+            timeout=timeout,
         )
         self.sem = asyncio.Semaphore(self.max_concurrent)
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.client:
-            await self.client.aclose()
+            await self.client.close()
 
     def merge_ranges(
         self, requests: List[COGTileRequest], gap_threshold: int = 1024
@@ -155,9 +160,9 @@ async def _read_range(self, url: str, start: int, end: int) -> bytes:
         for attempt in range(3):
             try:
                 async with self.sem:
-                    response = await self.client.get(url, headers=headers)
-                    response.raise_for_status()
-                    return response.content
+                    async with self.client.get(url, headers=headers) as response:
+                        response.raise_for_status()
+                        return await response.read()
             except Exception:
                 if attempt == 2:
                     raise
@@ -363,7 +368,7 @@ async def read_cog_tile_data(
     url: str,
     metadata: CogMetadata,
     geometry: Optional[Polygon] = None,
-    max_concurrent: int = 50,
+    max_concurrent: int = 150,  # Increased from 50
     debug: bool = False,
 ) -> Tuple[np.ndarray, Optional[Affine]]:
     """Read COG data, optionally masked by geometry."""

src/rasteret/stac/indexer.py

@@ -51,7 +51,7 @@ def __init__(
         name: Optional[str] = None,
         cloud_provider: Optional[CloudProvider] = None,
         cloud_config: Optional[CloudConfig] = None,
-        max_concurrent: int = 100,
+        max_concurrent: int = 300,  # Increased from 100
     ):
         self.data_source = data_source
         self.stac_api = stac_api
@@ -60,7 +60,7 @@ def __init__(
         self.cloud_config = cloud_config
         self.name = name
         self.max_concurrent = max_concurrent
-        self.batch_size = 50
+        self.batch_size = 100  # Increased from 50
 
     @property
     def band_map(self) -> Dict[str, str]:

src/rasteret/stac/parser.py

@@ -22,7 +22,7 @@
 import time
 from typing import Dict, List, Optional, Set, Any
 
-import httpx
+import aiohttp
 from cachetools import TTLCache, LRUCache
 
 from rasteret.types import CogMetadata
@@ -80,8 +80,8 @@ class AsyncCOGHeaderParser:
 
     def __init__(
         self,
-        max_concurrent: int = 100,
-        batch_size: int = 50,
+        max_concurrent: int = 300,  # Increased from 100
+        batch_size: int = 100,  # Increased from 50
         cache_ttl: int = 3600,  # 1 hour
         retry_attempts: int = 3,
         cloud_provider: Optional[CloudProvider] = None,
@@ -93,12 +93,10 @@ def __init__(
         self.cloud_config = cloud_config
         self.batch_size = batch_size
 
-        # Connection optimization
-        self.connector = httpx.Limits(
-            max_keepalive_connections=max_concurrent,
-            max_connections=max_concurrent,
-            keepalive_expiry=120,
-        )
+        # Connection optimization for aiohttp
+        self.connector_limit = max_concurrent
+        self.connector_limit_per_host = max_concurrent
+        self.keepalive_timeout = 120
 
         # Rate limiting
         self.semaphore = asyncio.Semaphore(max_concurrent)
@@ -118,17 +116,26 @@ def __init__(
         }
 
     async def __aenter__(self):
-        self.client = httpx.AsyncClient(
-            limits=self.connector,
-            timeout=30.0,
-            http2=True,
-            headers={"Connection": "keep-alive", "Keep-Alive": "timeout=120"},
+        # Create aiohttp connector with optimized settings
+        connector = aiohttp.TCPConnector(
+            limit=self.connector_limit,
+            limit_per_host=self.connector_limit_per_host,
+            keepalive_timeout=self.keepalive_timeout,
+            enable_cleanup_closed=True,
+        )
+
+        # Create aiohttp client session
+        timeout = aiohttp.ClientTimeout(total=30.0, connect=10.0)
+        self.client = aiohttp.ClientSession(
+            connector=connector,
+            timeout=timeout,
+            headers={"Connection": "keep-alive"},
         )
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.client:
-            await self.client.aclose()
+            await self.client.close()
 
     async def process_cog_headers_batch(
         self,
@@ -187,15 +194,15 @@ async def _fetch_byte_range(self, url: str, start: int, size: int) -> bytes:
             for attempt in range(self.retry_attempts):
                 try:
                     async with self.semaphore:
-                        response = await self.client.get(url, headers=headers)
-                        if response.status_code != 206:
-                            raise IOError(
-                                f"Range request failed: {response.status_code}"
-                            )
-
-                        data = response.content
-                        self.header_cache[cache_key] = data
-                        return data
+                        async with self.client.get(url, headers=headers) as response:
+                            if response.status != 206:
+                                raise IOError(
+                                    f"Range request failed: {response.status}"
+                                )
+
+                            data = await response.read()
+                            self.header_cache[cache_key] = data
+                            return data
 
                 except Exception as e:
                     if attempt == self.retry_attempts - 1:

src/rasteret/tests/test_stac_indexer.py

@@ -79,7 +79,9 @@ async def test_stac_search(self, mock_pystac):
         # Configure mock chain
         mock_pystac.Client.open.return_value = mock_client
         mock_client.search.return_value = mock_search
-        mock_search.items.return_value = iter(self.mock_stac_items)
+        mock_search.items_as_dicts.return_value = [
+            item.to_dict() for item in self.mock_stac_items
+        ]
 
         # Create indexer
         indexer = StacToGeoParquetIndexer(
@@ -88,7 +90,9 @@ async def test_stac_search(self, mock_pystac):
 
         # Test search
         items = await indexer._search_stac(
-            bbox=[-180, -90, 180, 90], date_range=["2023-01-01", "2023-12-31"]
+            bbox=[-180, -90, 180, 90],
+            date_range=["2023-01-01", "2023-12-31"],
+            query=None,
         )
 
         # Verify results
@@ -98,7 +102,7 @@ async def test_stac_search(self, mock_pystac):
         # Verify mock calls
         mock_pystac.Client.open.assert_called_once()
         mock_client.search.assert_called_once()
-        mock_search.items.assert_called_once()
+        mock_search.items_as_dicts.assert_called_once()
 
     @patch("rasteret.stac.indexer.AsyncCOGHeaderParser")
     @patch("rasteret.stac.indexer.pystac_client")
@@ -109,7 +113,9 @@ async def test_index_creation(self, mock_pystac, mock_parser):
 
         mock_pystac.Client.open.return_value = mock_client
         mock_client.search.return_value = mock_search
-        mock_search.items.return_value = iter(self.mock_stac_items)
+        mock_search.items_as_dicts.return_value = [
+            item.to_dict() for item in self.mock_stac_items
+        ]
 
         # Setup COG parser mock
         mock_parser_instance = AsyncMock()