FrankChen021 commented on code in PR #19643: URL: https://github.com/apache/druid/pull/19643#discussion_r3520194723
########## processing/src/main/java/org/apache/druid/query/QueryColumnUsageAnalyzer.java: ########## @@ -0,0 +1,365 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.query; + +import org.apache.druid.query.aggregation.AggregatorFactory; +import org.apache.druid.query.aggregation.FilteredAggregatorFactory; +import org.apache.druid.query.dimension.DimensionSpec; +import org.apache.druid.query.filter.DimFilter; +import org.apache.druid.query.groupby.GroupByQuery; +import org.apache.druid.query.planning.JoinDataSourceAnalysis; +import org.apache.druid.query.planning.PreJoinableClause; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.query.timeseries.TimeseriesQuery; +import org.apache.druid.query.topn.TopNQuery; +import org.apache.druid.query.union.UnionQuery; +import org.apache.druid.segment.VirtualColumn; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.join.JoinPrefixUtils; + +import javax.annotation.Nullable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.EnumSet; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; + +/** + * Analyzes a native {@link Query} to determine, per base table, which columns it references and in + * what role (projection, filter, group-by, aggregation, join key). Columns are read from the query + * plan -- not by parsing SQL -- and attributed to base tables across joins (using the planner's join + * prefixes), unions, sub-queries, and datasource wrappers (restricted/filtered/unnest). Columns with + * no base table (for example lookups or inline data) are dropped rather than fabricated. + * + * <p>This is intended as a single, planner-aligned source of truth for "which columns does this query + * read, and how" so that lineage emitters, column-statistics/optimizer work, and other consumers do + * not each re-derive it and drift from the parser/planner. + * + * <p>It is complementary to {@link Query#getRequiredColumns()}: that method returns a flat, + * un-attributed set, may be {@code null}, unconditionally injects {@code __time}, does not expand + * virtual columns to their base columns, and (by contract) excludes join-condition and pushed-down + * filter columns. This analyzer attributes per base table, assigns usage roles, expands virtual + * columns, and includes join-condition / {@code leftFilter} / filtered-aggregator predicate columns. + */ +public final class QueryColumnUsageAnalyzer +{ + private QueryColumnUsageAnalyzer() + { + } + + /** + * The role in which a column is referenced by a query. Names align with OpenLineage transformation + * subtypes where they exist (GROUP_BY, AGGREGATION, FILTER, JOIN); PROJECTION denotes a + * selected/projected column. + */ + public enum ColumnUsage + { + PROJECTION, + GROUP_BY, + AGGREGATION, + FILTER, + JOIN + } + + /** + * Returns a map of base-table name to (column name to the set of roles it was used in) for the given + * query. Returns an empty map when no base-table columns can be determined (unsupported query type, + * {@code SELECT *} with no filter, or a datasource with no base tables). Never returns {@code null}. + * + * <p>Result maps are sorted by table and column, and role sets iterate in enum-declaration order, so + * the output is deterministic. This method reads the query structure only and performs no I/O. The + * returned maps and role sets are freshly allocated, mutable, and owned by the caller. + */ + public static Map<String, Map<String, EnumSet<ColumnUsage>>> analyze(Query<?> query) + { + Map<String, Map<String, EnumSet<ColumnUsage>>> result = new TreeMap<>(); + collectInto(result, query); + return result; + } + + /** + * 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 branch queries -- {@code getDataSources()} would discard the + * branches' dimensions/aggregators/filters, so the full branch queries are used instead. + */ + private static void collectInto(Map<String, Map<String, EnumSet<ColumnUsage>>> result, Query<?> query) + { + if (query instanceof UnionQuery) { + for (Query<?> branch : ((UnionQuery) query).getQueries()) { + collectInto(result, branch); + } + 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 static void attribute( + Map<String, Map<String, EnumSet<ColumnUsage>>> result, + DataSource dataSource, + Map<String, EnumSet<ColumnUsage>> 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<ColumnUsage>> 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) { + // FilteredDataSource carries a pushed-down filter over its base; its columns are FILTER usages. + FilteredDataSource filtered = (FilteredDataSource) dataSource; + Map<String, EnumSet<ColumnUsage>> baseRoles = copyRoles(roles); + addFilterColumns(baseRoles, null, filtered.getFilter()); + attribute(result, filtered.getBase(), baseRoles); + } else if (dataSource instanceof UnnestDataSource) { + UnnestDataSource unnest = (UnnestDataSource) dataSource; + VirtualColumn unnestColumn = unnest.getVirtualColumn(); + Map<String, EnumSet<ColumnUsage>> baseRoles = copyRoles(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()); Review Comment: [P2] Preserve unnest output roles For `UnnestDataSource`, roles collected on the synthetic unnest output are discarded here and the source columns are re-added only as `PROJECTION`; `getUnnestFilter()` is also ignored. Queries such as `SELECT d3, COUNT(*) ... WHERE d3 = 'a' GROUP BY d3` will report the underlying array column as projected instead of `GROUP_BY` and `FILTER`. Preserve the removed role set and merge unnest-filter required columns as `FILTER` when mapping to `unnestColumn.requiredColumns()`. ########## extensions-contrib/openlineage-emitter/src/main/java/org/apache/druid/extensions/openlineage/OpenLineageRequestLogger.java: ########## @@ -222,7 +234,9 @@ public void logNativeQuery(RequestLogLine requestLogLine) throws IOException queryId = UNKNOWN_QUERY_ID; } - emit(buildRunEvent(queryId, queryType, requestLogLine, inputs, null)); + Map<String, Map<String, EnumSet<QueryColumnUsageAnalyzer.ColumnUsage>>> columnsByTable = + columnLineageEnabled ? extractColumnsByTable(requestLogLine.getQuery()) : null; Review Comment: [P2] Handle top-level UnionQuery inputs A native top-level `UnionQuery` has `getDataSource()` deliberately throwing, but `logNativeQuery` always calls `getDataSource().getTableNames()` before the new `QueryColumnUsageAnalyzer` can handle `UnionQuery`. SQL plans with `queryType` `union` will fail request logging and emit no OpenLineage event. Special-case `UnionQuery` input table extraction, or derive inputs from the analyzer or branch datasources, and add a logger-level `UnionQuery` test. -- 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]
