feat(ui) Display inputFields in column-lineage visualization (#6303)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/mappers/InputFieldsMapper.java

@@ -24,6 +24,9 @@ public class InputFieldsMapper {
             if (field.hasSchemaField()) {
                 fieldResult.setSchemaField(SchemaFieldMapper.map(field.getSchemaField(), entityUrn));
             }
+            if (field.hasSchemaFieldUrn()) {
+                fieldResult.setSchemaFieldUrn(field.getSchemaFieldUrn().toString());
+            }
             return fieldResult;
         }).collect(Collectors.toList()));
 

datahub-web-react/src/app/entity/EntityRegistry.tsx

@@ -142,6 +142,7 @@ export default class EntityRegistry {
                 siblingPlatforms: genericEntityProperties?.siblingPlatforms,
                 fineGrainedLineages: genericEntityProperties?.fineGrainedLineages,
                 schemaMetadata: genericEntityProperties?.schemaMetadata,
+                inputFields: genericEntityProperties?.inputFields,
             } as FetchedEntity) || undefined
         );
     }

datahub-web-react/src/app/lineage/LineageEntityColumns.tsx

@@ -10,7 +10,11 @@ import { centerY, EXPAND_COLLAPSE_COLUMNS_TOGGLE_HEIGHT, iconX, NUM_COLUMNS_PER_
 import ColumnNode from './ColumnNode';
 import NodeColumnsHeader from './NodeColumnsHeader';
 import usePrevious from '../shared/usePrevious';
-import { filterColumns, haveDisplayedFieldsChanged } from './utils/columnLineageUtils';
+import {
+    convertInputFieldsToSchemaFields,
+    filterColumns,
+    haveDisplayedFieldsChanged,
+} from './utils/columnLineageUtils';
 import { useResetPageIndexAfterSelect } from './utils/useResetPageIndexAfterSelect';
 
 const StyledPagination = styled(Pagination)`
@@ -32,7 +36,10 @@ export default function LineageEntityColumns({ node, onHover }: Props) {
 
     const titleHeight = getTitleHeight(expandTitles ? node.data.expandedName || node.data.name : undefined);
 
-    const fields = columnsByUrn[node.data.urn || ''] || node.data.schemaMetadata?.fields;
+    const fields =
+        columnsByUrn[node.data.urn || ''] ||
+        node.data.schemaMetadata?.fields ||
+        convertInputFieldsToSchemaFields(node.data.inputFields);
 
     const displayedFields = fields?.slice(
         pageIndex * NUM_COLUMNS_PER_PAGE,
@@ -60,7 +67,8 @@ export default function LineageEntityColumns({ node, onHover }: Props) {
     }, [displayedFields, node?.data?.urn, setVisibleColumnsByUrn, previousDisplayedFields]);
 
     const hasColumnPagination =
-        node.data.schemaMetadata?.fields && node.data.schemaMetadata?.fields.length > NUM_COLUMNS_PER_PAGE;
+        (node.data.schemaMetadata?.fields && node.data.schemaMetadata?.fields.length > NUM_COLUMNS_PER_PAGE) ||
+        (node.data.inputFields?.fields && node.data.inputFields.fields.length > NUM_COLUMNS_PER_PAGE);
 
     return (
         <>

datahub-web-react/src/app/lineage/LineageEntityNode.tsx

@@ -14,6 +14,7 @@ import { useGetEntityLineageLazyQuery } from '../../graphql/lineage.generated';
 import { useIsSeparateSiblingsMode } from '../entity/shared/siblingUtils';
 import { centerX, centerY, iconHeight, iconWidth, iconX, iconY, textX, width } from './constants';
 import LineageEntityColumns from './LineageEntityColumns';
+import { convertInputFieldsToSchemaFields } from './utils/columnLineageUtils';
 
 const CLICK_DELAY_THRESHOLD = 1000;
 const DRAG_DISTANCE_THRESHOLD = 20;
@@ -100,7 +101,7 @@ export default function LineageEntityNode({
 
     const nodeHeight = nodeHeightFromTitleLength(
         expandTitles ? node.data.expandedName || node.data.name : undefined,
-        node.data.schemaMetadata,
+        node.data.schemaMetadata?.fields || convertInputFieldsToSchemaFields(node.data.inputFields),
         showColumns,
         areColumnsCollapsed,
     );
@@ -320,7 +321,9 @@ export default function LineageEntityNode({
                         {unexploredHiddenChildren > 1 ? 'dependencies' : 'dependency'}
                     </UnselectableText>
                 ) : null}
-                {showColumns && node.data.schemaMetadata && <LineageEntityColumns node={node} onHover={onHover} />}
+                {showColumns && (node.data.schemaMetadata || node.data.inputFields) && (
+                    <LineageEntityColumns node={node} onHover={onHover} />
+                )}
             </Group>
         </PointerGroup>
     );

datahub-web-react/src/app/lineage/types.ts

@@ -15,6 +15,7 @@ import {
     DataPlatform,
     FineGrainedLineage,
     SchemaMetadata,
+    InputFields,
 } from '../../types.generated';
 
 export type EntitySelectParams = {
@@ -47,6 +48,7 @@ export type FetchedEntity = {
     siblingPlatforms?: Maybe<DataPlatform[]>;
     fineGrainedLineages?: [FineGrainedLineage];
     schemaMetadata?: SchemaMetadata;
+    inputFields?: InputFields;
 };
 
 export type NodeData = {
@@ -66,6 +68,7 @@ export type NodeData = {
     status?: Maybe<Status>;
     siblingPlatforms?: Maybe<DataPlatform[]>;
     schemaMetadata?: SchemaMetadata;
+    inputFields?: InputFields;
 };
 
 export type VizNode = {

datahub-web-react/src/app/lineage/utils/__tests__/columnLineageUtils.test.tsx

@@ -0,0 +1,45 @@
+import { getFieldPathFromSchemaFieldUrn, getSourceUrnFromSchemaFieldUrn } from '../columnLineageUtils';
+
+describe('getSourceUrnFromSchemaFieldUrn', () => {
+    it('should get the source urn for a chart schemaField', () => {
+        const schemaFieldUrn = 'urn:li:schemaField:(urn:li:chart:(looker,dashboard_elements.1),goal)';
+        const sourceUrn = getSourceUrnFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('urn:li:chart:(looker,dashboard_elements.1)');
+    });
+
+    it('should get the source urn for a dataset schemaField', () => {
+        const schemaFieldUrn =
+            'urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD),user_name)';
+        const sourceUrn = getSourceUrnFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD)');
+    });
+
+    it('should get the source urn for a nested schemaField', () => {
+        const schemaFieldUrn =
+            'urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD),user.name.test)';
+        const sourceUrn = getSourceUrnFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD)');
+    });
+});
+
+describe('getFieldPathFromSchemaFieldUrn', () => {
+    it('should get the fieldPath from a chart schemaField urn', () => {
+        const schemaFieldUrn = 'urn:li:schemaField:(urn:li:chart:(looker,dashboard_elements.1),goal)';
+        const sourceUrn = getFieldPathFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('goal');
+    });
+
+    it('should get the fieldPath for a dataset schemaField', () => {
+        const schemaFieldUrn =
+            'urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD),user_name)';
+        const sourceUrn = getFieldPathFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('user_name');
+    });
+
+    it('should get the fieldPath for a nested schemaField', () => {
+        const schemaFieldUrn =
+            'urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,fct_cypress_users_created,PROD),user.name.test)';
+        const sourceUrn = getFieldPathFromSchemaFieldUrn(schemaFieldUrn);
+        expect(sourceUrn).toBe('user.name.test');
+    });
+});

datahub-web-react/src/app/lineage/utils/columnLineageUtils.ts

@@ -1,5 +1,5 @@
 import { ColumnEdge, FetchedEntity, NodeData } from '../types';
-import { SchemaField } from '../../../types.generated';
+import { InputFields, SchemaField } from '../../../types.generated';
 import { downgradeV2FieldPath } from '../../entity/dataset/profile/schema/utils/utils';
 
 export function getHighlightedColumnsForNode(highlightedEdges: ColumnEdge[], fields: SchemaField[], nodeUrn: string) {
@@ -59,6 +59,10 @@ export function sortColumnsByDefault(
     };
 }
 
+export function convertInputFieldsToSchemaFields(inputFields?: InputFields) {
+    return inputFields?.fields?.map((field) => field?.schemaField) as SchemaField[] | undefined;
+}
+
 export function populateColumnsByUrn(
     columnsByUrn: Record<string, SchemaField[]>,
     fetchedEntities: { [x: string]: FetchedEntity },
@@ -71,6 +75,13 @@ export function populateColumnsByUrn(
                 ...populatedColumnsByUrn,
                 [urn]: convertFieldsToV1FieldPath(fetchedEntity.schemaMetadata.fields),
             };
+        } else if (fetchedEntity.inputFields?.fields && !columnsByUrn[urn]) {
+            populatedColumnsByUrn = {
+                ...populatedColumnsByUrn,
+                [urn]: convertFieldsToV1FieldPath(
+                    convertInputFieldsToSchemaFields(fetchedEntity.inputFields) as SchemaField[],
+                ),
+            };
         }
     });
     setColumnsByUrn(populatedColumnsByUrn);
@@ -104,3 +115,10 @@ export function filterColumns(
         }));
     }
 }
+
+export function getSourceUrnFromSchemaFieldUrn(schemaFieldUrn: string) {
+    return schemaFieldUrn.replace('urn:li:schemaField:(', '').split(')')[0].concat(')');
+}
+export function getFieldPathFromSchemaFieldUrn(schemaFieldUrn: string) {
+    return schemaFieldUrn.replace('urn:li:schemaField:(', '').split(')')[1].replace(',', '');
+}

datahub-web-react/src/app/lineage/utils/constructFetchedNode.ts

@@ -63,6 +63,7 @@ export default function constructFetchedNode(
             status: fetchedNode.status,
             siblingPlatforms: fetchedNode.siblingPlatforms,
             schemaMetadata: fetchedNode.schemaMetadata,
+            inputFields: fetchedNode.inputFields,
         };
 
         // eslint-disable-next-line no-param-reassign

datahub-web-react/src/app/lineage/utils/constructTree.ts

@@ -24,6 +24,7 @@ export default function constructTree(
         unexploredChildren: 0,
         siblingPlatforms: fetchedEntity?.siblingPlatforms,
         schemaMetadata: fetchedEntity?.schemaMetadata,
+        inputFields: fetchedEntity?.inputFields,
     };
     const lineageConfig = entityRegistry.getLineageVizConfig(entityAndType.type, entityAndType.entity);
     let children: EntityAndType[] = [];

datahub-web-react/src/app/lineage/utils/extendAsyncEntities.ts

@@ -1,6 +1,7 @@
 import { SchemaFieldRef } from '../../../types.generated';
 import EntityRegistry from '../../entity/EntityRegistry';
 import { EntityAndType, FetchedEntities, FetchedEntity } from '../types';
+import { getFieldPathFromSchemaFieldUrn, getSourceUrnFromSchemaFieldUrn } from './columnLineageUtils';
 
 const breakFieldUrn = (ref: SchemaFieldRef) => {
     const before = ref.urn;
@@ -55,6 +56,22 @@ function extendColumnLineage(lineageVizConfig: FetchedEntity, fineGrainedMap: an
             });
         });
     }
+    if (lineageVizConfig.inputFields?.fields && lineageVizConfig.inputFields.fields.length > 0) {
+        lineageVizConfig.inputFields.fields.forEach((inputField) => {
+            if (inputField?.schemaFieldUrn && inputField.schemaField) {
+                const sourceUrn = getSourceUrnFromSchemaFieldUrn(inputField.schemaFieldUrn);
+                if (sourceUrn !== lineageVizConfig.urn) {
+                    updateFineGrainedMap(
+                        fineGrainedMap,
+                        sourceUrn,
+                        getFieldPathFromSchemaFieldUrn(inputField.schemaFieldUrn),
+                        lineageVizConfig.urn,
+                        inputField.schemaField.fieldPath,
+                    );
+                }
+            }
+        });
+    }
 }
 
 export default function extendAsyncEntities(

datahub-web-react/src/app/lineage/utils/layoutTree.ts

@@ -9,6 +9,7 @@ import {
     width as nodeWidth,
 } from '../constants';
 import { Direction, NodeData, VizEdge, VizNode } from '../types';
+import { convertInputFieldsToSchemaFields } from './columnLineageUtils';
 import { getTitleHeight, nodeHeightFromTitleLength } from './titleUtils';
 
 type ProcessArray = {
@@ -95,7 +96,7 @@ function layoutNodesForOneDirection(
                 currentXPosition +=
                     nodeHeightFromTitleLength(
                         expandTitles ? node.expandedName || node.name : undefined,
-                        node.schemaMetadata,
+                        node.schemaMetadata?.fields || convertInputFieldsToSchemaFields(node.inputFields),
                         showColumns,
                         !!collapsedColumnsNodes[node?.urn || 'no-op'], // avoid indexing on undefined if node is undefined
                     ) + VERTICAL_SPACE_BETWEEN_NODES;

datahub-web-react/src/app/lineage/utils/titleUtils.ts

@@ -1,4 +1,4 @@
-import { SchemaMetadata } from '../../../types.generated';
+import { SchemaField } from '../../../types.generated';
 import { COLUMN_HEIGHT, EXPAND_COLLAPSE_COLUMNS_TOGGLE_HEIGHT, NUM_COLUMNS_PER_PAGE } from '../constants';
 
 interface OptionalOptions {
@@ -93,18 +93,17 @@ export function getTitleHeight(title?: string) {
 
 export function nodeHeightFromTitleLength(
     title?: string,
-    schemaMetadata?: SchemaMetadata,
+    fields?: SchemaField[],
     showColumns?: boolean,
     collapsed?: boolean,
 ) {
     let showColumnBuffer = 0;
     let columnPaginationBuffer = 0;
-    if (showColumns && schemaMetadata) {
+    if (showColumns && fields) {
         if (!collapsed) {
             showColumnBuffer =
-                Math.min(schemaMetadata.fields.length, NUM_COLUMNS_PER_PAGE) * COLUMN_HEIGHT +
+                Math.min(fields.length, NUM_COLUMNS_PER_PAGE) * COLUMN_HEIGHT + EXPAND_COLLAPSE_COLUMNS_TOGGLE_HEIGHT;
-                EXPAND_COLLAPSE_COLUMNS_TOGGLE_HEIGHT;
+            if (fields.length > NUM_COLUMNS_PER_PAGE) {
-            if (schemaMetadata.fields.length > NUM_COLUMNS_PER_PAGE) {
                 columnPaginationBuffer = 40;
             }
         } else {

datahub-web-react/src/app/lineage/utils/useSortColumnsBySelectedField.ts

@@ -1,9 +1,11 @@
 import { useContext, useEffect } from 'react';
+import { SchemaField } from '../../../types.generated';
 import usePrevious from '../../shared/usePrevious';
 import { NUM_COLUMNS_PER_PAGE } from '../constants';
 import { FetchedEntity } from '../types';
 import {
     convertFieldsToV1FieldPath,
+    convertInputFieldsToSchemaFields,
     getHighlightedColumnsForNode,
     sortColumnsByDefault,
     sortRelatedLineageColumns,
@@ -19,7 +21,7 @@ export default function useSortColumnsBySelectedField(fetchedEntities: { [x: str
 
         if (selectedField && previousSelectedField !== selectedField) {
             Object.entries(columnsByUrn).forEach(([urn, columns]) => {
-                if (selectedField.urn !== urn && columns.length >= NUM_COLUMNS_PER_PAGE) {
+                if (selectedField.urn !== urn && columns.length > NUM_COLUMNS_PER_PAGE) {
                     const highlightedColumnsForNode = getHighlightedColumnsForNode(highlightedEdges, columns, urn);
 
                     if (highlightedColumnsForNode.length > 0) {
@@ -43,6 +45,15 @@ export default function useSortColumnsBySelectedField(fetchedEntities: { [x: str
                         convertFieldsToV1FieldPath(fetchedEntity.schemaMetadata.fields),
                         urn,
                     );
+                } else if (fetchedEntity && fetchedEntity.inputFields) {
+                    updatedColumnsByUrn = sortColumnsByDefault(
+                        updatedColumnsByUrn,
+                        columns,
+                        convertFieldsToV1FieldPath(
+                            convertInputFieldsToSchemaFields(fetchedEntity.inputFields) as SchemaField[],
+                        ),
+                        urn,
+                    );
                 }
             });
             setColumnsByUrn(updatedColumnsByUrn);

datahub-web-react/src/graphql/lineage.graphql

@@ -259,6 +259,11 @@ fragment fullLineageResults on EntityLineageResult {
                     ...schemaMetadataFields
                 }
             }
+            ... on Chart {
+                inputFields @include(if: $showColumns) {
+                    ...inputFieldsFields
+                }
+            }
         }
     }
 }
@@ -300,6 +305,11 @@ query getEntityLineage($urn: String!, $separateSiblings: Boolean, $showColumns:
                 }
             }
         }
+        ... on Chart {
+            inputFields @include(if: $showColumns) {
+                ...inputFieldsFields
+            }
+        }
         ... on EntityWithRelationships {
             upstream: lineage(
                 input: { direction: UPSTREAM, start: 0, count: 100, separateSiblings: $separateSiblings }