earthcatalog 0.2.0__py3-none-any.whl

earthcatalog/tests/test_spatial_resolver.py

@@ -0,0 +1,413 @@
+"""Tests for spatial partition resolver functionality."""
+
+import json
+import shutil
+import tempfile
+from pathlib import Path
+
+import pytest
+from shapely.geometry import Point, box
+
+# Import from earthcatalog package
+from earthcatalog.spatial_resolver import (
+    SpatialPartitionResolver,
+    resolve_and_query,
+    spatial_resolver,
+)
+
+
+class TestSpatialPartitionResolver:
+    """Test SpatialPartitionResolver class."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.temp_dir = tempfile.mkdtemp()
+        self.catalog_path = Path(self.temp_dir) / "catalog"
+        self.catalog_path.mkdir(parents=True, exist_ok=True)
+
+        # Create mock schema configurations
+        self.h3_schema = {
+            "spatial_partitioning": {
+                "grid_system": "h3",
+                "resolution": 6,
+                "coordinate_system": "EPSG:4326",
+                "example_paths": [
+                    "testmission/partition=h3/level=6/86283082fffffff/year=2024/month=01/items.parquet",
+                    "testmission/partition=h3/level=6/global/year=2024/month=01/items.parquet",
+                ],
+            },
+            "global_partitioning": {
+                "enabled": True,
+                "threshold": 50,
+                "description": "Items spanning >50 H3 cells go to /global/",
+            },
+        }
+
+        self.geojson_schema = {
+            "spatial_partitioning": {
+                "grid_system": "geojson",
+                "custom_grid": True,
+                "geojson_source": str(Path(__file__).parent.parent / "examples" / "custom_tiles.geojson"),
+                "custom_tiles": {
+                    "total_tiles": 3,
+                    "tile_ids": ["region_a", "region_b", "region_c"],
+                },
+            },
+            "global_partitioning": {"enabled": False},
+        }
+
+    def teardown_method(self):
+        """Clean up test fixtures."""
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+    def test_resolver_initialization(self):
+        """Test resolver initialization."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        assert resolver.grid_system == "h3"
+        assert resolver.global_enabled is True
+        assert resolver.global_threshold == 50
+        assert resolver.catalog_path == self.catalog_path
+
+    def test_h3_point_resolution(self):
+        """Test H3 point resolution."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # San Francisco coordinates
+        sf_point = Point(-122.4194, 37.7749)
+        partitions = resolver.resolve_partitions(sf_point)
+
+        assert len(partitions) == 1
+        assert isinstance(partitions[0], str)
+        assert len(partitions[0]) == 15  # H3 cell ID length
+        assert "global" not in partitions  # Small query shouldn't include global
+
+    def test_h3_polygon_resolution(self):
+        """Test H3 polygon resolution."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # San Francisco Bay Area
+        sf_bay = box(-122.5, 37.5, -122.0, 38.0)
+        partitions = resolver.resolve_partitions(sf_bay, overlap=True)
+
+        assert len(partitions) > 10  # Should cover multiple cells
+        assert all(isinstance(p, str) for p in partitions if p != "global")
+        # Moderate area might include global depending on exact cell count
+
+    def test_h3_buffer_cells(self):
+        """Test H3 resolution with buffer cells."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        sf_point = Point(-122.4194, 37.7749)
+
+        # No buffer
+        partitions_no_buffer = resolver.resolve_partitions(sf_point, buffer_cells=0)
+
+        # With buffer
+        partitions_with_buffer = resolver.resolve_partitions(sf_point, buffer_cells=2)
+
+        assert len(partitions_no_buffer) == 1
+        assert len(partitions_with_buffer) > len(partitions_no_buffer)
+        assert partitions_no_buffer[0] in partitions_with_buffer
+
+    def test_global_partition_threshold_logic(self):
+        """Test global partition inclusion based on threshold."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Small area - should not include global
+        small_area = Point(-122.4194, 37.7749).buffer(0.01)
+        small_partitions = resolver.resolve_partitions(small_area)
+        assert "global" not in small_partitions
+
+        # Large area - should include global due to threshold
+        large_area = box(-125.0, 32.0, -115.0, 42.0)  # Entire California
+        large_partitions = resolver.resolve_partitions(large_area)
+        assert "global" in large_partitions
+        assert len([p for p in large_partitions if p != "global"]) > 50
+
+    def test_global_partition_geometry_size_logic(self):
+        """Test global partition inclusion based on geometry size."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Very large geometry (continental scale)
+        continental = box(-130.0, 25.0, -65.0, 50.0)
+        partitions = resolver.resolve_partitions(continental)
+
+        assert "global" in partitions
+        # Should trigger large geometry threshold even if cell count logic doesn't
+
+    def test_global_partition_manual_override(self):
+        """Test manual global partition inclusion/exclusion."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        small_area = Point(-122.4194, 37.7749).buffer(0.01)
+        large_area = box(-125.0, 32.0, -115.0, 42.0)
+
+        # Force global on small area
+        partitions_force_global = resolver.resolve_partitions(small_area, include_global=True)
+        assert "global" in partitions_force_global
+
+        # Disable global on large area
+        partitions_no_global = resolver.resolve_partitions(large_area, include_global=False)
+        assert "global" not in partitions_no_global
+
+    def test_global_partition_disabled(self):
+        """Test behavior when global partitioning is disabled."""
+        schema_no_global = {
+            "spatial_partitioning": {"grid_system": "h3", "resolution": 6},
+            "global_partitioning": {"enabled": False},
+        }
+
+        resolver = SpatialPartitionResolver(schema_no_global, str(self.catalog_path))
+
+        # Even large areas shouldn't include global when disabled
+        large_area = box(-125.0, 32.0, -115.0, 42.0)
+        partitions = resolver.resolve_partitions(large_area)
+
+        assert "global" not in partitions
+        assert resolver.global_enabled is False
+
+    def test_geojson_resolution(self):
+        """Test GeoJSON grid resolution."""
+        resolver = SpatialPartitionResolver(self.geojson_schema, str(self.catalog_path))
+
+        # Area that should intersect region_a (Northern region: -120 to -110, 40 to 50)
+        northern_area = box(-119, 41, -111, 49)
+        partitions = resolver.resolve_partitions(northern_area)
+
+        assert "region_a" in partitions
+        assert len(partitions) >= 1
+
+        # Large area spanning multiple regions
+        large_area = box(-119, 32, -105, 48)
+        spanning_partitions = resolver.resolve_partitions(large_area)
+
+        assert len(spanning_partitions) > 1
+        assert "region_a" in spanning_partitions
+        assert "region_b" in spanning_partitions
+
+    def test_path_existence_checking(self):
+        """Test existing partition path filtering."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Create some test directories in file structure
+        test_partitions = ["86283082fffffff", "862830827ffffff", "global"]
+        existing_partitions = test_partitions[:2]  # Only create first two
+
+        for partition_id in existing_partitions:
+            partition_path = self.catalog_path / "testmission" / "partition=h3" / "level=6" / partition_id
+            partition_path.mkdir(parents=True, exist_ok=True)
+
+        # Test path filtering
+        existing_paths = resolver.get_existing_partition_paths(test_partitions)
+
+        assert len(existing_paths) == len(existing_partitions)
+        for partition_id in existing_partitions:
+            assert any(partition_id in path for path in existing_paths)
+
+    def test_query_path_generation(self):
+        """Test query path pattern generation with Hive-style temporal partitioning."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Create test directories in file structure
+        test_partitions = ["86283082fffffff", "global"]
+        for partition_id in test_partitions:
+            partition_path = self.catalog_path / "testmission" / "partition=h3" / "level=6" / partition_id
+            partition_path.mkdir(parents=True, exist_ok=True)
+
+        # Test without temporal filter - should glob all items files
+        query_patterns = resolver.generate_query_paths(test_partitions)
+        assert len(query_patterns) == len(test_partitions)
+        assert all("/**/items.parquet" in pattern for pattern in query_patterns)
+
+        # Test with temporal filter - should convert to Hive-style path
+        temporal_patterns = resolver.generate_query_paths(test_partitions, "2024-01")
+        assert len(temporal_patterns) == len(test_partitions)
+        assert all("year=2024/month=01/items.parquet" in pattern for pattern in temporal_patterns)
+
+    def test_different_grid_systems(self):
+        """Test resolution across different grid systems."""
+        test_point = Point(-122.4194, 37.7749)
+
+        grid_schemas = {
+            "h3": {"spatial_partitioning": {"grid_system": "h3", "resolution": 6}},
+            "s2": {"spatial_partitioning": {"grid_system": "s2", "level": 13}},
+            "utm": {"spatial_partitioning": {"grid_system": "utm", "precision": 1}},
+            "latlon": {"spatial_partitioning": {"grid_system": "latlon", "cell_size_degrees": 0.1}},
+        }
+
+        for grid_name, schema in grid_schemas.items():
+            schema["global_partitioning"] = {"enabled": False}  # Disable for simplicity
+
+            try:
+                resolver = SpatialPartitionResolver(schema, str(self.catalog_path))
+                partitions = resolver.resolve_partitions(test_point)
+
+                assert len(partitions) >= 1
+                assert all(isinstance(p, str) for p in partitions)
+
+            except ImportError:
+                # Skip if optional dependencies not available
+                pytest.skip(f"Optional dependency for {grid_name} grid not available")
+
+    def test_effective_global_threshold(self):
+        """Test effective global threshold calculation."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Should use configured threshold
+        threshold = resolver._get_effective_global_threshold()
+        assert threshold == 50
+
+        # Test with custom thresholds in schema
+        schema_with_custom = self.h3_schema.copy()
+        schema_with_custom["custom_thresholds"] = {"h3": {"6": 100, "7": 200}}
+
+        resolver_custom = SpatialPartitionResolver(schema_with_custom, str(self.catalog_path))
+        custom_threshold = resolver_custom._get_effective_global_threshold()
+        assert custom_threshold == 100  # Should use custom threshold for H3 resolution 6
+
+    def test_parse_temporal_filter_to_hive(self):
+        """Test temporal filter to Hive path conversion."""
+        resolver = SpatialPartitionResolver(self.h3_schema, str(self.catalog_path))
+
+        # Test year only
+        assert resolver._parse_temporal_filter_to_hive("2024") == "year=2024"
+
+        # Test year-month
+        assert resolver._parse_temporal_filter_to_hive("2024-01") == "year=2024/month=01"
+
+        # Test year-month-day
+        assert resolver._parse_temporal_filter_to_hive("2024-01-15") == "year=2024/month=01/day=15"
+
+        # Test year with wildcard
+        assert resolver._parse_temporal_filter_to_hive("2024-*") == "year=2024/*"
+
+        # Test year-month with wildcard
+        assert resolver._parse_temporal_filter_to_hive("2024-01-*") == "year=2024/month=01/*"
+
+        # Test empty or wildcard only
+        assert resolver._parse_temporal_filter_to_hive("*") == "*"
+        assert resolver._parse_temporal_filter_to_hive("") == "*"
+
+        # Test single-digit month padding
+        assert resolver._parse_temporal_filter_to_hive("2024-1") == "year=2024/month=01"
+
+
+class TestConvenienceFunctions:
+    """Test convenience functions."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        self.temp_dir = tempfile.mkdtemp()
+        self.schema_file = Path(self.temp_dir) / "schema.json"
+        self.catalog_path = Path(self.temp_dir) / "catalog"
+
+        # Create test schema
+        schema = {
+            "spatial_partitioning": {"grid_system": "h3", "resolution": 7},
+            "global_partitioning": {"enabled": True, "threshold": 50},
+        }
+
+        with open(self.schema_file, "w") as f:
+            json.dump(schema, f)
+
+    def teardown_method(self):
+        """Clean up test fixtures."""
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+    def test_create_resolver_from_schema_file(self):
+        """Test creating resolver from schema file (using modern spatial_resolver function)."""
+        resolver = spatial_resolver(str(self.schema_file), str(self.catalog_path))
+
+        assert isinstance(resolver, SpatialPartitionResolver)
+        assert resolver.grid_system == "h3"
+        assert resolver.spatial_config["resolution"] == 7
+
+    def test_spatial_resolver_function(self):
+        """Test the new spatial_resolver function."""
+        resolver = spatial_resolver(str(self.schema_file), str(self.catalog_path))
+
+        assert isinstance(resolver, SpatialPartitionResolver)
+        assert resolver.grid_system == "h3"
+        assert resolver.spatial_config["resolution"] == 7
+
+        # Test with catalog_path defaulting to schema directory
+        resolver2 = spatial_resolver(str(self.schema_file))
+        assert isinstance(resolver2, SpatialPartitionResolver)
+
+    def test_remote_schema_support(self):
+        """Test remote URL support in spatial_resolver function."""
+        # Test error when catalog_path not provided for various remote URLs
+        remote_urls = [
+            "s3://bucket/schema.json",
+            "s3a://bucket/schema.json",
+            "s3n://bucket/schema.json",
+            "gcs://bucket/schema.json",
+            "gs://bucket/schema.json",
+            "https://example.com/schema.json",
+            "abfs://container/schema.json",
+        ]
+
+        for url in remote_urls:
+            with pytest.raises(ValueError, match="catalog_path must be explicitly provided"):
+                spatial_resolver(url)
+
+    def test_resolve_and_query_function(self):
+        """Test the convenience resolve_and_query function."""
+        # Create a test geometry
+        test_geometry = box(-122.5, 37.7, -122.0, 38.0)
+
+        partition_ids, query = resolve_and_query(
+            schema_path=str(self.schema_file),
+            catalog_path=str(self.catalog_path),
+            aoi_geometry=test_geometry,
+            temporal_filter="2024-*",
+            overlap=True,
+        )
+
+        assert isinstance(partition_ids, list)
+        assert len(partition_ids) > 0
+        assert isinstance(query, str)
+        # Query will be empty if no existing partition paths, which is expected in test
+
+
+class TestErrorHandling:
+    """Test error handling and edge cases."""
+
+    def test_unsupported_grid_system(self):
+        """Test error handling for unsupported grid systems."""
+        schema = {
+            "spatial_partitioning": {"grid_system": "invalid_grid"},
+            "global_partitioning": {"enabled": False},
+        }
+
+        resolver = SpatialPartitionResolver(schema, "/tmp")
+
+        with pytest.raises(ValueError, match="Unsupported grid system"):
+            resolver.resolve_partitions(Point(0, 0))
+
+    def test_missing_dependencies(self):
+        """Test handling of missing optional dependencies."""
+        # This test would need to mock the import failures
+        # Implementation depends on how you want to handle missing dependencies
+        pass
+
+    def test_invalid_geometry(self):
+        """Test handling of invalid geometries."""
+        from shapely.errors import GeometryTypeError
+
+        schema = {
+            "spatial_partitioning": {"grid_system": "h3", "resolution": 6},
+            "global_partitioning": {"enabled": False},
+        }
+
+        resolver = SpatialPartitionResolver(schema, "/tmp")
+
+        # Test with invalid geometry dict
+        with pytest.raises(GeometryTypeError):
+            resolver.resolve_partitions({"type": "Invalid"})
+
+
+if __name__ == "__main__":
+    # Allow running tests directly
+    pytest.main([__file__])