feat(lineage) Update Lineage tab and Impact Analysis feature (#5121)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -6,6 +6,7 @@ import com.datahub.authentication.user.NativeUserService;
 import com.datahub.authorization.AuthorizationConfiguration;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.common.VersionedUrn;
+import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.analytics.resolver.AnalyticsChartTypeResolver;
 import com.linkedin.datahub.graphql.analytics.resolver.GetChartsResolver;
@@ -179,6 +180,7 @@ import com.linkedin.datahub.graphql.types.assertion.AssertionType;
 import com.linkedin.datahub.graphql.types.auth.AccessTokenMetadataType;
 import com.linkedin.datahub.graphql.types.chart.ChartType;
 import com.linkedin.datahub.graphql.types.common.mappers.OperationMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper;
 import com.linkedin.datahub.graphql.types.container.ContainerType;
 import com.linkedin.datahub.graphql.types.corpgroup.CorpGroupType;
 import com.linkedin.datahub.graphql.types.corpuser.CorpUserType;
@@ -633,9 +635,22 @@ public class GmsGraphQLEngine {
             .dataFetcher("getRootGlossaryNodes", new GetRootGlossaryNodesResolver(this.entityClient))
             .dataFetcher("entityExists", new EntityExistsResolver(this.entityService))
             .dataFetcher("getNativeUserInviteToken", new GetNativeUserInviteTokenResolver(this.nativeUserService))
+            .dataFetcher("entity", getEntityResolver())
         );
     }
 
+    private DataFetcher getEntityResolver() {
+        return new EntityTypeResolver(entityTypes,
+            (env) -> {
+                try {
+                    Urn urn = Urn.createFromString(env.getArgument(URN_FIELD_NAME));
+                    return UrnToEntityMapper.map(urn);
+                } catch (Exception e) {
+                    throw new RuntimeException("Failed to get entity", e);
+                }
+            });
+    }
+
     private DataFetcher getResolver(LoadableType<?, String> loadableType) {
         return getResolver(loadableType, this::getUrnField);
     }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchAcrossLineageResolver.java

@@ -3,6 +3,7 @@ package com.linkedin.datahub.graphql.resolvers.search;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.FacetFilterInput;
 import com.linkedin.datahub.graphql.generated.LineageDirection;
 import com.linkedin.datahub.graphql.generated.SearchAcrossLineageInput;
 import com.linkedin.datahub.graphql.generated.SearchAcrossLineageResults;
@@ -15,6 +16,7 @@ import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.net.URISyntaxException;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
@@ -57,6 +59,7 @@ public class SearchAcrossLineageResolver
 
     final int start = input.getStart() != null ? input.getStart() : DEFAULT_START;
     final int count = input.getCount() != null ? input.getCount() : DEFAULT_COUNT;
+    final Integer maxHops = getMaxHops(input.getFilters());
 
     com.linkedin.metadata.graph.LineageDirection resolvedDirection =
         com.linkedin.metadata.graph.LineageDirection.valueOf(lineageDirection.toString());
@@ -67,7 +70,7 @@ public class SearchAcrossLineageResolver
             urn, resolvedDirection, input.getTypes(), input.getQuery(), input.getFilters(), start, count);
         return UrnSearchAcrossLineageResultsMapper.map(
             _entityClient.searchAcrossLineage(urn, resolvedDirection, entityNames, sanitizedQuery,
-                ResolverUtils.buildFilter(input.getFilters()), null, start, count,
+                maxHops, ResolverUtils.buildFilter(input.getFilters()), null, start, count,
                 ResolverUtils.getAuthentication(environment)));
       } catch (RemoteInvocationException e) {
         log.error(
@@ -79,4 +82,21 @@ public class SearchAcrossLineageResolver
       }
     });
   }
+
+//  Assumption is that filter values for degree are either null, 3+, 2, or 1.
+  private Integer getMaxHops(List<FacetFilterInput> filters) {
+    Set<String> degreeFilterValues = filters.stream()
+        .filter(filter -> filter.getField().equals("degree"))
+        .map(FacetFilterInput::getValue)
+        .collect(Collectors.toSet());
+    Integer maxHops = null;
+    if (!degreeFilterValues.contains("3+")) {
+      if (degreeFilterValues.contains("2")) {
+        maxHops = 2;
+      } else if (degreeFilterValues.contains("1")) {
+        maxHops = 1;
+      }
+    }
+    return maxHops;
+  }
 }

datahub-graphql-core/src/main/resources/entity.graphql

@@ -168,6 +168,11 @@ type Query {
     Gets the current invite token. If the optional regenerate param is set to true, generate a new invite token.
     """
     getNativeUserInviteToken: InviteToken
+
+    """
+    Gets an entity based on its urn
+    """
+    entity(urn: String!): Entity
 }
 
 """

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

@@ -128,6 +128,7 @@ export default class EntityRegistry {
                         entity: relationship.entity as EntityInterface,
                         type: (relationship.entity as EntityInterface).type,
                     })),
+                numDownstreamChildren: genericEntityProperties?.downstream?.total,
                 upstreamChildren: genericEntityProperties?.upstream?.relationships
                     ?.filter((relationship) => relationship.entity)
                     // eslint-disable-next-line @typescript-eslint/dot-notation
@@ -136,6 +137,7 @@ export default class EntityRegistry {
                         entity: relationship.entity as EntityInterface,
                         type: (relationship.entity as EntityInterface).type,
                     })),
+                numUpstreamChildren: genericEntityProperties?.upstream?.total,
                 status: genericEntityProperties?.status,
             } as FetchedEntity) || undefined
         );

datahub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearch.tsx

