diff --git a/catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/DatasetLineage.tsx b/catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/EntityLineage.tsx
similarity index 88%
rename from catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/DatasetLineage.tsx
rename to catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/EntityLineage.tsx
index 11ae6a4fb0f..1414ba90eaa 100644
--- a/catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/DatasetLineage.tsx
+++ b/catalog-rest-service/src/main/resources/ui/src/components/dataset-lineage/EntityLineage.tsx
@@ -10,6 +10,7 @@ import ReactFlow, {
   Controls,
   Edge,
   Elements,
+  FlowElement,
   Node,
   OnLoadParams,
   Position,
@@ -45,7 +46,11 @@ const onNodeContextMenu = (_event: ReactMouseEvent, _node: Node | Edge) => {
 const getDataLabel = (v = '', separator = '.') => {
   const length = v.split(separator).length;
 
-  return v.split(separator)[length - 1];
+  return (
+    <p className="tw-break-words description-text">
+      {v.split(separator)[length - 1]}
+    </p>
+  );
 };
 
 const positionX = 150;
@@ -71,11 +76,11 @@ const getLineageData = (entityLineage: EntityLineage) => {
   const getNodes = (
     id: string,
     pos: 'from' | 'to',
-    depth: number
+    depth: number,
+    NodesArr: Array<EntityReference & { lDepth: number }> = []
   ): Array<EntityReference & { lDepth: number }> => {
-    let upDepth = 0;
-    let downDepth = 0;
     if (pos === 'to') {
+      let upDepth = NodesArr.filter((nd) => nd.lDepth === depth).length;
       const UPNodes: Array<EntityReference> = [];
       const updatedUpStreamEdge = upstreamEdges.map((up) => {
         if (up.toEntity === id) {
@@ -118,6 +123,7 @@ const getLineageData = (entityLineage: EntityLineage) => {
 
       return UPNodes?.map((upNd) => ({ lDepth: depth, ...upNd })) || [];
     } else {
+      let downDepth = NodesArr.filter((nd) => nd.lDepth === depth).length;
       const DOWNNodes: Array<EntityReference> = [];
       const updatedDownStreamEdge = downstreamEdges.map((down) => {
         if (down.fromEntity === id) {
@@ -169,7 +175,7 @@ const getLineageData = (entityLineage: EntityLineage) => {
     depth = 1,
     upNodesArr: Array<EntityReference & { lDepth: number }> = []
   ) => {
-    const upNodes = getNodes(Entity.id, 'to', depth);
+    const upNodes = getNodes(Entity.id, 'to', depth, upNodesArr);
     upNodesArr.push(...upNodes);
     upNodes.forEach((up) => {
       if (
@@ -187,7 +193,7 @@ const getLineageData = (entityLineage: EntityLineage) => {
     depth = 1,
     downNodesArr: Array<EntityReference & { lDepth: number }> = []
   ) => {
-    const downNodes = getNodes(Entity.id, 'from', depth);
+    const downNodes = getNodes(Entity.id, 'from', depth, downNodesArr);
     downNodesArr.push(...downNodes);
     downNodes.forEach((down) => {
       if (
@@ -216,19 +222,27 @@ const getLineageData = (entityLineage: EntityLineage) => {
       data: { label: getDataLabel(mainNode.name as string) },
       position: { x: x, y: y },
     },
-    ...UPStreamNodes,
-    ...DOWNStreamNodes,
+    ...UPStreamNodes.map((up) => {
+      return lineageEdges.find(
+        (ed: FlowElement) => (ed as Edge).target === up.id
+      )
+        ? up
+        : { ...up, type: 'input' };
+    }),
+    ...DOWNStreamNodes.map((down) => {
+      return lineageEdges.find((ed: FlowElement) =>
+        (ed as Edge).source.includes(down.id)
+      )
+        ? down
+        : { ...down, type: 'output' };
+    }),
     ...lineageEdges,
   ];
 
   return lineageData;
 };
 
-const DatasetLineage = ({
-  entityLineage,
-}: {
-  entityLineage: EntityLineage;
-}) => {
+const Entitylineage = ({ entityLineage }: { entityLineage: EntityLineage }) => {
   const [elements, setElements] = useState<Elements>(
     getLineageData(entityLineage) as Elements
   );
@@ -269,4 +283,4 @@ const DatasetLineage = ({
   );
 };
 
-export default DatasetLineage;
+export default Entitylineage;
diff --git a/catalog-rest-service/src/main/resources/ui/src/pages/Pipeline-details/index.tsx b/catalog-rest-service/src/main/resources/ui/src/pages/Pipeline-details/index.tsx
index 84e0ff39cdd..0a1efff7451 100644
--- a/catalog-rest-service/src/main/resources/ui/src/pages/Pipeline-details/index.tsx
+++ b/catalog-rest-service/src/main/resources/ui/src/pages/Pipeline-details/index.tsx
@@ -21,7 +21,7 @@ import RichTextEditorPreviewer from '../../components/common/rich-text-editor/Ri
 import TabsPane from '../../components/common/TabsPane/TabsPane';
 import { TitleBreadcrumbProps } from '../../components/common/title-breadcrumb/title-breadcrumb.interface';
 import PageContainer from '../../components/containers/PageContainer';
-import DatasetLineage from '../../components/dataset-lineage/DatasetLineage';
+import Entitylineage from '../../components/dataset-lineage/EntityLineage';
 import Loader from '../../components/Loader/Loader';
 import { ModalWithMarkdownEditor } from '../../components/Modals/ModalWithMarkdownEditor/ModalWithMarkdownEditor';
 import ManageTab from '../../components/my-data-details/ManageTab';
@@ -619,7 +619,7 @@ const MyPipelinePage = () => {
               )}
               {activeTab === 2 && (
                 <div className="tw-w-full tw-h-96">
-                  <DatasetLineage entityLineage={entityLineage} />
+                  <Entitylineage entityLineage={entityLineage} />
                 </div>
               )}
               {activeTab === 3 && (
diff --git a/catalog-rest-service/src/main/resources/ui/src/pages/my-data-details/index.tsx b/catalog-rest-service/src/main/resources/ui/src/pages/my-data-details/index.tsx
index ba363f94e07..873d833706a 100644
--- a/catalog-rest-service/src/main/resources/ui/src/pages/my-data-details/index.tsx
+++ b/catalog-rest-service/src/main/resources/ui/src/pages/my-data-details/index.tsx
@@ -38,7 +38,7 @@ import EntityPageInfo from '../../components/common/entityPageInfo/EntityPageInf
 import TabsPane from '../../components/common/TabsPane/TabsPane';
 import { TitleBreadcrumbProps } from '../../components/common/title-breadcrumb/title-breadcrumb.interface';
 import PageContainer from '../../components/containers/PageContainer';
-import DatasetLineage from '../../components/dataset-lineage/DatasetLineage';
+import Entitylineage from '../../components/dataset-lineage/EntityLineage';
 import FrequentlyJoinedTables from '../../components/my-data-details/FrequentlyJoinedTables';
 import ManageTab from '../../components/my-data-details/ManageTab';
 import ProfilerTable from '../../components/my-data-details/ProfilerTable';
@@ -562,7 +562,7 @@ const MyDataDetailsPage = () => {
             )}
             {activeTab === 3 && (
               <div className="tw-w-full tw-h-96">
-                <DatasetLineage entityLineage={entityLineage} />
+                <Entitylineage entityLineage={entityLineage} />
               </div>
             )}
             {activeTab === 4 && (