Re: [PR] feat: add column-level lineage to the OpenLineage (druid)

["FrankChen021 (via GitHub)" <gi...@apache.org>]

FrankChen021 commented on code in PR #19643:
URL: https://github.com/apache/druid/pull/19643#discussion_r3505744946


##########
extensions-contrib/openlineage-emitter/src/main/java/org/apache/druid/extensions/openlineage/OpenLineageRequestLogger.java:
##########
@@ -309,11 +347,275 @@ static String extractMsqOutputDatasource(String sql)
     }
   }
 
+  /**
+   * Column usage roles for the custom {@code druid_columnUsage} dataset facet. Names align with
+   * OpenLineage transformation subtypes where they exist (GROUP_BY, AGGREGATION, FILTER, JOIN);
+   * PROJECTION corresponds to OpenLineage IDENTITY/TRANSFORMATION (a selected/projected column).
+   */
+  enum ColumnRole
+  {
+    PROJECTION, GROUP_BY, AGGREGATION, FILTER, JOIN
+  }
+
+  /**
+   * Resolves the per-base-table column-usage map for a native query, used to attach the {@code schema}
+   * and {@code druid_columnUsage} dataset facets to input datasets. Handles all datasource shapes:
+   * tables, joins (attributing columns per side via the planner's join prefixes), unions, sub-queries
+   * (recursing into the sub-query's own columns), and datasource wrappers (restricted/filtered/unnest).
+   * Columns that have no base table (lookups, inline data) are dropped rather than fabricated.
+   *
+   * <p>Returns {@code null} (yielding table-level lineage only) when no base-table columns can be
+   * determined, or on any error -- it never fabricates or mis-attributes columns.
+   */
+  @Nullable
+  private Map<String, Map<String, EnumSet<ColumnRole>>> extractColumnsByTable(Query<?> query)
+  {
+    try {
+      Map<String, Map<String, EnumSet<ColumnRole>>> result = new TreeMap<>();
+      collectInto(result, query);
+      return result.isEmpty() ? null : result;
+    }
+    // StackOverflowError (an Error, not an Exception) is caught too so that a pathologically deep
+    // query plan degrades to table-level lineage rather than breaking the request-logging path.
+    catch (Exception | StackOverflowError e) {
+      log.debug(e, "Failed to extract column lineage; falling back to table-level lineage");
+      return null;
+    }
+  }
+
+  /**
+   * Walks {@code query}'s column-bearing parts and attributes the referenced columns to the base
+   * tables of its datasource. {@link UnionQuery} (whose {@code getDataSource()} is undefined) is
+   * handled by recursing into each of its branches.
+   */
+  private void collectInto(Map<String, Map<String, EnumSet<ColumnRole>>> result, Query<?> query)
+  {
+    if (query instanceof UnionQuery) {
+      for (DataSource branch : ((UnionQuery) query).getDataSources()) {
+        attribute(result, branch, Collections.emptyMap());
+      }
+      return;
+    }
+    attribute(result, query.getDataSource(), collectColumnRoles(query));
+  }
+
+  /**
+   * Attributes {@code roles} (column to roles, expressed in {@code dataSource}'s output namespace) to
+   * the underlying base tables, recursing through the datasource tree. Sub-queries are recursed into
+   * via {@link #collectInto} (their columns come from their own parts, not the outer references).
+   */
+  private void attribute(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      DataSource dataSource,
+      Map<String, EnumSet<ColumnRole>> roles
+  )
+  {
+    if (dataSource instanceof TableDataSource) {
+      // Also covers GlobalTableDataSource (a TableDataSource subclass); it is a real named table, so
+      // handling it here is intentional. Keep this branch ahead of the wrapper branches below.
+      String table = ((TableDataSource) dataSource).getName();
+      for (Map.Entry<String, EnumSet<ColumnRole>> entry : roles.entrySet()) {
+        addRoles(result, table, entry.getKey(), entry.getValue());
+      }
+    } else if (dataSource instanceof RestrictedDataSource) {
+      attribute(result, ((RestrictedDataSource) dataSource).getBase(), roles);
+    } else if (dataSource instanceof FilteredDataSource) {
+      attribute(result, ((FilteredDataSource) dataSource).getBase(), roles);

Review Comment:
   [P2] FilteredDataSource filters are omitted from column lineage
   
   This wrapper is unwrapped without adding FilteredDataSource.getFilter() columns as FILTER roles. Druid creates this datasource for pushed-down filters, so a filtered datasource feeding another wrapper such as UNNEST can emit column facets that omit the filter column entirely unless it is also projected or aggregated.



##########
extensions-contrib/openlineage-emitter/src/main/java/org/apache/druid/extensions/openlineage/OpenLineageRequestLogger.java:
##########
@@ -309,11 +347,275 @@ static String extractMsqOutputDatasource(String sql)
     }
   }
 
