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.26"
+version = "0.0.27"
 description = ""
 authors = ["Wyvern AI <[email protected]>"]
 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 FeatureDataFrame
 from wyvern.entities.identifier_entities import ProductEntity
 from wyvern.entities.request import BaseWyvernRequest
 from wyvern.wyvern_request import WyvernRequest
@@ -65,7 +65,7 @@ def __init__(self):
             headers={},
             entity_store={},
             events=[],
-            feature_map=FeatureMap(feature_map={}),
+            feature_df=FeatureDataFrame(),
             model_output_map={},
         ),
     )

tests/feature_store/test_real_time_features.py

@@ -8,7 +8,7 @@
     RankingRealtimeFeatureComponent,
 )
 from wyvern.components.features.feature_store import feature_store_retrieval_component
-from wyvern.entities.feature_entities import FeatureMap
+from wyvern.entities.feature_entities import FeatureDataFrame
 from wyvern.feature_store.historical_feature_util import separate_real_time_features
 from wyvern.service import WyvernService
 
@@ -29,7 +29,7 @@ def mock_feature_store(mocker):
     mocker.patch.object(
         feature_store_retrieval_component,
         "fetch_features_from_feature_store",
-        return_value=FeatureMap(feature_map={}),
+        return_value=FeatureDataFrame(),
     )
 
 

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, FeatureDataFrame
 from wyvern.entities.identifier import Identifier
 from wyvern.entities.identifier_entities import ProductEntity, WyvernEntity
 from wyvern.entities.model_entities import ModelInput, ModelOutput
@@ -387,7 +387,7 @@ async def test_hydrate(mock_redis):
         entity_store={},
         model_output_map={},
         events=[],
-        feature_map=FeatureMap(feature_map={}),
+        feature_df=FeatureDataFrame(),
     )
     request_context.set(test_wyvern_request)
 
@@ -448,7 +448,7 @@ async def test_hydrate__duplicate_brand(mock_redis__duplicate_brand):
         headers={},
         entity_store={},
         events=[],
-        feature_map=FeatureMap(feature_map={}),
+        feature_df=FeatureDataFrame(),
         model_output_map={},
     )
     request_context.set(test_wyvern_request)

wyvern/__init__.py

@@ -19,7 +19,6 @@
     SingleEntityPipelineResponse,
 )
 from wyvern.entities.candidate_entities import CandidateSetEntity
-from wyvern.entities.feature_entities import FeatureData, FeatureMap
 from wyvern.entities.identifier import CompositeIdentifier, Identifier, IdentifierType
 from wyvern.entities.identifier_entities import (
     ProductEntity,
@@ -44,8 +43,6 @@
     "CandidateSetEntity",
     "ChainedModelInput",
     "CompositeIdentifier",
-    "FeatureData",
-    "FeatureMap",
     "Identifier",
     "IdentifierType",
     "ModelComponent",