Utilize Polars.DataFrame for performance in ModelbitComponent

.pre-commit-config.yaml

@@ -62,6 +62,7 @@ repos:
           - types-tqdm
           - nest-asyncio
           - aiohttp
+          - polars
         exclude: "^tests/"
 
   # Check for spelling

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wyvern-ai"
-version = "0.0.24"
+version = "0.0.25"
 description = ""
 authors = ["Wyvern AI <info@wyvern.ai>"]
 readme = "README.md"
@@ -33,6 +33,7 @@ aiohttp = {extras = ["speedups"], version = "^3.8.5"}
 requests = "^2.31.0"
 platformdirs = "^3.8"
 posthog = "^3.0.2"
+polars = "^0.19.6"
 
 
 [tool.poetry.group.dev.dependencies]

tests/components/business_logic/test_pinning_business_logic.py

@@ -13,7 +13,7 @@
     PinningBusinessLogicComponent,
 )
 from wyvern.entities.candidate_entities import ScoredCandidate
-from wyvern.entities.feature_entities import FeatureMap
+from wyvern.entities.feature_entities import FeatureMap, FeatureMapPolars
 from wyvern.entities.identifier_entities import ProductEntity
 from wyvern.entities.request import BaseWyvernRequest
 from wyvern.wyvern_request import WyvernRequest
@@ -66,6 +66,7 @@ def __init__(self):
             entity_store={},
             events=[],
             feature_map=FeatureMap(feature_map={}),
+            feature_map_polars=FeatureMapPolars(feature_map=FeatureMap(feature_map={})),
             model_output_map={},
         ),
     )

tests/scenarios/test_product_ranking.py

@@ -19,7 +19,7 @@
 from wyvern.core.compression import wyvern_encode
 from wyvern.core.http import aiohttp_client
 from wyvern.entities.candidate_entities import CandidateSetEntity
-from wyvern.entities.feature_entities import FeatureData, FeatureMap
+from wyvern.entities.feature_entities import FeatureData, FeatureMap, FeatureMapPolars
 from wyvern.entities.identifier import Identifier
 from wyvern.entities.identifier_entities import ProductEntity, WyvernEntity
 from wyvern.entities.model_entities import ModelInput, ModelOutput
@@ -388,6 +388,7 @@ async def test_hydrate(mock_redis):
         model_output_map={},
         events=[],
         feature_map=FeatureMap(feature_map={}),
+        feature_map_polars=FeatureMapPolars(feature_map=FeatureMap(feature_map={})),
     )
     request_context.set(test_wyvern_request)
 
@@ -449,6 +450,7 @@ async def test_hydrate__duplicate_brand(mock_redis__duplicate_brand):
         entity_store={},
         events=[],
         feature_map=FeatureMap(feature_map={}),
+        feature_map_polars=FeatureMapPolars(feature_map=FeatureMap(feature_map={})),
         model_output_map={},
     )
     request_context.set(test_wyvern_request)

wyvern/components/component.py

@@ -8,6 +8,8 @@
 from typing import Dict, Generic, List, Optional, Set, Union
 from uuid import uuid4
 
+import polars as pl
+
 from wyvern import request_context
 from wyvern.entities.identifier import Identifier
 from wyvern.wyvern_typing import INPUT_TYPE, OUTPUT_TYPE, WyvernFeature
@@ -142,6 +144,19 @@ def manifest_feature_names(self) -> Set[str]:
         """
         return set()
 
+    def get_features(
+        self,
+        identifier_type: str,
+        identifier_list: List[str],
+        feature_names: List[str],
+    ) -> pl.DataFrame:
+        current_request = request_context.ensure_current_request()
+        return current_request.feature_map_polars.get_features(
+            identifier_type,
+            identifier_list,
+            feature_names,
+        )
+
     def get_feature(
         self,
         identifier: Identifier,

wyvern/components/features/feature_retrieval_pipeline.py

@@ -22,7 +22,7 @@
     RealtimeFeatureRequest,
 )
 from wyvern.entities.candidate_entities import CandidateSetEntity
-from wyvern.entities.feature_entities import FeatureData, FeatureMap
+from wyvern.entities.feature_entities import FeatureData, FeatureMap, FeatureMapPolars
 from wyvern.entities.feature_entity_helpers import feature_map_create, feature_map_join
 from wyvern.entities.identifier_entities import WyvernEntity
 from wyvern.wyvern_typing import REQUEST_ENTITY
@@ -160,8 +160,6 @@ async def execute(
                 **kwargs,
             )
         )
-        current_request = request_context.ensure_current_request()
-        current_request.feature_map = feature_retrieval_response
 
         """
         TODO (suchintan):
@@ -327,4 +325,9 @@ async def execute(
                 real_time_feature_responses,
             )
 
+        current_request = request_context.ensure_current_request()
+        current_request.feature_map = feature_responses
+        current_request.feature_map_polars = FeatureMapPolars(
+            feature_map=feature_responses,
+        )
         return feature_responses