@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { useState, useEffect } from 'react';
 import * as QueryString from 'query-string';
 import { useHistory, useLocation, useParams } from 'react-router';
 import { message } from 'antd';
@@ -57,6 +57,8 @@ type Props = {
     fixedFilter?: FacetFilterInput | null;
     fixedQuery?: string | null;
     placeholderText?: string | null;
+    defaultShowFilters?: boolean;
+    defaultFilters?: Array<FacetFilterInput>;
     useGetSearchResults?: (params: GetSearchResultsParams) => {
         data: SearchResultsInterface | undefined | null;
         loading: boolean;
@@ -70,6 +72,8 @@ export const EmbeddedListSearch = ({
     fixedFilter,
     fixedQuery,
     placeholderText,
+    defaultShowFilters,
+    defaultFilters,
     useGetSearchResults = useWrappedSearchResults,
 }: Props) => {
     const history = useHistory();
@@ -89,7 +93,7 @@ export const EmbeddedListSearch = ({
         .filter((filter) => filter.field === ENTITY_FILTER_NAME)
         .map((filter) => filter.value.toUpperCase() as EntityType);
 
-    const [showFilters, setShowFilters] = useState(false);
+    const [showFilters, setShowFilters] = useState(defaultShowFilters || false);
 
     const { refetch } = useGetSearchResults({
         variables: {
@@ -157,6 +161,14 @@ export const EmbeddedListSearch = ({
         setShowFilters(!showFilters);
     };
 
+    useEffect(() => {
+        if (defaultFilters) {
+            onChangeFilters(defaultFilters);
+        }
+        // only want to run once on page load
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, []);
+
     // Filter out the persistent filter values
     const filteredFilters = data?.facets?.filter((facet) => facet.field !== fixedFilter?.field) || [];
 

datahub-web-react/src/app/entity/shared/containers/profile/nav/EntityProfileNavBar.tsx

@@ -27,8 +27,8 @@ export const EntityProfileNavBar = ({ urn, entityType }: Props) => {
                 breadcrumbLinksEnabled={isBrowsable}
                 type={entityType}
                 path={browseData?.browsePaths?.[0]?.path || []}
-                upstreams={lineage?.upstreamChildren?.length || 0}
+                upstreams={lineage?.numUpstreamChildren || 0}
-                downstreams={lineage?.downstreamChildren?.length || 0}
+                downstreams={lineage?.numDownstreamChildren || 0}
             />
         </AffixWithHeight>
     );

datahub-web-react/src/app/entity/shared/tabs/Lineage/ImpactAnalysis.tsx

@@ -18,9 +18,10 @@ const ImpactAnalysisWrapper = styled.div`
 
 type Props = {
     urn: string;
+    direction: LineageDirection;
 };
 
-export const ImpactAnalysis = ({ urn }: Props) => {
+export const ImpactAnalysis = ({ urn, direction }: Props) => {
     const location = useLocation();
 
     const params = QueryString.parse(location.search, { arrayFormat: 'comma' });
@@ -38,7 +39,7 @@ export const ImpactAnalysis = ({ urn }: Props) => {
         variables: {
             input: {
                 urn,
-                direction: LineageDirection.Downstream,
+                direction,
                 types: entityFilters,
                 query,
                 start: (page - 1) * SearchCfg.RESULTS_PER_PAGE,
@@ -63,8 +64,10 @@ export const ImpactAnalysis = ({ urn }: Props) => {
             <EmbeddedListSearch
                 useGetSearchResults={generateUseSearchResultsViaRelationshipHook({
                     urn,
-                    direction: LineageDirection.Downstream,
+                    direction,
                 })}
+                defaultShowFilters
+                defaultFilters={[{ field: 'degree', value: '1' }]}
             />
         </ImpactAnalysisWrapper>
     );

datahub-web-react/src/app/entity/shared/tabs/Lineage/LineageTab.tsx

@@ -1,71 +1,66 @@
 import React, { useCallback, useState } from 'react';
 import { Button } from 'antd';
 import { useHistory } from 'react-router';
-import { BarsOutlined, PartitionOutlined } from '@ant-design/icons';
+import { ArrowDownOutlined, ArrowUpOutlined, PartitionOutlined } from '@ant-design/icons';
-import { VscGraphLeft } from 'react-icons/vsc';
+import styled from 'styled-components/macro';
-import styled from 'styled-components';
 
-import { useEntityData, useLineageData } from '../../EntityContext';
+import { useEntityData } from '../../EntityContext';
 import TabToolbar from '../../components/styled/TabToolbar';
 import { getEntityPath } from '../../containers/profile/utils';
 import { useEntityRegistry } from '../../../../useEntityRegistry';
-import { LineageTable } from './LineageTable';
 import { ImpactAnalysis } from './ImpactAnalysis';
-import { useAppConfig } from '../../../../useAppConfig';
+import { LineageDirection } from '../../../../../types.generated';
 
-const ImpactAnalysisIcon = styled(VscGraphLeft)`
+const StyledTabToolbar = styled(TabToolbar)`
-    transform: scaleX(-1);
+    justify-content: space-between;
-    font-size: 18px;
+`;
+
+const StyledButton = styled(Button)<{ isSelected: boolean }>`
+    ${(props) =>
+        props.isSelected &&
+        `
+        color: #1890ff;
+        &:focus {
+            color: #1890ff;
+        }    
+    `}
 `;
 
 export const LineageTab = () => {
     const { urn, entityType } = useEntityData();
     const history = useHistory();
     const entityRegistry = useEntityRegistry();
-    const lineage = useLineageData();
+    const [lineageDirection, setLineageDirection] = useState<string>(LineageDirection.Downstream);
-    const [showImpactAnalysis, setShowImpactAnalysis] = useState(false);
-    const appConfig = useAppConfig();
 
     const routeToLineage = useCallback(() => {
         history.push(getEntityPath(entityType, urn, entityRegistry, true));
     }, [history, entityType, urn, entityRegistry]);
 
-    const upstreamEntities = lineage?.upstreamChildren?.map((result) => result.entity);
-    const downstreamEntities = lineage?.downstreamChildren?.map((result) => result.entity);
     return (
         <>
-            <TabToolbar>
+            <StyledTabToolbar>
                 <div>
-                    <Button type="text" onClick={routeToLineage}>
+                    <StyledButton
-                        <PartitionOutlined />
+                        type="text"
-                        Visualize Lineage
+                        isSelected={lineageDirection === LineageDirection.Downstream}
-                    </Button>
+                        onClick={() => setLineageDirection(LineageDirection.Downstream)}
-                    {appConfig.config.lineageConfig.supportsImpactAnalysis &&
+                    >
-                        (showImpactAnalysis ? (
+                        <ArrowDownOutlined /> Downstream
-                            <Button type="text" onClick={() => setShowImpactAnalysis(false)}>
+                    </StyledButton>
-                                <span className="anticon">
+                    <StyledButton
-                                    <BarsOutlined />
+                        type="text"
-                                </span>
+                        isSelected={lineageDirection === LineageDirection.Upstream}
-                                Direct Dependencies
+                        onClick={() => setLineageDirection(LineageDirection.Upstream)}
-                            </Button>
+                    >
-                        ) : (
+                        <ArrowUpOutlined /> Upstream
-                            <Button type="text" onClick={() => setShowImpactAnalysis(true)}>
+                    </StyledButton>
-                                <span className="anticon">
-                                    <ImpactAnalysisIcon />
-                                </span>
-                                Impact Analysis
-                            </Button>
-                        ))}
                 </div>
-            </TabToolbar>
+                <Button type="text" onClick={routeToLineage}>
-            {showImpactAnalysis ? (
+                    <PartitionOutlined />
-                <ImpactAnalysis urn={urn} />
+                    Visualize Lineage
-            ) : (
+                </Button>
-                <>
+            </StyledTabToolbar>
-                    <LineageTable data={upstreamEntities} title={`${upstreamEntities?.length || 0} Upstream`} />
+            <ImpactAnalysis urn={urn} direction={lineageDirection as LineageDirection} />
-                    <LineageTable data={downstreamEntities} title={`${downstreamEntities?.length || 0} Downstream`} />
-                </>
-            )}
         </>
     );
 };

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

@@ -10,7 +10,7 @@ import { ANTD_GRAY } from '../entity/shared/constants';
 import { capitalizeFirstLetter } from '../shared/textUtil';
 import { nodeHeightFromTitleLength } from './utils/nodeHeightFromTitleLength';
 import { LineageExplorerContext } from './utils/LineageExplorerContext';
-import useLazyGetEntityQuery from './utils/useLazyGetEntityQuery';
+import { useGetEntityLineageLazyQuery } from '../../graphql/lineage.generated';
 
 const CLICK_DELAY_THRESHOLD = 1000;
 const DRAG_DISTANCE_THRESHOLD = 20;
@@ -89,13 +89,17 @@ export default function LineageEntityNode({
     const { expandTitles } = useContext(LineageExplorerContext);
     const [isExpanding, setIsExpanding] = useState(false);
     const [expandHover, setExpandHover] = useState(false);
-    const { getAsyncEntity, asyncData } = useLazyGetEntityQuery();
+    const [getAsyncEntityLineage, { data: asyncLineageData }] = useGetEntityLineageLazyQuery();
 
     useEffect(() => {
-        if (asyncData) {
+        if (asyncLineageData && asyncLineageData.entity) {
-            onExpandClick(asyncData);
+            const entityAndType = {
+                type: asyncLineageData.entity.type,
+                entity: { ...asyncLineageData.entity },
+            } as EntityAndType;
+            onExpandClick(entityAndType);
         }
-    }, [asyncData, onExpandClick]);
+    }, [asyncLineageData, onExpandClick]);
 
     const entityRegistry = useEntityRegistry();
     const unexploredHiddenChildren =
@@ -148,7 +152,8 @@ export default function LineageEntityNode({
                         onClick={() => {
                             setIsExpanding(true);
                             if (node.data.urn && node.data.type) {
-                                getAsyncEntity(node.data.urn, node.data.type);
+                                // getAsyncEntity(node.data.urn, node.data.type);
+                                getAsyncEntityLineage({ variables: { urn: node.data.urn } });
                             }
                         }}
                         onMouseOver={() => {

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

@@ -1,4 +1,4 @@
-import React, { useCallback, useEffect, useRef, useState } from 'react';
+import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useHistory } from 'react-router';
 
 import { Alert, Button, Drawer } from 'antd';
@@ -11,10 +11,10 @@ import CompactContext from '../shared/CompactContext';
 import { EntityAndType, EntitySelectParams, FetchedEntities } from './types';
 import LineageViz from './LineageViz';
 import extendAsyncEntities from './utils/extendAsyncEntities';
-import useGetEntityQuery from './utils/useGetEntityQuery';
 import { EntityType } from '../../types.generated';
 import { capitalizeFirstLetter } from '../shared/textUtil';
 import { ANTD_GRAY } from '../entity/shared/constants';
+import { GetEntityLineageQuery, useGetEntityLineageQuery } from '../../graphql/lineage.generated';
 
 const DEFAULT_DISTANCE_FROM_TOP = 106;
 
@@ -45,6 +45,16 @@ function usePrevious(value) {
     return ref.current;
 }
 
+export function getEntityAndType(lineageData?: GetEntityLineageQuery) {
+    if (lineageData && lineageData.entity) {
+        return {
+            type: lineageData.entity.type,
+            entity: { ...lineageData.entity },
+        } as EntityAndType;
+    }
+    return null;
+}
+
 type Props = {
     urn: string;
     type: EntityType;
@@ -56,7 +66,8 @@ export default function LineageExplorer({ urn, type }: Props) {
 
     const entityRegistry = useEntityRegistry();
 
-    const { loading, error, data } = useGetEntityQuery(urn, type);
+    const { loading, error, data } = useGetEntityLineageQuery({ variables: { urn } });
+    const entityData: EntityAndType | null | undefined = useMemo(() => getEntityAndType(data), [data]);
 
     const [isDrawerVisible, setIsDrawVisible] = useState(false);
     const [selectedEntity, setSelectedEntity] = useState<EntitySelectParams | undefined>(undefined);
@@ -89,10 +100,10 @@ export default function LineageExplorer({ urn, type }: Props) {
     };
 
     useEffect(() => {
-        if (type && data) {
+        if (type && entityData) {
-            maybeAddAsyncLoadedEntity(data);
+            maybeAddAsyncLoadedEntity(entityData);
         }
-    }, [data, asyncEntities, setAsyncEntities, maybeAddAsyncLoadedEntity, urn, previousUrn, type]);
+    }, [entityData, asyncEntities, setAsyncEntities, maybeAddAsyncLoadedEntity, urn, previousUrn, type]);
 
     if (error || (!loading && !error && !data)) {
         return <Alert type="error" message={error?.message || 'Entity failed to load'} />;
@@ -109,7 +120,7 @@ export default function LineageExplorer({ urn, type }: Props) {
                     <LineageViz
                         selectedEntity={selectedEntity}
                         fetchedEntities={asyncEntities}
-                        entityAndType={data}
+                        entityAndType={entityData}
                         onEntityClick={(params: EntitySelectParams) => {
                             setIsDrawVisible(true);
                             setSelectedEntity(params);

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

@@ -1,3 +1,4 @@
+import { FullLineageResultsFragment } from '../../graphql/lineage.generated';
 import {
     Chart,
     Dashboard,
@@ -34,7 +35,9 @@ export type FetchedEntity = {
     icon?: string;
     // children?: Array<string>;
     upstreamChildren?: Array<EntityAndType>;
+    numUpstreamChildren?: number;
     downstreamChildren?: Array<EntityAndType>;
+    numDownstreamChildren?: number;
     fullyFetched?: boolean;
     platform?: string;
     status?: Maybe<Status>;
@@ -129,3 +132,9 @@ export type EntityAndType =
           type: EntityType.MlprimaryKey;
           entity: MlPrimaryKey;
       };
+
+export interface LineageResult {
+    urn: string;
+    upstream?: Maybe<{ __typename?: 'EntityLineageResult' } & FullLineageResultsFragment>;
+    downstream?: Maybe<{ __typename?: 'EntityLineageResult' } & FullLineageResultsFragment>;
+}

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

@@ -1,206 +0,0 @@
-import { useMemo } from 'react';
-import { useGetChartQuery } from '../../../graphql/chart.generated';
-import { useGetDashboardQuery } from '../../../graphql/dashboard.generated';
-import { useGetDatasetQuery } from '../../../graphql/dataset.generated';
-import { useGetDataJobQuery } from '../../../graphql/dataJob.generated';
-import { useGetMlFeatureTableQuery } from '../../../graphql/mlFeatureTable.generated';
-import { useGetMlFeatureQuery } from '../../../graphql/mlFeature.generated';
-import { useGetMlPrimaryKeyQuery } from '../../../graphql/mlPrimaryKey.generated';
-import { EntityType } from '../../../types.generated';
-import { EntityAndType } from '../types';
-import { useGetMlModelQuery } from '../../../graphql/mlModel.generated';
-import { useGetMlModelGroupQuery } from '../../../graphql/mlModelGroup.generated';
-
-export default function useGetEntityQuery(urn: string, entityType?: EntityType) {
-    const allResults = {
-        [EntityType.Dataset]: useGetDatasetQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.Dataset,
-        }),
-        [EntityType.Chart]: useGetChartQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.Chart,
-        }),
-        [EntityType.Dashboard]: useGetDashboardQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.Dashboard,
-        }),
-        [EntityType.DataJob]: useGetDataJobQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.DataJob,
-        }),
-        [EntityType.MlfeatureTable]: useGetMlFeatureTableQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.MlfeatureTable,
-        }),
-        [EntityType.Mlfeature]: useGetMlFeatureQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.Mlfeature,
-        }),
-        [EntityType.MlprimaryKey]: useGetMlPrimaryKeyQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.MlprimaryKey,
-        }),
-        [EntityType.Mlmodel]: useGetMlModelQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.Mlmodel,
-        }),
-        [EntityType.MlmodelGroup]: useGetMlModelGroupQuery({
-            variables: { urn },
-            skip: entityType !== EntityType.MlmodelGroup,
-        }),
-    };
-
-    const returnEntityAndType: EntityAndType | undefined = useMemo(() => {
-        let returnData;
-        switch (entityType) {
-            case EntityType.Dataset:
-                returnData = allResults[EntityType.Dataset].data?.dataset;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Dataset,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Chart:
-                returnData = allResults[EntityType.Chart]?.data?.chart;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Chart,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Dashboard:
-                returnData = allResults[EntityType.Dashboard]?.data?.dashboard;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Dashboard,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.DataJob:
-                returnData = allResults[EntityType.DataJob]?.data?.dataJob;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.DataJob,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlfeatureTable:
-                returnData = allResults[EntityType.MlfeatureTable]?.data?.mlFeatureTable;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlfeatureTable,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Mlfeature:
-                returnData = allResults[EntityType.Mlfeature]?.data?.mlFeature;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Mlfeature,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlprimaryKey:
-                returnData = allResults[EntityType.MlprimaryKey]?.data?.mlPrimaryKey;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlprimaryKey,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Mlmodel:
-                returnData = allResults[EntityType.Mlmodel]?.data?.mlModel;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Mlmodel,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlmodelGroup:
-                returnData = allResults[EntityType.MlmodelGroup]?.data?.mlModelGroup;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlmodelGroup,
-                    } as EntityAndType;
-                }
-                break;
-            default:
-                break;
-        }
-        return undefined;
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-    }, [
-        urn,
-        entityType,
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Dataset],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Chart],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Dashboard],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.DataJob],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlfeatureTable],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Mlmodel],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlmodelGroup],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Mlfeature],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlprimaryKey],
-    ]);
-
-    const returnObject = useMemo(() => {
-        if (!entityType) {
-            return {
-                loading: false,
-                error: null,
-                data: null,
-            };
-        }
-
-        return {
-            data: returnEntityAndType,
-            loading: allResults[entityType].loading,
-            error: allResults[entityType].error,
-        };
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-    }, [
-        urn,
-        entityType,
-        returnEntityAndType,
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Dataset],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Chart],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Dashboard],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.DataJob],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlfeatureTable],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Mlmodel],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlmodelGroup],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.Mlfeature],
-        // eslint-disable-next-line react-hooks/exhaustive-deps
-        allResults[EntityType.MlprimaryKey],
-    ]);
-
-    return returnObject;
-}

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