+  /**
+   * Column usage roles for the custom {@code druid_columnUsage} dataset facet. Names align with
+   * OpenLineage transformation subtypes where they exist (GROUP_BY, AGGREGATION, FILTER, JOIN);
+   * PROJECTION corresponds to OpenLineage IDENTITY/TRANSFORMATION (a selected/projected column).
+   */
+  enum ColumnRole
+  {
+    PROJECTION, GROUP_BY, AGGREGATION, FILTER, JOIN
+  }
+
+  /**
+   * Resolves the per-base-table column-usage map for a native query, used to attach the {@code schema}
+   * and {@code druid_columnUsage} dataset facets to input datasets. Handles all datasource shapes:
+   * tables, joins (attributing columns per side via the planner's join prefixes), unions, sub-queries
+   * (recursing into the sub-query's own columns), and datasource wrappers (restricted/filtered/unnest).
+   * Columns that have no base table (lookups, inline data) are dropped rather than fabricated.
+   *
+   * <p>Returns {@code null} (yielding table-level lineage only) when no base-table columns can be
+   * determined, or on any error -- it never fabricates or mis-attributes columns.
+   */
+  @Nullable
+  private Map<String, Map<String, EnumSet<ColumnRole>>> extractColumnsByTable(Query<?> query)
+  {
+    try {
+      Map<String, Map<String, EnumSet<ColumnRole>>> result = new TreeMap<>();
+      collectInto(result, query);
+      return result.isEmpty() ? null : result;
+    }
+    // StackOverflowError (an Error, not an Exception) is caught too so that a pathologically deep
+    // query plan degrades to table-level lineage rather than breaking the request-logging path.
+    catch (Exception | StackOverflowError e) {
+      log.debug(e, "Failed to extract column lineage; falling back to table-level lineage");
+      return null;
+    }
+  }
+
+  /**
+   * Walks {@code query}'s column-bearing parts and attributes the referenced columns to the base
+   * tables of its datasource. {@link UnionQuery} (whose {@code getDataSource()} is undefined) is
+   * handled by recursing into each of its branches.
+   */
+  private void collectInto(Map<String, Map<String, EnumSet<ColumnRole>>> result, Query<?> query)
+  {
+    if (query instanceof UnionQuery) {
+      for (DataSource branch : ((UnionQuery) query).getDataSources()) {
+        attribute(result, branch, Collections.emptyMap());
+      }
+      return;
+    }
+    attribute(result, query.getDataSource(), collectColumnRoles(query));
+  }
+
+  /**
+   * Attributes {@code roles} (column to roles, expressed in {@code dataSource}'s output namespace) to
+   * the underlying base tables, recursing through the datasource tree. Sub-queries are recursed into
+   * via {@link #collectInto} (their columns come from their own parts, not the outer references).
+   */
+  private void attribute(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      DataSource dataSource,
+      Map<String, EnumSet<ColumnRole>> roles
+  )
+  {
+    if (dataSource instanceof TableDataSource) {
+      // Also covers GlobalTableDataSource (a TableDataSource subclass); it is a real named table, so
+      // handling it here is intentional. Keep this branch ahead of the wrapper branches below.
+      String table = ((TableDataSource) dataSource).getName();
+      for (Map.Entry<String, EnumSet<ColumnRole>> entry : roles.entrySet()) {
+        addRoles(result, table, entry.getKey(), entry.getValue());
+      }
+    } else if (dataSource instanceof RestrictedDataSource) {
+      attribute(result, ((RestrictedDataSource) dataSource).getBase(), roles);
+    } else if (dataSource instanceof FilteredDataSource) {
+      attribute(result, ((FilteredDataSource) dataSource).getBase(), roles);
+    } else if (dataSource instanceof UnnestDataSource) {
+      UnnestDataSource unnest = (UnnestDataSource) dataSource;
+      VirtualColumn unnestColumn = unnest.getVirtualColumn();
+      Map<String, EnumSet<ColumnRole>> baseRoles = new TreeMap<>(roles);
+      // The unnest output column is synthetic (not a base column); drop it and instead record the
+      // underlying column(s) being unnested as projected from the base.
+      baseRoles.remove(unnestColumn.getOutputName());
+      for (String required : unnestColumn.requiredColumns()) {
+        baseRoles.computeIfAbsent(required, k -> EnumSet.noneOf(ColumnRole.class)).add(ColumnRole.PROJECTION);
+      }
+      attribute(result, unnest.getBase(), baseRoles);
+    } else if (dataSource instanceof UnionDataSource) {
+      // Union members share the same output schema, so the referenced columns apply to each.
+      for (DataSource member : dataSource.getChildren()) {
+        attribute(result, member, roles);
+      }
+    } else if (dataSource instanceof QueryDataSource) {
+      // The outer references are this sub-query's OUTPUT columns, not base columns; the sub-query's
+      // own base-table columns are captured by recursing into its parts.
+      collectInto(result, ((QueryDataSource) dataSource).getQuery());
+    } else if (dataSource instanceof JoinDataSource) {
+      attributeJoin(result, (JoinDataSource) dataSource, roles);
+    }
+    // LookupDataSource, InlineDataSource and any other shape have no base table: drop (never fabricate).
+  }
+
+  /**
+   * Splits {@code roles} (plus the join-condition columns, tagged {@link ColumnRole#JOIN}) across the
+   * base datasource and each joinable clause using the planner's join prefixes, then recurses. Clauses
+   * are matched longest-prefix-first; right-side columns arrive already prefixed and are un-prefixed
+   * before attribution to the clause's datasource (which may itself be a table, join, or sub-query).
+   */
+  private void attributeJoin(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      JoinDataSource join,
+      Map<String, EnumSet<ColumnRole>> roles
+  )
+  {
+    JoinDataSourceAnalysis analysis = JoinDataSourceAnalysis.constructAnalysis(join);
+    DataSource base = analysis.getBaseDataSource();
+    List<PreJoinableClause> clauses = new ArrayList<>(analysis.getPreJoinableClauses());
+    // Longest prefix first so that, e.g., "j0.x" matches clause "j0." rather than a shorter prefix.
+    clauses.sort((a, b) -> Integer.compare(b.getPrefix().length(), a.getPrefix().length()));
+
+    Map<String, EnumSet<ColumnRole>> all = new TreeMap<>(roles);
+    for (PreJoinableClause clause : clauses) {
+      for (String column : clause.getCondition().getRequiredColumns()) {
+        all.computeIfAbsent(column, k -> EnumSet.noneOf(ColumnRole.class)).add(ColumnRole.JOIN);
+      }
+    }
+
+    Map<DataSource, Map<String, EnumSet<ColumnRole>>> partitioned = new LinkedHashMap<>();
+    for (Map.Entry<String, EnumSet<ColumnRole>> entry : all.entrySet()) {
+      String column = entry.getKey();
+      DataSource target = base;
+      String resolved = column;
+      for (PreJoinableClause clause : clauses) {
+        if (JoinPrefixUtils.isPrefixedBy(column, clause.getPrefix())) {
+          target = clause.getDataSource();
+          resolved = JoinPrefixUtils.unprefix(column, clause.getPrefix());
+          break;
+        }
+      }
+      partitioned.computeIfAbsent(target, k -> new TreeMap<>())
+                 .computeIfAbsent(resolved, k -> EnumSet.noneOf(ColumnRole.class))
+                 .addAll(entry.getValue());
+    }
+    for (Map.Entry<DataSource, Map<String, EnumSet<ColumnRole>>> entry : partitioned.entrySet()) {
+      attribute(result, entry.getKey(), entry.getValue());
+    }
+  }
+
+  private static void addRoles(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      String table,
+      String column,
+      EnumSet<ColumnRole> roles
+  )
+  {
+    result.computeIfAbsent(table, k -> new TreeMap<>())
+          .computeIfAbsent(column, k -> EnumSet.noneOf(ColumnRole.class))
+          .addAll(roles);
+  }
+
+  /**
+   * Walks the query's column-bearing parts (projected columns, dimensions, aggregator inputs, filter
+   * columns) and records each referenced base column with the role(s) it was used in. Virtual columns
+   * are expanded transitively to their underlying base columns, carrying the consuming role. Only
+   * explicitly-referenced columns are captured (notably {@code __time} is included only when it
+   * actually appears in a part, never its implicit interval usage). Returns an empty map for
+   * unsupported query types and for a bare {@code SELECT *} (a Scan with neither explicit columns nor
+   * a filter); a {@code SELECT *} that carries a filter still contributes its filter columns.
+   */
+  private Map<String, EnumSet<ColumnRole>> collectColumnRoles(Query<?> query)
+  {
+    Map<String, EnumSet<ColumnRole>> roles = new TreeMap<>();
+    if (query instanceof ScanQuery) {
+      ScanQuery scan = (ScanQuery) query;
+      VirtualColumns vcs = scan.getVirtualColumns();
+      if (scan.getColumns() != null) {
+        for (String column : scan.getColumns()) {
+          addColumn(roles, vcs, column, ColumnRole.PROJECTION);
+        }
+      }
+      addFilterColumns(roles, vcs, scan.getFilter());
+    } else if (query instanceof GroupByQuery) {
+      GroupByQuery groupBy = (GroupByQuery) query;
+      VirtualColumns vcs = groupBy.getVirtualColumns();
+      for (DimensionSpec dimension : groupBy.getDimensions()) {
+        addColumn(roles, vcs, dimension.getDimension(), ColumnRole.GROUP_BY);
+      }
+      addAggregatorColumns(roles, vcs, groupBy.getAggregatorSpecs());
+      addFilterColumns(roles, vcs, groupBy.getDimFilter());
+    } else if (query instanceof TopNQuery) {
+      TopNQuery topN = (TopNQuery) query;
+      VirtualColumns vcs = topN.getVirtualColumns();
+      addColumn(roles, vcs, topN.getDimensionSpec().getDimension(), ColumnRole.GROUP_BY);
+      addAggregatorColumns(roles, vcs, topN.getAggregatorSpecs());
+      addFilterColumns(roles, vcs, topN.getFilter());
+    } else if (query instanceof TimeseriesQuery) {
+      TimeseriesQuery timeseries = (TimeseriesQuery) query;
+      VirtualColumns vcs = timeseries.getVirtualColumns();
+      addAggregatorColumns(roles, vcs, timeseries.getAggregatorSpecs());
+      addFilterColumns(roles, vcs, timeseries.getFilter());
+    } else {
+      // Unsupported query type: emit table-level lineage only.
+      return Collections.emptyMap();
+    }
+    return roles;
+  }
+
+  private void addAggregatorColumns(
+      Map<String, EnumSet<ColumnRole>> roles,
+      VirtualColumns vcs,
+      List<AggregatorFactory> aggregators
+  )
+  {
+    for (AggregatorFactory aggregator : aggregators) {
+      for (String column : aggregator.requiredFields()) {

Review Comment:
   [P2] Filtered aggregator predicates are labeled as aggregations
   
   FilteredAggregatorFactory.requiredFields() includes both the delegate aggregator inputs and the filter predicate columns, but this loop tags every required field as AGGREGATION. For an aggregate such as SUM(cnt) FILTER (WHERE dim2 = 'a'), dim2 is a filter predicate and should receive FILTER, not AGGREGATION.



##########
extensions-contrib/openlineage-emitter/src/main/java/org/apache/druid/extensions/openlineage/OpenLineageRequestLogger.java:
##########
@@ -309,11 +347,275 @@ static String extractMsqOutputDatasource(String sql)
     }
   }
 
+  /**
+   * Column usage roles for the custom {@code druid_columnUsage} dataset facet. Names align with
+   * OpenLineage transformation subtypes where they exist (GROUP_BY, AGGREGATION, FILTER, JOIN);
+   * PROJECTION corresponds to OpenLineage IDENTITY/TRANSFORMATION (a selected/projected column).
+   */
+  enum ColumnRole
+  {
+    PROJECTION, GROUP_BY, AGGREGATION, FILTER, JOIN
+  }
+
+  /**
+   * Resolves the per-base-table column-usage map for a native query, used to attach the {@code schema}
+   * and {@code druid_columnUsage} dataset facets to input datasets. Handles all datasource shapes:
+   * tables, joins (attributing columns per side via the planner's join prefixes), unions, sub-queries
+   * (recursing into the sub-query's own columns), and datasource wrappers (restricted/filtered/unnest).
+   * Columns that have no base table (lookups, inline data) are dropped rather than fabricated.
+   *
+   * <p>Returns {@code null} (yielding table-level lineage only) when no base-table columns can be
+   * determined, or on any error -- it never fabricates or mis-attributes columns.
+   */
+  @Nullable
+  private Map<String, Map<String, EnumSet<ColumnRole>>> extractColumnsByTable(Query<?> query)
+  {
+    try {
+      Map<String, Map<String, EnumSet<ColumnRole>>> result = new TreeMap<>();
+      collectInto(result, query);
+      return result.isEmpty() ? null : result;
+    }
+    // StackOverflowError (an Error, not an Exception) is caught too so that a pathologically deep
+    // query plan degrades to table-level lineage rather than breaking the request-logging path.
+    catch (Exception | StackOverflowError e) {
+      log.debug(e, "Failed to extract column lineage; falling back to table-level lineage");
+      return null;
+    }
+  }
+
+  /**
+   * Walks {@code query}'s column-bearing parts and attributes the referenced columns to the base
+   * tables of its datasource. {@link UnionQuery} (whose {@code getDataSource()} is undefined) is
+   * handled by recursing into each of its branches.
+   */
+  private void collectInto(Map<String, Map<String, EnumSet<ColumnRole>>> result, Query<?> query)
+  {
+    if (query instanceof UnionQuery) {
+      for (DataSource branch : ((UnionQuery) query).getDataSources()) {
+        attribute(result, branch, Collections.emptyMap());
+      }
+      return;
+    }
+    attribute(result, query.getDataSource(), collectColumnRoles(query));
+  }
+
+  /**
+   * Attributes {@code roles} (column to roles, expressed in {@code dataSource}'s output namespace) to
+   * the underlying base tables, recursing through the datasource tree. Sub-queries are recursed into
+   * via {@link #collectInto} (their columns come from their own parts, not the outer references).
+   */
+  private void attribute(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      DataSource dataSource,
+      Map<String, EnumSet<ColumnRole>> roles
+  )
+  {
+    if (dataSource instanceof TableDataSource) {
+      // Also covers GlobalTableDataSource (a TableDataSource subclass); it is a real named table, so
+      // handling it here is intentional. Keep this branch ahead of the wrapper branches below.
+      String table = ((TableDataSource) dataSource).getName();
+      for (Map.Entry<String, EnumSet<ColumnRole>> entry : roles.entrySet()) {
+        addRoles(result, table, entry.getKey(), entry.getValue());
+      }
+    } else if (dataSource instanceof RestrictedDataSource) {
+      attribute(result, ((RestrictedDataSource) dataSource).getBase(), roles);
+    } else if (dataSource instanceof FilteredDataSource) {
+      attribute(result, ((FilteredDataSource) dataSource).getBase(), roles);
+    } else if (dataSource instanceof UnnestDataSource) {
+      UnnestDataSource unnest = (UnnestDataSource) dataSource;
+      VirtualColumn unnestColumn = unnest.getVirtualColumn();
+      Map<String, EnumSet<ColumnRole>> baseRoles = new TreeMap<>(roles);
+      // The unnest output column is synthetic (not a base column); drop it and instead record the
+      // underlying column(s) being unnested as projected from the base.
+      baseRoles.remove(unnestColumn.getOutputName());
+      for (String required : unnestColumn.requiredColumns()) {
+        baseRoles.computeIfAbsent(required, k -> EnumSet.noneOf(ColumnRole.class)).add(ColumnRole.PROJECTION);
+      }
+      attribute(result, unnest.getBase(), baseRoles);
+    } else if (dataSource instanceof UnionDataSource) {
+      // Union members share the same output schema, so the referenced columns apply to each.
+      for (DataSource member : dataSource.getChildren()) {
+        attribute(result, member, roles);
+      }
+    } else if (dataSource instanceof QueryDataSource) {
+      // The outer references are this sub-query's OUTPUT columns, not base columns; the sub-query's
+      // own base-table columns are captured by recursing into its parts.
+      collectInto(result, ((QueryDataSource) dataSource).getQuery());
+    } else if (dataSource instanceof JoinDataSource) {
+      attributeJoin(result, (JoinDataSource) dataSource, roles);
+    }
+    // LookupDataSource, InlineDataSource and any other shape have no base table: drop (never fabricate).
+  }
+
+  /**
+   * Splits {@code roles} (plus the join-condition columns, tagged {@link ColumnRole#JOIN}) across the
+   * base datasource and each joinable clause using the planner's join prefixes, then recurses. Clauses
+   * are matched longest-prefix-first; right-side columns arrive already prefixed and are un-prefixed
+   * before attribution to the clause's datasource (which may itself be a table, join, or sub-query).
+   */
+  private void attributeJoin(
+      Map<String, Map<String, EnumSet<ColumnRole>>> result,
+      JoinDataSource join,
+      Map<String, EnumSet<ColumnRole>> roles
+  )
+  {
+    JoinDataSourceAnalysis analysis = JoinDataSourceAnalysis.constructAnalysis(join);
+    DataSource base = analysis.getBaseDataSource();
+    List<PreJoinableClause> clauses = new ArrayList<>(analysis.getPreJoinableClauses());
+    // Longest prefix first so that, e.g., "j0.x" matches clause "j0." rather than a shorter prefix.
+    clauses.sort((a, b) -> Integer.compare(b.getPrefix().length(), a.getPrefix().length()));
+
+    Map<String, EnumSet<ColumnRole>> all = new TreeMap<>(roles);

Review Comment:
   [P2] Join left filters are dropped
   
   JoinDataSourceAnalysis exposes the base-table leftFilter through getJoinBaseTableFilter(), but attributeJoin only adds join-condition columns before partitioning roles. SQL plans that push a predicate into JoinDataSource.leftFilter will under-report that filtered base-table column in the emitted lineage.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org