diff --git a/src/brad/routing/functionality_catalog.py b/src/brad/routing/functionality_catalog.py index 1b561be7..43abfdc9 100644 --- a/src/brad/routing/functionality_catalog.py +++ b/src/brad/routing/functionality_catalog.py @@ -1,10 +1,11 @@ -from typing import List import operator import yaml +from typing import List, Tuple, Dict from functools import reduce -from typing import Dict from importlib.resources import files, as_file + import brad.routing as routing +from brad.config.engine import Engine, EngineBitmapValues class Functionality: @@ -33,13 +34,19 @@ def __init__(self): redshift_functionalities = engine["functionalities"] # Convert to bitmaps - engine_functionality_strings = [ - athena_functionalities, - aurora_functionalities, - redshift_functionalities, - ] self.engine_functionalities = [ - Functionality.to_bitmap(f) for f in engine_functionality_strings + ( + EngineBitmapValues[Engine.Athena], + Functionality.to_bitmap(athena_functionalities), + ), + ( + EngineBitmapValues[Engine.Aurora], + Functionality.to_bitmap(aurora_functionalities), + ), + ( + EngineBitmapValues[Engine.Redshift], + Functionality.to_bitmap(redshift_functionalities), + ), ] @staticmethod @@ -53,10 +60,14 @@ def to_bitmap(functionalities: List[str]) -> int: 0, ) - def get_engine_functionalities(self) -> List[int]: + def get_engine_functionalities(self) -> List[Tuple[int, int]]: """ Return a bitmap for each engine that states what functionalities the - engine supports + engine supports. + + The first value in the tuple is the bitmask representing the the engine. + The second value in the tuple is the bitmap representing its supported + functionalities. """ return self.engine_functionalities diff --git a/src/brad/routing/geospatial_keywords.yml b/src/brad/routing/geospatial_keywords.yml index b341a86c..59c1e8a3 100644 --- a/src/brad/routing/geospatial_keywords.yml +++ b/src/brad/routing/geospatial_keywords.yml @@ -90,8 +90,6 @@ - ST_GeogFromText - ST_GeogFromWKB - ST_GeographyFromText -- '=' -- '&&' - ST_Intersection - ST_Intersects - ST_Length @@ -233,9 +231,6 @@ - ST_GeomFromText - ST_GeomFromWKB - ST_GeometryN -- '=' -- '&<|' -- '&&' - ST_HasArc - ST_IsClosed - ST_IsEmpty @@ -280,8 +275,6 @@ - ST_GMLToSQL - ST_GeomFromGML - ST_GeomFromKML -- '&&' -- ~= - ST_HausdorffDistance - ST_Intersection - ST_Intersects diff --git a/src/brad/routing/router.py b/src/brad/routing/router.py index 60e29623..7f3b579b 100644 --- a/src/brad/routing/router.py +++ b/src/brad/routing/router.py @@ -151,23 +151,21 @@ def _run_functionality_routing(self, query: QueryRep) -> int: compute the set of engines that are able to serve this query. """ - # Bitmap describing what functionality is required for running query + # Bitmap describing what special functionality is required for running + # the query. req_bitmap = query.get_required_functionality() # Bitmap for each engine which features it supports engine_support = self.functionality_catalog.get_engine_functionalities() - engines = [Engine.Athena, Engine.Aurora, Engine.Redshift] # Narrow down the valid engines that can run the query, based on the # engine functionality - valid_locations_list = [] - for engine, sup_bitmap in zip(engines, engine_support): - query_supported = (~req_bitmap | (req_bitmap & sup_bitmap)) == -1 + supported_engines_bitmap = 0 + for engine_mask, sup_bitmap in engine_support: + if (req_bitmap & sup_bitmap) == req_bitmap: + supported_engines_bitmap |= engine_mask - if query_supported: - valid_locations_list.append(engine) - - return Engine.to_bitmap(valid_locations_list) + return supported_engines_bitmap def _run_location_routing( self, query: QueryRep, location_bitmap: Dict[str, int]