@@ -1,181 +0,0 @@
-import { useCallback, useMemo, useState } from 'react';
-import { useGetChartLazyQuery } from '../../../graphql/chart.generated';
-import { useGetDashboardLazyQuery } from '../../../graphql/dashboard.generated';
-import { useGetDatasetLazyQuery } from '../../../graphql/dataset.generated';
-import { useGetDataJobLazyQuery } from '../../../graphql/dataJob.generated';
-import { useGetMlFeatureTableLazyQuery } from '../../../graphql/mlFeatureTable.generated';
-import { useGetMlFeatureLazyQuery } from '../../../graphql/mlFeature.generated';
-import { useGetMlPrimaryKeyLazyQuery } from '../../../graphql/mlPrimaryKey.generated';
-import { EntityType } from '../../../types.generated';
-import { EntityAndType } from '../types';
-import { useGetMlModelLazyQuery } from '../../../graphql/mlModel.generated';
-import { useGetMlModelGroupLazyQuery } from '../../../graphql/mlModelGroup.generated';
-
-export default function useLazyGetEntityQuery() {
-    const [fetchedEntityType, setFetchedEntityType] = useState<EntityType | undefined>(undefined);
-    const [getAsyncDataset, { data: asyncDatasetData }] = useGetDatasetLazyQuery();
-    const [getAsyncChart, { data: asyncChartData }] = useGetChartLazyQuery();
-    const [getAsyncDashboard, { data: asyncDashboardData }] = useGetDashboardLazyQuery();
-    const [getAsyncDataJob, { data: asyncDataJobData }] = useGetDataJobLazyQuery();
-    const [getAsyncMLFeatureTable, { data: asyncMLFeatureTable }] = useGetMlFeatureTableLazyQuery();
-    const [getAsyncMLFeature, { data: asyncMLFeature }] = useGetMlFeatureLazyQuery();
-    const [getAsyncMLPrimaryKey, { data: asyncMLPrimaryKey }] = useGetMlPrimaryKeyLazyQuery();
-    const [getAsyncMlModel, { data: asyncMlModel }] = useGetMlModelLazyQuery();
-    const [getAsyncMlModelGroup, { data: asyncMlModelGroup }] = useGetMlModelGroupLazyQuery();
-
-    const getAsyncEntity = useCallback(
-        (urn: string, type: EntityType) => {
-            if (type === EntityType.Dataset) {
-                setFetchedEntityType(type);
-                getAsyncDataset({ variables: { urn } });
-            }
-            if (type === EntityType.Chart) {
-                setFetchedEntityType(type);
-                getAsyncChart({ variables: { urn } });
-            }
-            if (type === EntityType.Dashboard) {
-                setFetchedEntityType(type);
-                getAsyncDashboard({ variables: { urn } });
-            }
-            if (type === EntityType.DataJob) {
-                setFetchedEntityType(type);
-                getAsyncDataJob({ variables: { urn } });
-            }
-            if (type === EntityType.MlfeatureTable) {
-                setFetchedEntityType(type);
-                getAsyncMLFeatureTable({ variables: { urn } });
-            }
-            if (type === EntityType.Mlfeature) {
-                setFetchedEntityType(type);
-                getAsyncMLFeature({ variables: { urn } });
-            }
-            if (type === EntityType.MlprimaryKey) {
-                setFetchedEntityType(type);
-                getAsyncMLPrimaryKey({ variables: { urn } });
-            }
-            if (type === EntityType.Mlmodel) {
-                setFetchedEntityType(type);
-                getAsyncMlModel({ variables: { urn } });
-            }
-            if (type === EntityType.MlmodelGroup) {
-                setFetchedEntityType(type);
-                getAsyncMlModelGroup({ variables: { urn } });
-            }
-        },
-        [
-            setFetchedEntityType,
-            getAsyncChart,
-            getAsyncDataset,
-            getAsyncDashboard,
-            getAsyncDataJob,
-            getAsyncMLFeatureTable,
-            getAsyncMLFeature,
-            getAsyncMLPrimaryKey,
-            getAsyncMlModel,
-            getAsyncMlModelGroup,
-        ],
-    );
-
-    const returnEntityAndType: EntityAndType | undefined = useMemo(() => {
-        let returnData;
-        switch (fetchedEntityType) {
-            case EntityType.Dataset:
-                returnData = asyncDatasetData?.dataset;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Dataset,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Chart:
-                returnData = asyncChartData?.chart;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Chart,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Dashboard:
-                returnData = asyncDashboardData?.dashboard;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Dashboard,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.DataJob:
-                returnData = asyncDataJobData?.dataJob;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.DataJob,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlfeatureTable:
-                returnData = asyncMLFeatureTable?.mlFeatureTable;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlfeatureTable,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Mlfeature:
-                returnData = asyncMLFeature?.mlFeature;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Mlfeature,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlprimaryKey:
-                returnData = asyncMLPrimaryKey?.mlPrimaryKey;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlprimaryKey,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.Mlmodel:
-                returnData = asyncMlModel?.mlModel;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.Mlmodel,
-                    } as EntityAndType;
-                }
-                break;
-            case EntityType.MlmodelGroup:
-                returnData = asyncMlModelGroup?.mlModelGroup;
-                if (returnData) {
-                    return {
-                        entity: returnData,
-                        type: EntityType.MlmodelGroup,
-                    } as EntityAndType;
-                }
-                break;
-            default:
-                break;
-        }
-        return undefined;
-    }, [
-        asyncDatasetData,
-        asyncChartData,
-        asyncDashboardData,
-        asyncDataJobData,
-        fetchedEntityType,
-        asyncMLFeatureTable,
-        asyncMLFeature,
-        asyncMLPrimaryKey,
-        asyncMlModel,
-        asyncMlModelGroup,
-    ]);
-
-    return { getAsyncEntity, asyncData: returnEntityAndType };
-}

