feat: Add feature view lineage tab and filtering to home page lineage (#5308)

franciscojavierarceo · devin-ai-integration[bot] · web-flow · commit 308255d0f1aa · 2025-04-30T14:23:18.000-04:00
* Add feature view lineage tab and filtering to home page lineage

Co-Authored-By: Francisco Javier Arceo &lt;arceofrancisco@gmail.com&gt;

* Fix TypeScript errors in lineage visualization components

Co-Authored-By: Francisco Javier Arceo &lt;arceofrancisco@gmail.com&gt;

* Format FeatureViewLineageTab.tsx

Co-Authored-By: Francisco Javier Arceo &lt;arceofrancisco@gmail.com&gt;

---------

Co-authored-by: Devin AI &lt;158243242+devin-ai-integration[bot]@users.noreply.github.com&gt;
diff --git a/ui/src/components/RegistryVisualization.tsx b/ui/src/components/RegistryVisualization.tsx
@@ -572,12 +572,14 @@ interface RegistryVisualizationProps {
   registryData: feast.core.Registry;
   relationships: EntityRelation[];
   indirectRelationships: EntityRelation[];
+  filterNode?: { type: FEAST_FCO_TYPES; name: string };
 }
 
 const RegistryVisualization: React.FC<RegistryVisualizationProps> = ({
   registryData,
   relationships,
   indirectRelationships,
+  filterNode,
 }) => {
   const [nodes, setNodes, onNodesChange] = useNodesState([]);
   const [edges, setEdges, onEdgesChange] = useEdgesState([]);
@@ -592,10 +594,22 @@ const RegistryVisualization: React.FC<RegistryVisualizationProps> = ({
       setLoading(true);
 
       // Only include indirect relationships if the toggle is on
-      const relationshipsToShow = showIndirectRelationships
+      let relationshipsToShow = showIndirectRelationships
         ? [...relationships, ...indirectRelationships]
         : relationships;
 
+      // Filter relationships based on filterNode if provided
+      if (filterNode) {
+        relationshipsToShow = relationshipsToShow.filter((rel) => {
+          return (
+            (rel.source.type === filterNode.type &&
+              rel.source.name === filterNode.name) ||
+            (rel.target.type === filterNode.type &&
+              rel.target.name === filterNode.name)
+          );
+        });
+      }
+
       // Filter out invalid relationships
       const validRelationships = relationshipsToShow.filter((rel) => {
         // Add additional validation as needed for your use case
@@ -625,6 +639,7 @@ const RegistryVisualization: React.FC<RegistryVisualizationProps> = ({
     indirectRelationships,
     showIndirectRelationships,
     showIsolatedNodes,
+    filterNode,
     setNodes,
     setEdges,
   ]);
diff --git a/ui/src/components/RegistryVisualizationTab.tsx b/ui/src/components/RegistryVisualizationTab.tsx
@@ -1,12 +1,56 @@
-import React, { useContext } from "react";
-import { EuiEmptyPrompt, EuiLoadingSpinner, EuiSpacer } from "@elastic/eui";
+import React, { useContext, useState } from "react";
+import {
+  EuiEmptyPrompt,
+  EuiLoadingSpinner,
+  EuiSpacer,
+  EuiSelect,
+  EuiFormRow,
+  EuiFlexGroup,
+  EuiFlexItem,
+} from "@elastic/eui";
 import useLoadRegistry from "../queries/useLoadRegistry";
 import RegistryPathContext from "../contexts/RegistryPathContext";
 import RegistryVisualization from "./RegistryVisualization";
+import { FEAST_FCO_TYPES } from "../parsers/types";
 
 const RegistryVisualizationTab = () => {
   const registryUrl = useContext(RegistryPathContext);
   const { isLoading, isSuccess, isError, data } = useLoadRegistry(registryUrl);
+  const [selectedObjectType, setSelectedObjectType] = useState("");
+  const [selectedObjectName, setSelectedObjectName] = useState("");
+
+  const getObjectOptions = (objects: any, type: string) => {
+    switch (type) {
+      case "dataSource":
+        const dataSources = new Set<string>();
+        objects.featureViews?.forEach((fv: any) => {
+          if (fv.spec?.batchSource?.name)
+            dataSources.add(fv.spec.batchSource.name);
+        });
+        objects.streamFeatureViews?.forEach((sfv: any) => {
+          if (sfv.spec?.batchSource?.name)
+            dataSources.add(sfv.spec.batchSource.name);
+          if (sfv.spec?.streamSource?.name)
+            dataSources.add(sfv.spec.streamSource.name);
+        });
+        return Array.from(dataSources);
+      case "entity":
+        return objects.entities?.map((entity: any) => entity.spec?.name) || [];
+      case "featureView":
+        return [
+          ...(objects.featureViews?.map((fv: any) => fv.spec?.name) || []),
+          ...(objects.onDemandFeatureViews?.map(
+            (odfv: any) => odfv.spec?.name,
+          ) || []),
+          ...(objects.streamFeatureViews?.map((sfv: any) => sfv.spec?.name) ||
+            []),
+        ];
+      case "featureService":
+        return objects.featureServices?.map((fs: any) => fs.spec?.name) || [];
+      default:
+        return [];
+    }
+  };
 
   return (
     <>
@@ -31,10 +75,58 @@ const RegistryVisualizationTab = () => {
       {isSuccess && data && (
         <>
           <EuiSpacer size="l" />
+          <EuiFlexGroup style={{ marginBottom: 16 }}>
+            <EuiFlexItem grow={false} style={{ width: 200 }}>
+              <EuiFormRow label="Filter by type">
+                <EuiSelect
+                  options={[
+                    { value: "", text: "All" },
+                    { value: "dataSource", text: "Data Source" },
+                    { value: "entity", text: "Entity" },
+                    { value: "featureView", text: "Feature View" },
+                    { value: "featureService", text: "Feature Service" },
+                  ]}
+                  value={selectedObjectType}
+                  onChange={(e) => {
+                    setSelectedObjectType(e.target.value);
+                    setSelectedObjectName(""); // Reset name when type changes
+                  }}
+                  aria-label="Select object type"
+                />
+              </EuiFormRow>
+            </EuiFlexItem>
+            <EuiFlexItem grow={false} style={{ width: 300 }}>
+              <EuiFormRow label="Select object">
+                <EuiSelect
+                  options={[
+                    { value: "", text: "All" },
+                    ...getObjectOptions(data.objects, selectedObjectType).map(
+                      (name: string) => ({
+                        value: name,
+                        text: name,
+                      }),
+                    ),
+                  ]}
+                  value={selectedObjectName}
+                  onChange={(e) => setSelectedObjectName(e.target.value)}
+                  aria-label="Select object"
+                  disabled={selectedObjectType === ""}
+                />
+              </EuiFormRow>
+            </EuiFlexItem>
+          </EuiFlexGroup>
           <RegistryVisualization
             registryData={data.objects}
             relationships={data.relationships}
             indirectRelationships={data.indirectRelationships}
+            filterNode={
+              selectedObjectType && selectedObjectName
+                ? {
+                    type: selectedObjectType as FEAST_FCO_TYPES,
+                    name: selectedObjectName,
+                  }
+                : undefined
+            }
           />
         </>
       )}
diff --git a/ui/src/pages/feature-views/FeatureViewLineageTab.tsx b/ui/src/pages/feature-views/FeatureViewLineageTab.tsx
@@ -0,0 +1,61 @@
+import React, { useContext } from "react";
+import { useParams } from "react-router-dom";
+import { EuiEmptyPrompt, EuiLoadingSpinner } from "@elastic/eui";
+import { feast } from "../../protos";
+import useLoadRegistry from "../../queries/useLoadRegistry";
+import RegistryPathContext from "../../contexts/RegistryPathContext";
+import RegistryVisualization from "../../components/RegistryVisualization";
+import { FEAST_FCO_TYPES } from "../../parsers/types";
+
+interface FeatureViewLineageTabProps {
+  data: feast.core.IFeatureView;
+}
+
+const FeatureViewLineageTab = ({ data }: FeatureViewLineageTabProps) => {
+  const registryUrl = useContext(RegistryPathContext);
+  const {
+    isLoading,
+    isSuccess,
+    isError,
+    data: registryData,
+  } = useLoadRegistry(registryUrl);
+  const { featureViewName } = useParams();
+
+  const filterNode = {
+    type: FEAST_FCO_TYPES.featureView,
+    name: featureViewName || data.spec?.name || "",
+  };
+
+  return (
+    <>
+      {isLoading && (
+        <div style={{ display: "flex", justifyContent: "center", padding: 25 }}>
+          <EuiLoadingSpinner size="xl" />
+        </div>
+      )}
+      {isError && (
+        <EuiEmptyPrompt
+          iconType="alert"
+          color="danger"
+          title={<h2>Error Loading Registry Data</h2>}
+          body={
+            <p>
+              There was an error loading the Registry Data. Please check that{" "}
+              <code>feature_store.yaml</code> file is available and well-formed.
+            </p>
+          }
+        />
+      )}
+      {isSuccess && registryData && (
+        <RegistryVisualization
+          registryData={registryData.objects}
+          relationships={registryData.relationships}
+          indirectRelationships={registryData.indirectRelationships}
+          filterNode={filterNode}
+        />
+      )}
+    </>
+  );
+};
+
+export default FeatureViewLineageTab;
diff --git a/ui/src/pages/feature-views/RegularFeatureViewInstance.tsx b/ui/src/pages/feature-views/RegularFeatureViewInstance.tsx
@@ -6,6 +6,7 @@ import { FeatureViewIcon } from "../../graphics/FeatureViewIcon";
 
 import { useMatchExact, useMatchSubpath } from "../../hooks/useMatchSubpath";
 import RegularFeatureViewOverviewTab from "./RegularFeatureViewOverviewTab";
+import FeatureViewLineageTab from "./FeatureViewLineageTab";
 
 import {
   useRegularFeatureViewCustomTabs,
@@ -33,6 +34,14 @@ const RegularFeatureInstance = ({ data }: RegularFeatureInstanceProps) => {
     },
   ];
 
+  tabs.push({
+    label: "Lineage",
+    isSelected: useMatchSubpath("lineage"),
+    onClick: () => {
+      navigate("lineage");
+    },
+  });
+
   let statisticsIsSelected = useMatchSubpath("statistics");
   if (enabledFeatureStatistics) {
     tabs.push({
@@ -62,6 +71,10 @@ const RegularFeatureInstance = ({ data }: RegularFeatureInstanceProps) => {
             path="/"
             element={<RegularFeatureViewOverviewTab data={data} />}
           />
+          <Route
+            path="/lineage"
+            element={<FeatureViewLineageTab data={data} />}
+          />
           {TabRoutes}
         </Routes>
       </EuiPageTemplate.Section>
