Fixed psql test

aligned/sources/psql.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Any
+from typing import TYPE_CHECKING, Any
 
 from aligned.data_source.batch_data_source import BatchDataSource, ColumnFeatureMappable
 from aligned.feature_source import WritableFeatureSource
@@ -13,7 +13,6 @@
 if TYPE_CHECKING:
     from aligned.compiler.feature_factory import FeatureFactory
     from aligned.enricher import Enricher
-    from aligned.entity_data_source import EntityDataSource
     from aligned.schemas.feature import EventTimestamp
 
 
@@ -50,11 +49,6 @@ def data_enricher(self, sql: str, values: dict | None = None) -> Enricher:
 
         return SqlDatabaseEnricher(self.env_var, sql, values)
 
-    def entity_source(self, timestamp_column: str, sql: Callable[[str], str]) -> EntityDataSource:
-        from aligned.compiler.model import SqlEntityDataSource
-
-        return SqlEntityDataSource(sql, self.env_var, timestamp_column)
-
     def fetch(self, query: str) -> RetrivalJob:
         from aligned.psql.jobs import PostgreSqlJob
 

conftest.py

@@ -447,6 +447,7 @@ def titanic_model(titanic_feature_view: FeatureView) -> ModelContractWrapper:
         ],
     )
     class Titanic:
+        passenger_id = Int32().as_entity()
 
         will_survive = features.survived.as_classification_label()  # type: ignore
 

test_data/titanic-sets.json

@@ -1 +1 @@
-{"raw_data": [], "train_test": [], "train_test_validation": [{"id": "titanic_test", "name": null, "request_result": {"entities": [{"name": "passenger_id", "dtype": {"name": "int32"}, "description": null, "tags": null, "constraints": null}], "features": [{"name": "age", "dtype": {"name": "float"}, "description": "A float as some have decimals", "tags": null, "constraints": [{"name": "lower_bound_inc", "value": 0.0}, {"name": "upper_bound_inc", "value": 100.0}]}, {"name": "is_mr", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "name", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "contains", "dtype": {"name": "bool"}, "key": "name", "value": "Mr."}, "depth": 1}, {"name": "has_siblings", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sibsp", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "int32"}}], "transformation": {"name": "not-equals", "dtype": {"name": "bool"}, "key": "sibsp", "value": {"name": "int", "value": 0}}, "depth": 1}, {"name": "is_female", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sex", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "equals", "dtype": {"name": "bool"}, "key": "sex", "value": {"name": "string", "value": "female"}}, "depth": 1}, {"name": "sex", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}, {"name": "in_domain", "values": ["male", "female"]}]}, {"name": "cabin", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}]}, {"name": "sibsp", "dtype": {"name": "int32"}, "description": "Number of siblings on titanic", "tags": null, "constraints": [{"name": "lower_bound_inc", "value": 0.0}, {"name": "optional"}, {"name": "upper_bound_inc", "value": 20.0}]}, {"name": "is_male", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sex", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "equals", "dtype": {"name": "bool"}, "key": "sex", "value": {"name": "string", "value": "male"}}, "depth": 1}, {"name": "name", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}]}, {"name": "survived", "dtype": {"name": "bool"}, "description": "If the passenger survived", "tags": null, "constraints": null}], "event_timestamp": null}, "train_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-train.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "test_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-test.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "validation_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-validate.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "train_size_fraction": 0.6, "test_size_fraction": 0.20000000000000007, "validate_size_fraction": 0.19999999999999996, "target": ["survived"], "description": null, "tags": null}], "active_learning": []}
+{"raw_data": [], "train_test": [], "train_test_validation": [{"id": "titanic_test", "name": null, "request_result": {"entities": [{"name": "passenger_id", "dtype": {"name": "int32"}, "description": null, "tags": null, "constraints": null}], "features": [{"name": "name", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}]}, {"name": "sex", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}, {"name": "in_domain", "values": ["male", "female"]}]}, {"name": "is_female", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sex", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "equals", "dtype": {"name": "bool"}, "key": "sex", "value": {"name": "string", "value": "female"}}, "depth": 1}, {"name": "is_mr", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "name", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "contains", "dtype": {"name": "bool"}, "key": "name", "value": "Mr."}, "depth": 1}, {"name": "is_male", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sex", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "string"}}], "transformation": {"name": "equals", "dtype": {"name": "bool"}, "key": "sex", "value": {"name": "string", "value": "male"}}, "depth": 1}, {"name": "age", "dtype": {"name": "float"}, "description": "A float as some have decimals", "tags": null, "constraints": [{"name": "lower_bound_inc", "value": 0.0}, {"name": "upper_bound_inc", "value": 100.0}]}, {"name": "sibsp", "dtype": {"name": "int32"}, "description": "Number of siblings on titanic", "tags": null, "constraints": [{"name": "lower_bound_inc", "value": 0.0}, {"name": "optional"}, {"name": "upper_bound_inc", "value": 20.0}]}, {"name": "has_siblings", "dtype": {"name": "bool"}, "description": null, "tags": null, "constraints": null, "depending_on": [{"name": "sibsp", "location": {"name": "titanic", "location": "feature_view"}, "dtype": {"name": "int32"}}], "transformation": {"name": "not-equals", "dtype": {"name": "bool"}, "key": "sibsp", "value": {"name": "int", "value": 0}}, "depth": 1}, {"name": "cabin", "dtype": {"name": "string"}, "description": null, "tags": null, "constraints": [{"name": "optional"}]}, {"name": "survived", "dtype": {"name": "bool"}, "description": "If the passenger survived", "tags": null, "constraints": null}], "event_timestamp": null}, "train_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-train.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "test_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-test.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "validation_dataset": {"mapping_keys": {}, "type_name": "csv", "path": "test_data/titanic-validate.csv", "csv_config": {"seperator": ",", "compression": "infer", "should_write_index": false}}, "train_size_fraction": 0.6, "test_size_fraction": 0.20000000000000007, "validate_size_fraction": 0.19999999999999996, "target": ["survived"], "description": null, "tags": null}], "active_learning": []}