datahub-web-react/src/app/search/SearchFilters.tsx

@@ -4,7 +4,7 @@ import { useEffect, useState } from 'react';
 import { FacetMetadata } from '../../types.generated';
 import { SearchFilter } from './SearchFilter';
 
-const TOP_FILTERS = ['entity', 'tags', 'glossaryTerms', 'domains', 'owners'];
+const TOP_FILTERS = ['degree', 'entity', 'tags', 'glossaryTerms', 'domains', 'owners'];
 
 export const SearchFilterWrapper = styled.div`
     max-height: 100%;

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

@@ -63,10 +63,10 @@ query getChart($urn: String!) {
             ...parentContainersFields
         }
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         status {
             removed

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

@@ -8,10 +8,10 @@ query getDashboard($urn: String!) {
             ...fullRelationshipResults
         }
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
     }
 }

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

@@ -50,10 +50,10 @@ query getDataFlow($urn: String!) {
     dataFlow(urn: $urn) {
         ...dataFlowFields
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         childJobs: relationships(input: { types: ["IsPartOf"], direction: INCOMING, start: 0, count: 100 }) {
             start

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

@@ -20,10 +20,10 @@ query getDataJob($urn: String!) {
             ...fullRelationshipResults
         }
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         status {
             removed

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

@@ -126,10 +126,10 @@ query getDataset($urn: String!) {
             lastUpdatedTimestamp
         }
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         ...viewProperties
         autoRenderAspects: aspects(input: { autoRenderOnly: true }) {

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

@@ -1,4 +1,4 @@
-fragment relationshipFields on EntityWithRelationships {
+fragment lineageNodeProperties on EntityWithRelationships {
     urn
     type
     ... on DataJob {
@@ -139,6 +139,10 @@ fragment relationshipFields on EntityWithRelationships {
     ... on MLPrimaryKey {
         ...nonRecursiveMLPrimaryKey
     }
+}
+
+fragment relationshipFields on EntityWithRelationships {
+    ...lineageNodeProperties
     upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
         ...leafLineageResults
     }
@@ -171,3 +175,25 @@ fragment leafLineageResults on EntityLineageResult {
         }
     }
 }
+
+fragment partialLineageResults on EntityLineageResult {
+    start
+    count
+    total
+}
+
+query getEntityLineage($urn: String!) {
+    entity(urn: $urn) {
+        urn
+        type
+        ...lineageNodeProperties
+        ... on EntityWithRelationships {
+            upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
+                ...fullLineageResults
+            }
+            downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
+                ...fullLineageResults
+            }
+        }
+    }
+}

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

@@ -2,10 +2,10 @@ query getMLFeature($urn: String!) {
     mlFeature(urn: $urn) {
         ...nonRecursiveMLFeature
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         featureTables: relationships(input: { types: ["Contains"], direction: INCOMING, start: 0, count: 100 }) {
             ...fullRelationshipResults

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

@@ -2,10 +2,10 @@ query getMLFeatureTable($urn: String!) {
     mlFeatureTable(urn: $urn) {
         ...nonRecursiveMLFeatureTable
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
     }
 }

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

@@ -2,10 +2,10 @@ query getMLModel($urn: String!) {
     mlModel(urn: $urn) {
         ...nonRecursiveMLModel
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         features: relationships(input: { types: ["Consumes"], direction: OUTGOING, start: 0, count: 100 }) {
             ...fullRelationshipResults

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

@@ -22,10 +22,10 @@ query getMLModelGroup($urn: String!) {
             ...fullRelationshipResults
         }
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
     }
 }

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

@@ -2,10 +2,10 @@ query getMLPrimaryKey($urn: String!) {
     mlPrimaryKey(urn: $urn) {
         ...nonRecursiveMLPrimaryKey
         upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
-            ...fullLineageResults
+            ...partialLineageResults
         }
         featureTables: relationships(input: { types: ["KeyedBy"], direction: INCOMING, start: 0, count: 100 }) {
             ...fullRelationshipResults

metadata-io/src/main/java/com/linkedin/metadata/graph/GraphService.java

@@ -94,8 +94,7 @@ public interface GraphService {
   default EntityLineageResult getLineage(@Nonnull Urn entityUrn, @Nonnull LineageDirection direction, int offset,
       int count, int maxHops) {
     if (maxHops > 1) {
-      throw new UnsupportedOperationException(
+      maxHops = 1;
-          String.format("More than 1 hop is not supported for %s", this.getClass().getSimpleName()));
     }
     List<LineageRegistry.EdgeInfo> edgesToFetch =
         getLineageRegistry().getLineageRelationships(entityUrn.getEntityType(), direction);

metadata-io/src/main/java/com/linkedin/metadata/search/LineageSearchService.java

@@ -58,6 +58,7 @@ public class LineageSearchService {
    * @param direction Direction of the relationship
    * @param entities list of entities to search (If empty, searches across all entities)
    * @param input the search input text
+   * @param maxHops the maximum number of hops away to search for. If null, defaults to 1000
    * @param inputFilters the request map with fields and values as filters to be applied to search hits
    * @param sortCriterion {@link SortCriterion} to be applied to search results
    * @param from index to start the search from
@@ -67,12 +68,13 @@ public class LineageSearchService {
   @Nonnull
   @WithSpan
   public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull LineageDirection direction,
-      @Nonnull List<String> entities, @Nullable String input, @Nullable Filter inputFilters,
+      @Nonnull List<String> entities, @Nullable String input, @Nullable Integer maxHops, @Nullable Filter inputFilters,
       @Nullable SortCriterion sortCriterion, int from, int size) {
     // Cache multihop result for faster performance
     EntityLineageResult lineageResult = cache.get(Pair.of(sourceUrn, direction), EntityLineageResult.class);
     if (lineageResult == null) {
-      lineageResult = _graphService.getLineage(sourceUrn, direction, 0, MAX_RELATIONSHIPS, 1000);
+      maxHops = maxHops != null ? maxHops : 1000;
+      lineageResult = _graphService.getLineage(sourceUrn, direction, 0, MAX_RELATIONSHIPS, maxHops);
     }
 
     // Filter hopped result based on the set of entities to return and inputFilters before sending to search

metadata-io/src/test/java/com/linkedin/metadata/search/LineageSearchServiceTest.java

@@ -134,11 +134,11 @@ public class LineageSearchServiceTest {
         anyInt())).thenReturn(mockResult(Collections.emptyList()));
     LineageSearchResult searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME),
-            "test", null, null, 0, 10);
+            "test", null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     clearCache();
 
@@ -147,11 +147,11 @@ public class LineageSearchServiceTest {
         mockResult(ImmutableList.of(new LineageRelationship().setEntity(TEST_URN).setType("test").setDegree(1))));
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(ENTITY_NAME),
-            "test", null, null, 0, 10);
+            "test", null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     clearCache();
 
@@ -168,7 +168,7 @@ public class LineageSearchServiceTest {
         anyInt())).thenReturn(mockResult(Collections.emptyList()));
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     assertEquals(searchResult.getEntities().size(), 0);
     clearCache();
@@ -178,21 +178,21 @@ public class LineageSearchServiceTest {
         mockResult(ImmutableList.of(new LineageRelationship().setEntity(urn).setType("test").setDegree(1))));
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 1);
     assertEquals(searchResult.getEntities().get(0).getEntity(), urn);
     assertEquals(searchResult.getEntities().get(0).getDegree().intValue(), 1);
 
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            QueryUtils.newFilter("degree.keyword", "1"), null, 0, 10);
+            null, QueryUtils.newFilter("degree.keyword", "1"), null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 1);
     assertEquals(searchResult.getEntities().get(0).getEntity(), urn);
     assertEquals(searchResult.getEntities().get(0).getDegree().intValue(), 1);
 
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            QueryUtils.newFilter("degree.keyword", "2"), null, 0, 10);
+            null, QueryUtils.newFilter("degree.keyword", "2"), null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     assertEquals(searchResult.getEntities().size(), 0);
     clearCache();
@@ -208,7 +208,7 @@ public class LineageSearchServiceTest {
 
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 1);
     assertEquals(searchResult.getEntities().get(0).getEntity(), urn);
     clearCache();
@@ -218,7 +218,7 @@ public class LineageSearchServiceTest {
         mockResult(ImmutableList.of(new LineageRelationship().setEntity(urn2).setType("test").setDegree(1))));
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
     assertEquals(searchResult.getEntities().size(), 0);
     clearCache();
@@ -232,7 +232,7 @@ public class LineageSearchServiceTest {
         mockResult(ImmutableList.of(new LineageRelationship().setEntity(urn).setType("test").setDegree(1))));
     searchResult =
         _lineageSearchService.searchAcrossLineage(TEST_URN, LineageDirection.DOWNSTREAM, ImmutableList.of(), "test",
-            null, null, 0, 10);
+            null, null, null, 0, 10);
     assertEquals(searchResult.getNumEntities().intValue(), 0);
   }
 }

metadata-service/restli-api/src/main/idl/com.linkedin.entity.entities.restspec.json

@@ -257,6 +257,10 @@
         "name" : "input",
         "type" : "string",
         "optional" : true
+      }, {
+        "name" : "maxHops",
+        "type" : "int",
+        "optional" : true
       }, {
         "name" : "filter",
         "type" : "com.linkedin.metadata.query.filter.Filter",

metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json

@@ -5822,6 +5822,10 @@
           "name" : "input",
           "type" : "string",
           "optional" : true
+        }, {
+          "name" : "maxHops",
+          "type" : "int",
+          "optional" : true
         }, {
           "name" : "filter",
           "type" : "com.linkedin.metadata.query.filter.Filter",

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java

@@ -187,6 +187,7 @@ public interface EntityClient {
    * @param direction Direction of the relationship
    * @param entities list of entities to search (If empty, searches across all entities)
    * @param input the search input text
+   * @param maxHops the max number of hops away to search for. If null, searches all hops.
    * @param filter the request map with fields and values as filters to be applied to search hits
    * @param sortCriterion {@link SortCriterion} to be applied to search results
    * @param start index to start the search from
@@ -195,7 +196,7 @@ public interface EntityClient {
    */
   @Nonnull
   public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull LineageDirection direction,
-      @Nonnull List<String> entities, @Nonnull String input, @Nullable Filter filter,
+      @Nonnull List<String> entities, @Nonnull String input, @Nullable Integer maxHops, @Nullable Filter filter,
       @Nullable SortCriterion sortCriterion, int start, int count, @Nonnull final Authentication authentication)
       throws RemoteInvocationException;
 

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/JavaEntityClient.java

@@ -305,10 +305,10 @@ public class JavaEntityClient implements EntityClient {
     @Nonnull
     @Override
     public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull LineageDirection direction,
-        @Nonnull List<String> entities, @Nullable String input, @Nullable Filter filter,
+        @Nonnull List<String> entities, @Nullable String input, @Nullable Integer maxHops, @Nullable Filter filter,
         @Nullable SortCriterion sortCriterion, int start, int count, @Nonnull final Authentication authentication)
         throws RemoteInvocationException {
-        return _lineageSearchService.searchAcrossLineage(sourceUrn, direction, entities, input, filter,
+        return _lineageSearchService.searchAcrossLineage(sourceUrn, direction, entities, input, maxHops, filter,
             sortCriterion, start, count);
     }
 

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/RestliEntityClient.java

@@ -416,7 +416,7 @@ public class RestliEntityClient extends BaseClient implements EntityClient {
   @Nonnull
   @Override
   public LineageSearchResult searchAcrossLineage(@Nonnull Urn sourceUrn, @Nonnull LineageDirection direction,
-      @Nonnull List<String> entities, @Nonnull String input, @Nullable Filter filter,
+      @Nonnull List<String> entities, @Nonnull String input, @Nullable Integer maxHops, @Nullable Filter filter,
       @Nullable SortCriterion sortCriterion, int start, int count, @Nonnull final Authentication authentication)
       throws RemoteInvocationException {
 

metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/EntityResource.java

@@ -274,16 +274,16 @@ public class EntityResource extends CollectionResourceTaskTemplate<String, Entit
   public Task<LineageSearchResult> searchAcrossLineage(@ActionParam(PARAM_URN) @Nonnull String urnStr,
       @ActionParam(PARAM_DIRECTION) String direction,
       @ActionParam(PARAM_ENTITIES) @Optional @Nullable String[] entities,
-      @ActionParam(PARAM_INPUT) @Optional @Nullable String input, @ActionParam(PARAM_FILTER) @Optional @Nullable Filter filter,
+      @ActionParam(PARAM_INPUT) @Optional @Nullable String input, @ActionParam(PARAM_MAX_HOPS) @Optional @Nullable Integer maxHops,
-      @ActionParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion, @ActionParam(PARAM_START) int start,
+      @ActionParam(PARAM_FILTER) @Optional @Nullable Filter filter, @ActionParam(PARAM_SORT) @Optional @Nullable SortCriterion sortCriterion,
-      @ActionParam(PARAM_COUNT) int count) throws URISyntaxException {
+      @ActionParam(PARAM_START) int start, @ActionParam(PARAM_COUNT) int count) throws URISyntaxException {
     Urn urn = Urn.createFromString(urnStr);
     List<String> entityList = entities == null ? Collections.emptyList() : Arrays.asList(entities);
     log.info("GET SEARCH RESULTS ACROSS RELATIONSHIPS for source urn {}, direction {}, entities {} with query {}",
         urnStr, direction, entityList, input);
     return RestliUtil.toTask(
         () -> _lineageSearchService.searchAcrossLineage(urn, LineageDirection.valueOf(direction), entityList,
-            input, filter, sortCriterion, start, count), "searchAcrossRelationships");
+            input, maxHops, filter, sortCriterion, start, count), "searchAcrossRelationships");
   }
 
   @Action(name = ACTION_LIST)

metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/restli/RestliConstants.java

@@ -19,6 +19,7 @@ public final class RestliConstants {
   public static final String ACTION_LIST_URNS_FROM_INDEX = "listUrnsFromIndex";
 
   public static final String PARAM_INPUT = "input";
+  public static final String PARAM_MAX_HOPS = "maxHops";
   public static final String PARAM_ASPECTS = "aspects";
   public static final String PARAM_FILTER = "filter";
   public static final String PARAM_GROUP = "group";

smoke-test/tests/cypress/cypress/integration/lineage/impact_analysis.js

@@ -2,7 +2,8 @@ describe('mutations', () => {
   it('can create and add a tag to dataset and visit new tag page', () => {
     cy.login();
     cy.visit('/dataset/urn:li:dataset:(urn:li:dataPlatform:kafka,SampleCypressKafkaDataset,PROD)/Lineage?is_lineage_mode=false');
-    cy.contains('Impact Analysis').click({ force: true });
+    // click to show more relationships now that we default to 1 degree of dependency
+    cy.contains('3+').click({ force: true });
 
     // impact analysis can take a beat- don't want to time out here
     cy.wait(5000);