lineage: fix expand collapse operation on nodes (#21309)

openmetadata-ui/src/main/resources/ui/playwright/e2e/Flow/LineageSettings.spec.ts

@@ -13,7 +13,9 @@
 import test, { expect } from '@playwright/test';
 import { get } from 'lodash';
 import { GlobalSettingOptions } from '../../constant/settings';
+import { ContainerClass } from '../../support/entity/ContainerClass';
 import { DashboardClass } from '../../support/entity/DashboardClass';
+import { MetricClass } from '../../support/entity/MetricClass';
 import { MlModelClass } from '../../support/entity/MlModelClass';
 import { SearchIndexClass } from '../../support/entity/SearchIndexClass';
 import { TableClass } from '../../support/entity/TableClass';
@@ -27,6 +29,7 @@ import {
 import {
   addPipelineBetweenNodes,
   fillLineageConfigForm,
+  performCollapse,
   performExpand,
   performZoomOut,
   verifyColumnLayerActive,
@@ -48,6 +51,8 @@ test.describe('Lineage Settings Tests', () => {
     const dashboard = new DashboardClass();
     const mlModel = new MlModelClass();
     const searchIndex = new SearchIndexClass();
+    const container = new ContainerClass();
+    const metric = new MetricClass();
 
     try {
       await Promise.all([
@@ -56,12 +61,16 @@ test.describe('Lineage Settings Tests', () => {
         dashboard.create(apiContext),
         mlModel.create(apiContext),
         searchIndex.create(apiContext),
+        container.create(apiContext),
+        metric.create(apiContext),
       ]);
 
       await addPipelineBetweenNodes(page, table, topic);
       await addPipelineBetweenNodes(page, topic, dashboard);
       await addPipelineBetweenNodes(page, dashboard, mlModel);
       await addPipelineBetweenNodes(page, mlModel, searchIndex);
+      await addPipelineBetweenNodes(page, searchIndex, container);
+      await addPipelineBetweenNodes(page, container, metric);
 
       await test.step(
         'Lineage config should throw error if upstream depth is less than 0',
@@ -168,7 +177,17 @@ test.describe('Lineage Settings Tests', () => {
           await verifyNodePresent(page, topic);
           await verifyNodePresent(page, mlModel);
           await performExpand(page, mlModel, false, searchIndex);
-          await performExpand(page, topic, true);
+          await performExpand(page, searchIndex, false, container);
+          await performExpand(page, container, false, metric);
+          await performExpand(page, topic, true, table);
+
+          // perform collapse
+          await performCollapse(page, mlModel, false, [
+            searchIndex,
+            container,
+            metric,
+          ]);
+          await performCollapse(page, dashboard, true, [table, topic]);
         }
       );
 

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/Lineage.spec.ts

@@ -275,6 +275,7 @@ test('Verify column lineage between table and topic', async ({ browser }) => {
 
   await table.visitEntityPage(page);
   await visitLineageTab(page);
+  await activateColumnLayer(page);
   await page.click('[data-testid="edit-lineage"]');
 
   await removeColumnLineage(page, sourceCol, targetCol);
@@ -389,6 +390,8 @@ test('Verify function data in edge drawer', async ({ browser }) => {
     await page.reload();
     await lineageReq;
 
+    await activateColumnLayer(page);
+
     await page
       .locator(
         `[data-testid="column-edge-${btoa(sourceColName)}-${btoa(

openmetadata-ui/src/main/resources/ui/playwright/utils/lineage.ts

@@ -236,6 +236,30 @@ export const performExpand = async (
   }
 };
 
+export const performCollapse = async (
+  page: Page,
+  node: EntityClass,
+  upstream: boolean,
+  hiddenEntity: EntityClass[]
+) => {
+  const nodeFqn = get(node, 'entityResponseData.fullyQualifiedName');
+  const handleDirection = upstream ? 'left' : 'right';
+  const collapseBtn = page
+    .locator(`[data-testid="lineage-node-${nodeFqn}"]`)
+    .locator(`.react-flow__handle-${handleDirection}`)
+    .getByTestId('minus-icon');
+
+  await collapseBtn.click();
+
+  for (const entity of hiddenEntity) {
+    const hiddenNodeFqn = get(entity, 'entityResponseData.fullyQualifiedName');
+    const hiddenNode = page.locator(
+      `[data-testid="lineage-node-${hiddenNodeFqn}"]`
+    );
+
+    await expect(hiddenNode).not.toBeVisible();
+  }
+};
 export const verifyNodePresent = async (page: Page, node: EntityClass) => {
   const nodeFqn = get(node, 'entityResponseData.fullyQualifiedName');
   const name = get(node, 'entityResponseData.name');

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNode.utils.tsx

@@ -208,3 +208,22 @@ export const getColumnContent = (
     </div>
   );
 };
+
+export function getNodeClassNames({
+  isSelected,
+  showDqTracing,
+  isTraced,
+}: {
+  isSelected: boolean;
+  showDqTracing: boolean;
+  isTraced: boolean;
+}) {
+  return classNames(
+    'lineage-node p-0',
+    isSelected ? 'custom-node-header-active' : 'custom-node-header-normal',
+    {
+      'data-quality-failed-custom-node-header': showDqTracing,
+      'custom-node-header-tracing': isTraced,
+    }
+  );
+}

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/CustomNodeV1.component.tsx

@@ -11,7 +11,6 @@
  *  limitations under the License.
  */
 
-import classNames from 'classnames';
 import React, { memo, useCallback, useEffect, useMemo, useState } from 'react';
 import { Handle, NodeProps, Position } from 'reactflow';
 import { useLineageProvider } from '../../../context/LineageProvider/LineageProvider';
@@ -20,7 +19,11 @@ import { LineageDirection } from '../../../generated/api/lineage/lineageDirectio
 import { LineageLayer } from '../../../generated/configuration/lineageSettings';
 import LineageNodeRemoveButton from '../../Lineage/LineageNodeRemoveButton';
 import './custom-node.less';
-import { getCollapseHandle, getExpandHandle } from './CustomNode.utils';
+import {
+  getCollapseHandle,
+  getExpandHandle,
+  getNodeClassNames,
+} from './CustomNode.utils';
 import './entity-lineage.style.less';
 import {
   ExpandCollapseHandlesProps,
@@ -100,7 +103,8 @@ const ExpandCollapseHandles = memo(
     isDownstreamNode,
     isUpstreamNode,
     isRootNode,
-    expandPerformed,
+    upstreamExpandPerformed,
+    downstreamExpandPerformed,
     upstreamLineageLength,
     onCollapse,
     onExpand,
@@ -114,18 +118,21 @@ const ExpandCollapseHandles = memo(
         {hasOutgoers &&
           (isDownstreamNode || isRootNode) &&
           getCollapseHandle(LineageDirection.Downstream, onCollapse)}
+
         {!hasOutgoers &&
-          !expandPerformed &&
+          !downstreamExpandPerformed &&
           getExpandHandle(LineageDirection.Downstream, () =>
             onExpand(LineageDirection.Downstream)
           )}
+
         {hasIncomers &&
           (isUpstreamNode || isRootNode) &&
           getCollapseHandle(LineageDirection.Upstream, () =>
             onCollapse(LineageDirection.Upstream)
           )}
+
         {!hasIncomers &&
-          !expandPerformed &&
+          !upstreamExpandPerformed &&
           upstreamLineageLength > 0 &&
           getExpandHandle(LineageDirection.Upstream, () =>
             onExpand(LineageDirection.Upstream)
@@ -166,17 +173,23 @@ const CustomNodeV1 = (props: NodeProps) => {
     id,
     fullyQualifiedName,
     upstreamLineage = [],
-    expandPerformed = false,
+    upstreamExpandPerformed = false,
+    downstreamExpandPerformed = false,
   } = node;
-  const [isTraced, setIsTraced] = useState<boolean>(false);
+  const [isTraced, setIsTraced] = useState(false);
 
-  const showDqTracing = useMemo(() => {
+  const showDqTracing = useMemo(
-    return (
+    () =>
-      (activeLayer.includes(LineageLayer.DataObservability) &&
+      activeLayer.includes(LineageLayer.DataObservability) &&
-        dataQualityLineage?.nodes?.some((dqNode) => dqNode.id === id)) ??
+      dataQualityLineage?.nodes?.some((dqNode) => dqNode.id === id),
-      false
+    [activeLayer, dataQualityLineage, id]
-    );
+  );
-  }, [activeLayer, dataQualityLineage, id]);
+
+  const containerClass = getNodeClassNames({
+    isSelected,
+    showDqTracing: showDqTracing ?? false,
+    isTraced,
+  });
 
   const onExpand = useCallback(
     (direction: LineageDirection) => {
@@ -197,33 +210,20 @@ const CustomNodeV1 = (props: NodeProps) => {
       return label;
     }
 
-    const renderRemoveBtn =
-      isSelected && isEditMode && !isRootNode ? (
-        <LineageNodeRemoveButton onRemove={() => removeNodeHandler(props)} />
-      ) : null;
-
     return (
       <>
         <LineageNodeLabelV1 node={node} />
-        {renderRemoveBtn}
+        {isSelected && isEditMode && !isRootNode && (
+          <LineageNodeRemoveButton onRemove={() => removeNodeHandler(props)} />
+        )}
       </>
     );
   }, [node.id, isNewNode, label, isSelected, isEditMode, isRootNode]);
 
-  const containerClass = useMemo(() => {
+  const expandCollapseProps = useMemo<ExpandCollapseHandlesProps>(
-    return classNames(
-      'lineage-node p-0',
-      isSelected ? 'custom-node-header-active' : 'custom-node-header-normal',
-      {
-        'data-quality-failed-custom-node-header': showDqTracing,
-        'custom-node-header-tracing': isTraced,
-      }
-    );
-  }, [isSelected, showDqTracing, isTraced]);
-
-  const expandCollapseProps = useMemo(
     () => ({
-      expandPerformed,
+      upstreamExpandPerformed,
+      downstreamExpandPerformed,
       hasIncomers,
       hasOutgoers,
       isDownstreamNode,
@@ -235,7 +235,8 @@ const CustomNodeV1 = (props: NodeProps) => {
       onExpand,
     }),
     [
-      expandPerformed,
+      upstreamExpandPerformed,
+      downstreamExpandPerformed,
       hasIncomers,
       hasOutgoers,
       isDownstreamNode,
@@ -248,6 +249,11 @@ const CustomNodeV1 = (props: NodeProps) => {
     ]
   );
 
+  const handlesElement = useMemo(
+    () => <ExpandCollapseHandles {...expandCollapseProps} />,
+    [expandCollapseProps]
+  );
+
   useEffect(() => {
     setIsTraced(tracedNodes.includes(id));
   }, [tracedNodes, id]);
@@ -257,9 +263,7 @@ const CustomNodeV1 = (props: NodeProps) => {
       className={containerClass}
       data-testid={`lineage-node-${fullyQualifiedName}`}>
       <NodeHandles
-        expandCollapseHandles={
+        expandCollapseHandles={handlesElement}
-          <ExpandCollapseHandles {...expandCollapseProps} />
-        }
         id={id}
         isConnectable={isConnectable}
         nodeType={nodeType}

openmetadata-ui/src/main/resources/ui/src/components/Entity/EntityLineage/EntityLineage.interface.ts

@@ -88,7 +88,8 @@ export interface ExpandCollapseHandlesProps {
   isDownstreamNode: boolean;
   isUpstreamNode: boolean;
   isRootNode: boolean;
-  expandPerformed: boolean;
+  upstreamExpandPerformed: boolean;
+  downstreamExpandPerformed: boolean;
   upstreamLineageLength: number;
   onCollapse: (direction?: LineageDirection) => void;
   onExpand: (direction: LineageDirection) => void;

openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.component.tsx

@@ -21,12 +21,7 @@ import React, {
 } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useHistory } from 'react-router-dom';
-import ReactFlow, {
+import ReactFlow, { Background, MiniMap, Panel } from 'reactflow';
-  Background,
-  MiniMap,
-  Panel,
-  ReactFlowProvider,
-} from 'reactflow';
 import {
   MAX_ZOOM_VALUE,
   MIN_ZOOM_VALUE,
@@ -73,7 +68,6 @@ const Lineage = ({
     onNodeDrop,
     onNodesChange,
     onEdgesChange,
-    entityLineage,
     onPaneClick,
     onConnect,
     onInitReactFlow,
@@ -162,7 +156,7 @@ const Lineage = ({
         data-testid="lineage-container"
         id="lineage-container" // ID is required for export PNG functionality
         ref={reactFlowWrapper}>
-        {entityLineage && (
+        {init ? (
           <>
             {isPlatformLineage ? null : (
               <CustomControlsComponent className="absolute top-1 right-1 p-xs" />
@@ -178,10 +172,6 @@ const Lineage = ({
                 isFullScreen ? onExitFullScreenViewClick : undefined
               }
             />
-          </>
-        )}
-        {init ? (
-          <ReactFlowProvider>
             <ReactFlow
               elevateEdgesOnSelect
               className="custom-react-flow"
@@ -222,7 +212,7 @@ const Lineage = ({
                 <LineageLayers entity={entity} entityType={entityType} />
               </Panel>
             </ReactFlow>
-          </ReactFlowProvider>
+          </>
         ) : (
           <div className="loading-card">
             <Loader />

openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.interface.ts

@@ -81,6 +81,7 @@ export interface LineageEntityReference extends EntityReference {
     parentId?: string;
     childrenLength?: number;
   };
-  expandPerformed?: boolean;
+  upstreamExpandPerformed?: boolean;
+  downstreamExpandPerformed?: boolean;
   direction?: LineageDirection;
 }

openmetadata-ui/src/main/resources/ui/src/components/Lineage/Lineage.test.tsx

@@ -17,7 +17,7 @@ import { MOCK_EXPLORE_SEARCH_RESULTS } from '../Explore/Explore.mock';
 import Lineage from './Lineage.component';
 import { EntityLineageResponse } from './Lineage.interface';
 
-let entityLineage: EntityLineageResponse | undefined = {
+const entityLineage: EntityLineageResponse | undefined = {
   entity: {
     name: 'fact_sale',
     fullyQualifiedName: 'sample_data.ecommerce_db.shopify.fact_sale',
@@ -90,6 +90,7 @@ jest.mock('../../context/LineageProvider/LineageProvider', () => ({
     activeLayer: [],
     entityLineage: entityLineage,
     updateEntityData: jest.fn(),
+    init: true,
   })),
 }));
 
@@ -126,19 +127,7 @@ describe('Lineage', () => {
     const customControlsComponent = screen.getByText('Controls Component');
     const lineageComponent = screen.getByTestId('lineage-container');
 
-    expect(customControlsComponent).toBeInTheDocument();
     expect(lineageComponent).toBeInTheDocument();
-  });
-
-  it('does not render CustomControlsComponent when entityLineage is falsy', () => {
-    const mockPropsWithoutEntity = {
-      ...mockProps,
-      entity: undefined,
-    };
-    entityLineage = undefined;
-    render(<Lineage {...mockPropsWithoutEntity} />);
-    const customControlsComponent = screen.getByText('Controls Component');
-
     expect(customControlsComponent).toBeInTheDocument();
   });
 });

openmetadata-ui/src/main/resources/ui/src/context/LineageProvider/LineageProvider.tsx

@@ -23,6 +23,7 @@ import React, {
   useContext,
   useEffect,
   useMemo,
+  useRef,
   useState,
 } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -55,6 +56,7 @@ import {
   EntityLineageResponse,
   LineageData,
   LineageEntityReference,
+  NodeData,
 } from '../../components/Lineage/Lineage.interface';
 import LineageNodeRemoveButton from '../../components/Lineage/LineageNodeRemoveButton';
 import { SourceType } from '../../components/SearchedData/SearchedData.interface';
@@ -152,6 +154,11 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
     {} as SourceType
   );
   const [activeLayer, setActiveLayer] = useState<LineageLayer[]>([]);
+
+  // Added this ref to compare the previous active layer with the current active layer.
+  // We need to redraw the lineage if the column level lineage is added or removed.
+  const prevActiveLayerRef = useRef<LineageLayer[]>([]);
+
   const [activeNode, setActiveNode] = useState<Node>();
   const [expandAllColumns, setExpandAllColumns] = useState(false);
   const [selectedColumn, setSelectedColumn] = useState<string>('');
@@ -264,6 +271,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         allNodes,
         lineageData.edges ?? [],
         decodedFqn,
+        activeLayer.includes(LineageLayer.ColumnLevelLineage),
         isFirstTime ? true : undefined
       );
 
@@ -365,7 +373,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
 
         setLineageData(res);
 
-        const { nodes, edges, entity } = parseLineageData(res, '');
+        const { nodes, edges, entity } = parseLineageData(res, '', decodedFqn);
         const updatedEntityLineage = {
           nodes,
           edges,
@@ -409,7 +417,7 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         });
         setLineageData(res);
 
-        const { nodes, edges, entity } = parseLineageData(res, fqn);
+        const { nodes, edges, entity } = parseLineageData(res, fqn, decodedFqn);
         const updatedEntityLineage = {
           nodes,
           edges,
@@ -510,9 +518,19 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
           direction,
         });
 
+        const currentNodes: Record<string, NodeData> = {};
+        entityLineage.nodes?.forEach((node) => {
+          currentNodes[node.fullyQualifiedName ?? ''] = {
+            entity: node,
+            paging: (node as LineageEntityReference).paging ?? {
+              entityDownstreamCount: 0,
+              entityUpstreamCount: 0,
+            },
+          };
+        });
         const concatenatedLineageData = {
           nodes: {
-            ...lineageData?.nodes,
+            ...currentNodes,
             ...res.nodes,
           },
           downstreamEdges: {
@@ -527,20 +545,56 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
 
         const { nodes: newNodes, edges: newEdges } = parseLineageData(
           concatenatedLineageData,
-          node.fullyQualifiedName ?? ''
+          node.fullyQualifiedName ?? '',
+          decodedFqn
         );
+
+        const uniqueNodes = [...(entityLineage.nodes ?? [])];
+        for (const nNode of newNodes ?? []) {
+          if (
+            !uniqueNodes.some(
+              (n) => n.fullyQualifiedName === nNode.fullyQualifiedName
+            )
+          ) {
+            uniqueNodes.push(nNode);
+          }
+        }
+
         const updatedEntityLineage = {
           entity: entityLineage.entity,
-          nodes: uniqWith(
+          nodes: uniqueNodes,
-            [...(entityLineage.nodes ?? []), ...newNodes],
-            isEqual
-          ),
           edges: uniqWith(
             [...(entityLineage.edges ?? []), ...newEdges],
             isEqual
           ),
         };
 
+        // remove the nodes and edges from the lineageData
+        const visibleNodes: Record<string, NodeData> = {};
+        uniqueNodes.forEach((node: EntityReference) => {
+          visibleNodes[node.fullyQualifiedName ?? ''] = {
+            entity: node,
+            paging: (node as LineageEntityReference).paging ?? {
+              entityDownstreamCount: 0,
+              entityUpstreamCount: 0,
+            },
+          };
+        });
+
+        const currentNode = updatedEntityLineage.nodes.find(
+          (n) => n.fullyQualifiedName === node.fullyQualifiedName
+        );
+
+        if (currentNode) {
+          if (direction === LineageDirection.Upstream) {
+            (currentNode as LineageEntityReference).upstreamExpandPerformed =
+              true;
+          } else {
+            (currentNode as LineageEntityReference).downstreamExpandPerformed =
+              true;
+          }
+        }
+
         updateLineageData(updatedEntityLineage, {
           shouldRedraw: true,
           centerNode: false,
@@ -887,7 +941,8 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
 
       const { nodes: newNodes, edges: newEdges } = parseLineageData(
         concatenatedLineageData,
-        parentNode.data.node.fullyQualifiedName
+        parentNode.data.node.fullyQualifiedName,
+        decodedFqn
       );
 
       updateLineageData(
@@ -1100,7 +1155,12 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
             );
 
             const { edges: createdEdges, columnsHavingLineage } =
-              createEdgesAndEdgeMaps(allNodes, allEdges, decodedFqn);
+              createEdgesAndEdgeMaps(
+                allNodes,
+                allEdges,
+                decodedFqn,
+                activeLayer.includes(LineageLayer.ColumnLevelLineage)
+              );
             setEdges(createdEdges);
             setColumnsHavingLineage(columnsHavingLineage);
 
@@ -1328,6 +1388,31 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
         );
       });
 
+      // Find the node in updatedNodes by ID and set expandPerformed: false
+      const currentNodeId = (node as Node).id;
+      const nodeToUpdate = updatedNodes.find((n) => n.id === currentNodeId);
+      if (nodeToUpdate) {
+        if (direction === LineageDirection.Upstream) {
+          (nodeToUpdate as LineageEntityReference).upstreamExpandPerformed =
+            false;
+        } else {
+          (nodeToUpdate as LineageEntityReference).downstreamExpandPerformed =
+            false;
+        }
+      }
+
+      // remove the nodes and edges from the lineageData
+      const visibleNodes: Record<string, NodeData> = {};
+      updatedNodes.forEach((node) => {
+        visibleNodes[node.fullyQualifiedName ?? ''] = {
+          entity: node,
+          paging: (node as LineageEntityReference).paging ?? {
+            entityDownstreamCount: 0,
+            entityUpstreamCount: 0,
+          },
+        };
+      });
+
       updateLineageData(
         {
           ...entityLineage,
@@ -1503,7 +1588,20 @@ const LineageProvider = ({ children }: LineageProviderProps) => {
   }, [isEditMode, deletePressed, backspacePressed, activeNode, selectedEdge]);
 
   useEffect(() => {
-    repositionLayout();
+    const prevActiveLayer = prevActiveLayerRef.current;
+
+    const prevHadColumn = prevActiveLayer.includes(
+      LineageLayer.ColumnLevelLineage
+    );
+    const currHasColumn = activeLayer.includes(LineageLayer.ColumnLevelLineage);
+
+    if (prevHadColumn !== currHasColumn) {
+      redraw();
+    } else {
+      repositionLayout();
+    }
+
+    prevActiveLayerRef.current = activeLayer;
   }, [activeLayer, expandAllColumns]);
 
   useEffect(() => {

openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx

@@ -16,7 +16,15 @@ import { graphlib, layout } from '@dagrejs/dagre';
 import { AxiosError } from 'axios';
 import ELK, { ElkExtendedEdge, ElkNode } from 'elkjs/lib/elk.bundled.js';
 import { t } from 'i18next';
-import { get, isEmpty, isNil, isUndefined, uniqueId } from 'lodash';
+import {
+  get,
+  isEmpty,
+  isEqual,
+  isNil,
+  isUndefined,
+  uniqueId,
+  uniqWith,
+} from 'lodash';
 import { EntityTags, LoadingState } from 'Models';
 import React, { MouseEvent as ReactMouseEvent } from 'react';
 import {
@@ -839,6 +847,7 @@ export const createEdgesAndEdgeMaps = (
   nodes: EntityReference[],
   edges: EdgeDetails[],
   entityFqn: string,
+  isColumnLayerActive: boolean,
   hidden?: boolean
 ) => {
   const lineageEdgesV1: Edge[] = [];
@@ -851,10 +860,6 @@ export const createEdgesAndEdgeMaps = (
     const sourceId = edge.fromEntity.id;
     const targetId = edge.toEntity.id;
 
-    // Update edge maps for fast lookup
-    outgoingMap.set(sourceId, (outgoingMap.get(sourceId) ?? 0) + 1);
-    incomingMap.set(targetId, (incomingMap.get(targetId) ?? 0) + 1);
-
     const sourceType = nodes.find((n) => sourceId === n.id);
     const targetType = nodes.find((n) => targetId === n.id);
 
@@ -862,7 +867,11 @@ export const createEdgesAndEdgeMaps = (
       return;
     }
 
-    if (!isUndefined(edge.columns)) {
+    // Update edge maps for fast lookup
+    outgoingMap.set(sourceId, (outgoingMap.get(sourceId) ?? 0) + 1);
+    incomingMap.set(targetId, (incomingMap.get(targetId) ?? 0) + 1);
+
+    if (!isUndefined(edge.columns) && isColumnLayerActive) {
       edge.columns?.forEach((e) => {
         const toColumn = e.toColumn ?? '';
         if (toColumn && e.fromColumns?.length) {
@@ -1133,7 +1142,7 @@ export const getConnectedNodesEdges = (
 
   return {
     nodes: outgoers,
-    edges: connectedEdges,
+    edges: uniqWith(connectedEdges, isEqual),
     nodeFqn: childNodeFqn,
   };
 };
@@ -1406,9 +1415,16 @@ const processNodeArray = (
     .map((node: NodeData) => ({
       ...node.entity,
       paging: node.paging,
-      expandPerformed:
+      upstreamExpandPerformed:
-        (node.entity as LineageEntityReference).expandPerformed ||
+        (node.entity as LineageEntityReference).upstreamExpandPerformed !==
-        node.entity.fullyQualifiedName === entityFqn,
+        undefined
+          ? (node.entity as LineageEntityReference).upstreamExpandPerformed
+          : node.entity.fullyQualifiedName === entityFqn,
+      downstreamExpandPerformed:
+        (node.entity as LineageEntityReference).downstreamExpandPerformed !==
+        undefined
+          ? (node.entity as LineageEntityReference).downstreamExpandPerformed
+          : node.entity.fullyQualifiedName === entityFqn,
     }))
     .flat();
 };
@@ -1509,7 +1525,8 @@ const processPagination = (
 
 export const parseLineageData = (
   data: LineageData,
-  entityFqn: string
+  entityFqn: string, // This contains fqn of node or entity that is being viewed in lineage page
+  rootFqn: string // This contains the fqn of the entity that is being viewed in lineage page
 ): {
   nodes: LineageEntityReference[];
   edges: EdgeDetails[];
@@ -1518,7 +1535,7 @@ export const parseLineageData = (
   const { nodes, downstreamEdges, upstreamEdges } = data;
 
   // Process nodes
-  const nodesArray = processNodeArray(nodes, entityFqn);
+  const nodesArray = processNodeArray(nodes, rootFqn);
   const processedNodes: LineageEntityReference[] = [...nodesArray];
 
   // Process edges