diff --git a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
index 4939be343fc..3545dfd213a 100644
--- a/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/components/EntityLineage/EntityLineage.component.tsx
@@ -197,6 +197,7 @@ const EntityLineageComponent: FunctionComponent<EntityLineageProp> = ({
 
   const [nodes, setNodes, onNodesChange] = useNodesState([]);
   const [edges, setEdges, onEdgesChange] = useEdgesState([]);
+  const edgesRef = useRef<Edge[]>([]);
   const [paginationData, setPaginationData] = useState({});
   const [entityLineage, setEntityLineage] = useState<EntityLineage>();
   const [updatedLineageData, setUpdatedLineageData] = useState<EntityLineage>();
@@ -368,11 +369,7 @@ const EntityLineageComponent: FunctionComponent<EntityLineageProp> = ({
     if (node.type === EntityLineageNodeType.LOAD_MORE) {
       selectLoadMoreNode(node);
     } else {
-      const selectedNode = [
-        ...(updatedLineageData?.nodes || []),
-        updatedLineageData?.entity,
-      ].find((n) => n && node.id.includes(n.id));
-
+      const selectedNode = node.data.node;
       if (!expandButton.current) {
         selectNodeHandler(true, {
           name: selectedNode?.name as string,
@@ -697,7 +694,10 @@ const EntityLineageComponent: FunctionComponent<EntityLineageProp> = ({
                   ...col,
                   type: isEditMode
                     ? EntityLineageNodeType.DEFAULT
-                    : getColumnType(edges, col.fullyQualifiedName || col.name),
+                    : getColumnType(
+                        edgesRef.current,
+                        col.fullyQualifiedName || col.name
+                      ),
                 };
               });
               prevNode.data.columns = cols;
@@ -1213,7 +1213,10 @@ const EntityLineageComponent: FunctionComponent<EntityLineageProp> = ({
               ...col,
               type: isEditMode
                 ? 'default'
-                : getColumnType(edges, col.fullyQualifiedName || col.name),
+                : getColumnType(
+                    edgesRef.current,
+                    col.fullyQualifiedName || col.name
+                  ),
             };
           });
           node.data.columns = cols;
@@ -1640,6 +1643,10 @@ const EntityLineageComponent: FunctionComponent<EntityLineageProp> = ({
     }
   }, [pipelineSearchValue]);
 
+  useEffect(() => {
+    edgesRef.current = edges;
+  }, [edges]);
+
   if (isLineageLoading || (nodes.length === 0 && !deleted)) {
     return <Loader />;
   }
diff --git a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
index 68aa9a9dd05..b90c525423f 100644
--- a/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
+++ b/openmetadata-ui/src/main/resources/ui/src/utils/EntityLineageUtils.tsx
@@ -257,16 +257,7 @@ export const getLineageData = (
       edge.lineageDetails.columnsLineage?.forEach((e) => {
         const toColumn = e.toColumn || '';
         if (toColumn && e.fromColumns && e.fromColumns.length > 0) {
-          const targetCol = nodes.find((n) => toColumn === n.id);
-          if (!targetCol) {
-            return;
-          }
-
           e.fromColumns.forEach((fromColumn) => {
-            const sourceCol = nodes.find((n) => fromColumn === n.id);
-            if (!sourceCol) {
-              return;
-            }
             lineageEdgesV1.push({
               id: `column-${fromColumn}-${toColumn}-edge-${edge.fromEntity}-${edge.toEntity}`,
               source: edge.fromEntity,