fix: Fixed the entity to on-demand feature view relationship

Signed-off-by: ntkathole <[email protected]>

Author: ntkathole

sdk/python/feast/api/registry/rest/entities.py

@@ -109,8 +109,8 @@ def get_entity(
         seen_ds_names = set()
         for rel in relationships:
             if rel.get("target", {}).get("type") == "dataSource":
-                ds_name = rel["target"]["name"]
-                if ds_name not in seen_ds_names:
+                ds_name = rel["target"].get("name")
+                if ds_name and ds_name not in seen_ds_names:
                     ds_obj = ds_map.get(ds_name)
                     if ds_obj:
                         data_source_objs.append(ds_obj)

sdk/python/feast/lineage/registry_lineage.py

@@ -172,18 +172,31 @@ def _parse_direct_relationships(self, registry: Registry) -> List[EntityRelation
 
         # OnDemand FeatureView: Feature -> OnDemandFeatureView relationships
         for odfv in registry.on_demand_feature_views:
-            if hasattr(odfv, "spec") and odfv.spec and hasattr(odfv.spec, "features"):
-                for feature in odfv.spec.features:
-                    relationships.append(
-                        EntityRelation(
-                            source=EntityReference(
-                                FeastObjectType.FEATURE, feature.name
-                            ),
+            if hasattr(odfv, "spec") and odfv.spec:
+                # Entity relationships
+                if hasattr(odfv.spec, "entities"):
+                    for entity_name in odfv.spec.entities:
+                        rel = EntityRelation(
+                            source=EntityReference(FeastObjectType.ENTITY, entity_name),
                             target=EntityReference(
                                 FeastObjectType.FEATURE_VIEW, odfv.spec.name
                             ),
                         )
-                    )
+                        relationships.append(rel)
+
+                # Feature -> OnDemandFeatureView relationships
+                if hasattr(odfv.spec, "features"):
+                    for feature in odfv.spec.features:
+                        relationships.append(
+                            EntityRelation(
+                                source=EntityReference(
+                                    FeastObjectType.FEATURE, feature.name
+                                ),
+                                target=EntityReference(
+                                    FeastObjectType.FEATURE_VIEW, odfv.spec.name
+                                ),
+                            )
+                        )
 
         # OnDemand FeatureView relationships
         for odfv in registry.on_demand_feature_views:

sdk/python/tests/unit/api/test_api_rest_registry.py

@@ -1653,6 +1653,7 @@ def test_metrics_popular_tags_via_rest(fastapi_test_app):
             assert "feature_views" in tag_info
             assert "total_feature_views" in tag_info
 
+
 def test_all_endpoints_return_404_for_invalid_objects(fastapi_test_app):
     """Test that all REST API endpoints return 404 errors when objects are not found."""
 

sdk/python/tests/unit/infra/test_registry_lineage.py

@@ -231,6 +231,61 @@ def test_on_demand_feature_view_lineage(self):
         assert isinstance(relationships, list)
         assert isinstance(indirect_relationships, list)
 
+    def test_on_demand_feature_view_entity_relationships(self):
+        """Test entity relationships with on-demand feature views."""
+        registry = Registry()
+
+        # Create entities
+        entity1_spec = EntitySpecV2(name="user_id")
+        entity1 = Entity(spec=entity1_spec)
+
+        entity2_spec = EntitySpecV2(name="dummy")
+        entity2 = Entity(spec=entity2_spec)
+
+        registry.entities.extend([entity1, entity2])
+
+        # Create on-demand feature view with entities
+        odfv_spec = OnDemandFeatureViewSpec(name="transformed_features")
+        odfv_spec.entities.append("user_id")
+        odfv_spec.entities.append("dummy")
+        odfv = OnDemandFeatureView(spec=odfv_spec)
+        registry.on_demand_feature_views.append(odfv)
+
+        # Generate lineage
+        lineage_generator = RegistryLineageGenerator()
+        direct_relationships, indirect_relationships = (
+            lineage_generator.generate_lineage(registry)
+        )
+
+        # Filter Entity -> OnDemandFeatureView relationships
+        entity_to_odfv_relationships = [
+            rel
+            for rel in direct_relationships
+            if rel.source.type == FeastObjectType.ENTITY
+            and rel.target.type == FeastObjectType.FEATURE_VIEW
+            and rel.target.name == "transformed_features"
+        ]
+
+        # Should have 2 relationships: user_id -> transformed_features, dummy -> transformed_features
+        assert len(entity_to_odfv_relationships) == 2
+
+        # Check specific relationships
+        relationship_pairs = {
+            (rel.source.name, rel.target.name) for rel in entity_to_odfv_relationships
+        }
+
+        expected_pairs = {
+            ("user_id", "transformed_features"),
+            ("dummy", "transformed_features"),
+        }
+
+        assert relationship_pairs == expected_pairs
+
+        # Test relationship types
+        for rel in entity_to_odfv_relationships:
+            assert rel.source.type == FeastObjectType.ENTITY
+            assert rel.target.type == FeastObjectType.FEATURE_VIEW
+
     def test_stream_feature_view_lineage(self):
         """Test lineage with stream feature views."""
         registry = Registry()

ui/src/parsers/parseEntityRelationships.ts

@@ -57,6 +57,21 @@ const parseEntityRelationships = (objects: feast.core.Registry) => {
   });
 
   objects.onDemandFeatureViews?.forEach((fv) => {
+    // Entity relationships
+    fv.spec?.entities?.forEach((ent) => {
+      links.push({
+        source: {
+          type: FEAST_FCO_TYPES["entity"],
+          name: ent,
+        },
+        target: {
+          type: FEAST_FCO_TYPES["featureView"],
+          name: fv.spec?.name!,
+        },
+      });
+    });
+
+    // Data source relationships
     Object.values(fv.spec?.sources!).forEach(
       (input: { [key: string]: any }) => {
         if (input.requestDataSource) {