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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
