Line data Source code
1 : //! Reading Trino plans into the query profile.
2 : //!
3 : //! Trino plans are a map from fragment id to the plan tree of the fragment, the JSON format of
4 : //! `EXPLAIN (TYPE DISTRIBUTED, FORMAT JSON)`. QueryInfo (`/v1/query/{id}`) carries the same JSON per stage
5 : //! (`stages.stages[].plan.jsonRepresentation`), both are read. Every fragment is a stage at runtime, the stage id
6 : //! `<queryId>.<fragmentId>` ends with the fragment id.
7 : //!
8 : //! A `RemoteSource` (or `RemoteMerge`) node reads the output of the fragments in `sourceFragmentIds`. The parser
9 : //! hangs the roots of those fragments below it, so the plan is one tree and the fragments nest.
10 :
11 : use super::query_profile::{Fragment, OperatorNode, ParsedOperatorNode, PathComponent, QueryProfile};
12 : use dashql_buffers::profile as fb;
13 : use dashql_buffers::status::StatusCode;
14 : use serde_json::{Map, Value};
15 : use std::collections::{BTreeMap, HashMap, HashSet};
16 :
17 : /// The engine-neutral kind of a Trino plan node.
18 : ///
19 : /// The names are the ones of Trino's `PlanPrinter`. It merges a filter and a projection into one node
20 : /// (`FilterProject`, `ScanFilterProject`, ...), and names joins by their type (`LeftJoin`, `SpatialInnerJoin`,
21 : /// `SourceOuterIndexJoin`). `Intersect` and `Except` keep all duplicates unless the descriptor says `isDistinct`,
22 : /// which Trino prints as `"true"` or `""`.
23 : ///
24 : /// A `LocalExchange` is classified by the `partitioning` of its descriptor: `SINGLE` and `COORDINATOR_ONLY` gather,
25 : /// `BROADCAST` broadcasts, `UNKNOWN` (passthrough, arbitrary) stays a plain exchange, and every other function (`HASH`,
26 : /// `ROUND_ROBIN`, a connector partitioning) repartitions. A `RemoteSource` or `RemoteMerge` reads other stages over
27 : /// the network.
28 111 : fn trino_operator_kind(op: &OperatorNode) -> fb::ProfileOperatorKind {
29 : use fb::ProfileOperatorKind as K;
30 111 : let Some(name) = op.operator_type.as_deref() else {
31 0 : return K::OTHER;
32 : };
33 111 : let descriptor = |key: &str| {
34 57 : let (_, json) = op.attributes.iter().find(|(attribute, _)| attribute == "descriptor")?;
35 19 : serde_json::from_str::<Value>(json).ok()?.get(key).cloned()
36 19 : };
37 111 : let distinct = || match descriptor("isDistinct") {
38 2 : Some(Value::String(s)) => s == "true",
39 0 : Some(Value::Bool(b)) => b,
40 0 : _ => true,
41 2 : };
42 111 : let local_exchange = || match descriptor("partitioning").as_ref().and_then(Value::as_str) {
43 16 : None | Some("UNKNOWN") => K::EXCHANGE,
44 16 : Some("SINGLE" | "COORDINATOR_ONLY") => K::GATHER_EXCHANGE,
45 7 : Some("BROADCAST") => K::BROADCAST_EXCHANGE,
46 6 : Some(_) => K::REPARTITION_EXCHANGE,
47 17 : };
48 1 : match name {
49 111 : "Output" => K::OUTPUT,
50 100 : "TableScan" | "ScanFilter" | "ScanProject" | "ScanFilterProject" | "IndexSource" => K::SCAN,
51 82 : "Values" => K::VIRTUAL_TABLE,
52 78 : "TableFunction" | "TableFunctionProcessor" => K::TABLE_FUNCTION,
53 78 : "Filter" | "FilterProject" => K::FILTER,
54 77 : "Project" => K::PROJECT,
55 75 : "Aggregate" | "MarkDistinct" => K::AGGREGATE,
56 67 : "Window" | "TopNRanking" | "RowNumber" | "PatternRecognition" => K::WINDOW,
57 66 : "Sort" | "PartialSort" | "TopN" | "TopNPartial" | "LocalMerge" => K::SORT,
58 54 : "Limit" | "LimitPartial" | "DistinctLimit" | "DistinctLimitPartial" | "Offset" => K::LIMIT,
59 53 : "InnerJoin" | "CrossJoin" | "SpatialInnerJoin" | "InnerIndexJoin" => K::JOIN,
60 44 : "LeftJoin" | "SpatialLeftJoin" | "SourceOuterIndexJoin" => K::LEFT_OUTER_JOIN,
61 42 : "RightJoin" => K::RIGHT_OUTER_JOIN,
62 42 : "FullJoin" => K::FULL_OUTER_JOIN,
63 41 : "SemiJoin" => K::LEFT_SEMI_JOIN,
64 41 : "Union" => K::UNION_ALL,
65 40 : "Intersect" if distinct() => K::INTERSECT,
66 0 : "Intersect" => K::INTERSECT_ALL,
67 39 : "Except" if distinct() => K::EXCEPT,
68 1 : "Except" => K::EXCEPT_ALL,
69 38 : "TableWriter" | "TableWriterPartial" | "TableCommit" | "TableFinish" | "MergeWriter" | "TableExecute" => {
70 1 : K::INSERT
71 : }
72 37 : "LocalExchange" => local_exchange(),
73 20 : "RemoteSource" | "RemoteMerge" | "RemoteExchange" => K::REMOTE_EXCHANGE,
74 1 : _ => K::OTHER,
75 : }
76 111 : }
77 :
78 : /// Serialize an attribute value as compact JSON
79 531 : fn to_json(value: &Value) -> String {
80 531 : serde_json::to_string(value).expect("a JSON value always serializes")
81 531 : }
82 :
83 : /// The fragments of a node that reads other fragments (`sourceFragmentIds: "[1, 2]"`)
84 222 : fn source_fragments(node: &Map<String, Value>) -> Vec<String> {
85 222 : let Some(ids) = node.get("descriptor").and_then(|d| d.get("sourceFragmentIds")).and_then(Value::as_str) else {
86 186 : return Vec::new();
87 : };
88 114 : let ids = ids.trim_matches(|c| c == '[' || c == ']');
89 42 : ids.split(',').map(str::trim).filter(|s| !s.is_empty()).map(String::from).collect()
90 222 : }
91 :
92 : /// The label of a node: the table of a scan (`tpch:sf1:lineitem` -> `lineitem`)
93 111 : fn label_of(node: &Map<String, Value>) -> Option<String> {
94 111 : let table = node.get("descriptor")?.get("table")?.as_str()?;
95 16 : Some(table.rsplit(':').next().unwrap_or(table).to_string())
96 111 : }
97 :
98 : /// The fragments of a Trino plan document by fragment id
99 15 : fn read_fragments(document: Value) -> Result<BTreeMap<String, Value>, StatusCode> {
100 15 : let Value::Object(mut object) = document else {
101 1 : return Err(StatusCode::VIEWMODEL_INPUT_JSON_PARSER_ERROR);
102 : };
103 : // A QueryInfo: the plans of the stages
104 14 : if let Some(stages) = object.remove("stages") {
105 5 : let mut fragments = BTreeMap::new();
106 5 : let stages = stages.get("stages").and_then(Value::as_array).cloned().unwrap_or_default();
107 11 : for stage in stages {
108 11 : let Some(stage_id) = stage.get("stageId").and_then(Value::as_str) else { continue };
109 11 : let fragment_id = stage_id.rsplit('.').next().unwrap_or(stage_id).to_string();
110 11 : let Some(plan) = stage.get("plan").and_then(|p| p.get("jsonRepresentation")).and_then(Value::as_str) else {
111 1 : continue;
112 : };
113 10 : let plan: Value = serde_json::from_str(plan).map_err(|_| StatusCode::VIEWMODEL_INPUT_JSON_PARSER_ERROR)?;
114 9 : fragments.insert(fragment_id, plan);
115 : }
116 4 : return Ok(fragments);
117 9 : }
118 : // EXPLAIN (TYPE DISTRIBUTED, FORMAT JSON): fragment id -> plan node
119 25 : if object.values().all(|v| v.get("name").is_some_and(Value::is_string)) {
120 8 : return Ok(object.into_iter().collect());
121 1 : }
122 1 : Err(StatusCode::VIEWMODEL_INPUT_JSON_PARSER_ERROR)
123 15 : }
124 :
125 :
126 : /// The ids of the fragments that a plan tree reads
127 33 : fn read_fragment_ids(root: &Value) -> Vec<String> {
128 33 : let mut ids = Vec::new();
129 33 : let mut pending = vec![root];
130 144 : while let Some(node) = pending.pop() {
131 111 : if let Some(node) = node.as_object() {
132 111 : ids.extend(source_fragments(node));
133 111 : pending.extend(node.get("children").and_then(Value::as_array).into_iter().flatten());
134 111 : }
135 : }
136 33 : ids
137 33 : }
138 :
139 : /// The parser state
140 : struct TrinoPlanReader<'a> {
141 : /// The plan trees of the fragments
142 : plans: &'a BTreeMap<String, Value>,
143 : /// The parsed operators
144 : operators: Vec<ParsedOperatorNode>,
145 : /// The fragments in the order they were read, parents first
146 : fragments: Vec<Fragment>,
147 : /// The index of every fragment that was read
148 : fragment_index: HashMap<String, u32>,
149 : }
150 :
151 : impl TrinoPlanReader<'_> {
152 : /// Read a fragment, returns its root operator. A fragment is read once, a second reference is ignored.
153 33 : fn read_fragment(&mut self, fragment_id: &str, parent: Option<u32>) -> Option<usize> {
154 33 : if self.fragment_index.contains_key(fragment_id) {
155 0 : return None;
156 33 : }
157 33 : let root = self.plans.get(fragment_id)?.as_object()?;
158 33 : let fragment = self.fragments.len() as u32;
159 33 : self.fragment_index.insert(fragment_id.to_string(), fragment);
160 33 : self.fragments.push(Fragment {
161 33 : engine_id: Some(fragment_id.to_string()),
162 33 : parent_fragment: parent,
163 33 : ..Default::default()
164 33 : });
165 33 : Some(self.read_node(root, fragment, Vec::new()))
166 33 : }
167 :
168 : /// Read a plan node and its children, children first
169 111 : fn read_node(&mut self, node: &Map<String, Value>, fragment: u32, path: Vec<PathComponent>) -> usize {
170 111 : let mut children = Vec::new();
171 111 : let child_nodes = node.get("children").and_then(Value::as_array).into_iter().flatten();
172 111 : for (i, child) in child_nodes.enumerate() {
173 78 : if let Some(child) = child.as_object() {
174 78 : let path = vec![
175 78 : PathComponent::MemberInObject { attribute: "children".into() },
176 78 : PathComponent::EntryInArray { index: i },
177 78 : ];
178 78 : children.push(self.read_node(child, fragment, path));
179 78 : }
180 : }
181 111 : for source in source_fragments(node) {
182 21 : if let Some(root) = self.read_fragment(&source, Some(fragment)) {
183 21 : self.operators[root].parent_child_path =
184 21 : vec![PathComponent::MemberInObject { attribute: format!("fragment {source}") }];
185 21 : children.push(root);
186 21 : }
187 : }
188 111 : let name = node.get("name").and_then(Value::as_str).unwrap_or("Unknown").to_string();
189 111 : let estimated_rows = node
190 111 : .get("estimates")
191 111 : .and_then(Value::as_array)
192 111 : .and_then(|e| e.first())
193 111 : .and_then(|e| e.get("outputRowCount"))
194 111 : .and_then(Value::as_f64);
195 111 : let attributes = node
196 111 : .iter()
197 624 : .filter(|(key, _)| *key != "children")
198 531 : .map(|(key, value)| (key.clone(), to_json(value)))
199 111 : .collect();
200 111 : self.operators.push(ParsedOperatorNode {
201 111 : parent_child_path: path,
202 111 : attributes,
203 111 : estimated_rows,
204 111 : operator_label: label_of(node),
205 111 : operator_type: Some(name),
206 111 : engine_id: node.get("id").and_then(Value::as_str).map(String::from),
207 111 : fragment: Some(fragment),
208 111 : child_operators: children,
209 111 : ..Default::default()
210 111 : });
211 111 : self.operators.len() - 1
212 111 : }
213 : }
214 :
215 : impl QueryProfile {
216 : /// Parse a Trino plan: the fragments of `EXPLAIN (TYPE DISTRIBUTED, FORMAT JSON)` or a QueryInfo with the plans
217 : /// of its stages. Fragment ids are the engine ids of the fragments and plan node ids the engine ids of the
218 : /// operators, so the progress of the stages finds its fragments.
219 16 : pub fn parse_trino_plan(&mut self, plan: &[u8]) -> Result<(), StatusCode> {
220 16 : self.reset();
221 16 : self.engine = fb::ProfileEngine::TRINO;
222 15 : let document: Value =
223 16 : serde_json::from_slice(plan).map_err(|_| StatusCode::VIEWMODEL_INPUT_JSON_PARSER_ERROR)?;
224 15 : let plans = read_fragments(document)?;
225 :
226 : // The root fragments are the fragments that no other fragment reads, usually fragment 0
227 12 : let read: HashSet<String> = plans.values().flat_map(read_fragment_ids).collect();
228 12 : let mut reader = TrinoPlanReader {
229 12 : plans: &plans,
230 12 : operators: Vec::new(),
231 12 : fragments: Vec::new(),
232 12 : fragment_index: HashMap::new(),
233 12 : };
234 12 : let mut roots = Vec::new();
235 33 : for fragment_id in plans.keys().filter(|id| !read.contains(*id)) {
236 12 : roots.extend(reader.read_fragment(fragment_id, None));
237 12 : }
238 12 : let TrinoPlanReader { operators, fragments, .. } = reader;
239 12 : self.flatten_operators(operators, roots);
240 12 : self.fragments = fragments;
241 :
242 : // The operators of the fragments; the root of a fragment has no parent in the same fragment
243 111 : for (op_id, op) in self.operators.iter_mut().enumerate() {
244 111 : op.kind = trino_operator_kind(op);
245 111 : if let Some(fragment) = op.fragment {
246 111 : self.fragments[fragment as usize].operators.push(op_id as u32);
247 111 : }
248 : }
249 111 : for (op_id, op) in self.operators.iter().enumerate() {
250 111 : let Some(fragment) = op.fragment else { continue };
251 111 : if op.parent.is_none_or(|parent| self.operators[parent].fragment != Some(fragment)) {
252 33 : self.fragments[fragment as usize].root_operator = Some(op_id as u32);
253 78 : }
254 : }
255 :
256 12 : self.identify_operator_edges();
257 12 : self.seal_loaded_structure();
258 12 : Ok(())
259 16 : }
260 : }
|