test_data/titanic-test.csv

@@ -1,21 +1,21 @@
-age,passenger_id,is_mr,has_siblings,is_female,sex,cabin,sibsp,is_male,name,survived
-22.0,61,True,False,False,male,,0,True,"Sirayanian, Mr. Orsen",False
-38.0,62,False,False,True,female,B28,0,False,"Icard, Miss. Amelie",True
-45.0,63,True,True,False,male,C83,1,True,"Harris, Mr. Henry Birkhardt",False
-4.0,64,False,True,False,male,,3,True,"Skoog, Master. Harald",False
-,65,True,False,False,male,,0,True,"Stewart, Mr. Albert A",False
-,66,False,True,False,male,,1,True,"Moubarek, Master. Gerios",True
-29.0,67,True,False,True,female,F33,0,False,"Nye, Mrs. (Elizabeth Ramell)",True
-19.0,68,True,False,False,male,,0,True,"Crease, Mr. Ernest James",False
-17.0,69,False,True,True,female,,4,False,"Andersson, Miss. Erna Alexandra",True
-26.0,70,True,True,False,male,,2,True,"Kink, Mr. Vincenz",False
-32.0,71,True,False,False,male,,0,True,"Jenkin, Mr. Stephen Curnow",False
-16.0,72,False,True,True,female,,5,False,"Goodwin, Miss. Lillian Amy",False
-21.0,73,True,False,False,male,,0,True,"Hood, Mr. Ambrose Jr",False
-26.0,74,True,True,False,male,,1,True,"Chronopoulos, Mr. Apostolos",False
-32.0,75,True,False,False,male,,0,True,"Bing, Mr. Lee",True
-25.0,76,True,False,False,male,F G73,0,True,"Moen, Mr. Sigurd Hansen",False
-,77,True,False,False,male,,0,True,"Staneff, Mr. Ivan",False
-,78,True,False,False,male,,0,True,"Moutal, Mr. Rahamin Haim",False
-0.83,79,False,False,False,male,,0,True,"Caldwell, Master. Alden Gates",True
-30.0,80,False,False,True,female,,0,False,"Dowdell, Miss. Elizabeth",True
+name,sex,is_female,is_mr,is_male,age,sibsp,has_siblings,passenger_id,cabin,survived
+"Sirayanian, Mr. Orsen",male,False,True,True,22.0,0,False,61,,False
+"Icard, Miss. Amelie",female,True,False,False,38.0,0,False,62,B28,True
+"Harris, Mr. Henry Birkhardt",male,False,True,True,45.0,1,True,63,C83,False
+"Skoog, Master. Harald",male,False,False,True,4.0,3,True,64,,False
+"Stewart, Mr. Albert A",male,False,True,True,,0,False,65,,False
+"Moubarek, Master. Gerios",male,False,False,True,,1,True,66,,True
+"Nye, Mrs. (Elizabeth Ramell)",female,True,True,False,29.0,0,False,67,F33,True
+"Crease, Mr. Ernest James",male,False,True,True,19.0,0,False,68,,False
+"Andersson, Miss. Erna Alexandra",female,True,False,False,17.0,4,True,69,,True
+"Kink, Mr. Vincenz",male,False,True,True,26.0,2,True,70,,False
+"Jenkin, Mr. Stephen Curnow",male,False,True,True,32.0,0,False,71,,False
+"Goodwin, Miss. Lillian Amy",female,True,False,False,16.0,5,True,72,,False
+"Hood, Mr. Ambrose Jr",male,False,True,True,21.0,0,False,73,,False
+"Chronopoulos, Mr. Apostolos",male,False,True,True,26.0,1,True,74,,False
+"Bing, Mr. Lee",male,False,True,True,32.0,0,False,75,,True
+"Moen, Mr. Sigurd Hansen",male,False,True,True,25.0,0,False,76,F G73,False
+"Staneff, Mr. Ivan",male,False,True,True,,0,False,77,,False
+"Moutal, Mr. Rahamin Haim",male,False,True,True,,0,False,78,,False
+"Caldwell, Master. Alden Gates",male,False,False,True,0.83,0,False,79,,True
+"Dowdell, Miss. Elizabeth",female,True,False,False,30.0,0,False,80,,True