Line data Source code
1 : //! Unit tests of reading Trino plans into the query profile.
2 :
3 : use super::QueryProfile;
4 : use dashql_buffers::profile as fb;
5 : use dashql_buffers::status::StatusCode;
6 :
7 3 : fn parse(plan: &str) -> Result<Vec<u8>, StatusCode> {
8 3 : let mut profile = QueryProfile::new();
9 3 : profile.parse_trino_plan(plan.as_bytes())?;
10 3 : profile.compute_layout();
11 3 : Ok(profile.pack_buffer())
12 3 : }
13 :
14 4 : fn root(buffer: &[u8]) -> fb::QueryProfile<'_> {
15 4 : flatbuffers::root::<fb::QueryProfile>(buffer).expect("valid query profile")
16 4 : }
17 :
18 200 : fn string<'a>(profile: &fb::QueryProfile<'a>, id: u32) -> &'a str {
19 200 : profile.string_dictionary().unwrap().get(id as usize)
20 200 : }
21 :
22 : /// A join of two fragments below a coordinator fragment
23 : const JOIN: &str = r#"{
24 : "0": {"id": "9", "name": "Output", "children": [
25 : {"id": "8", "name": "RemoteSource", "descriptor": {"sourceFragmentIds": "[1]"}, "children": []}]},
26 : "1": {"id": "7", "name": "InnerJoin", "estimates": [{"outputRowCount": "NaN"}], "children": [
27 : {"id": "5", "name": "ScanFilter", "descriptor": {"table": "tpch:sf1:orders"}, "children": []},
28 : {"id": "6", "name": "RemoteSource", "descriptor": {"sourceFragmentIds": "[2, 3]"}, "children": []}]},
29 : "2": {"id": "2", "name": "TableScan", "descriptor": {"table": "tpch:sf1:customer"},
30 : "estimates": [{"outputRowCount": 150000.0}], "children": []},
31 : "3": {"id": "3", "name": "Values", "children": []}
32 : }"#;
33 :
34 : #[test]
35 1 : fn stitches_fragments_into_one_tree() {
36 1 : let buffer = parse(JOIN).unwrap();
37 1 : let profile = root(&buffer);
38 1 : assert_eq!(profile.engine(), fb::ProfileEngine::TRINO);
39 1 : assert_eq!(profile.root_operators().unwrap().len(), 1);
40 1 : let operators = profile.operators().unwrap();
41 1 : assert_eq!(operators.len(), 7);
42 1 : assert_eq!(profile.operator_edges().unwrap().len(), 6);
43 :
44 : // The fragments nest along the remote sources, every fragment draws a contour around its own operators
45 1 : let fragments = profile.fragments().unwrap();
46 1 : let described: Vec<(&str, u32, u8)> =
47 4 : fragments.iter().map(|f| (string(&profile, f.engine_id()), f.parent_fragment(), f.hidden())).collect();
48 1 : assert_eq!(described, vec![("0", u32::MAX, 0), ("1", 0, 0), ("2", 1, 0), ("3", 1, 0)]);
49 4 : for (f, fragment) in fragments.iter().enumerate() {
50 4 : let root = operators.get(fragment.root_operator() as usize);
51 4 : assert_eq!(root.fragment(), f as u32);
52 : }
53 :
54 : // Kinds, labels and estimates
55 23 : let by_engine_id = |id: &str| operators.iter().find(|op| string(&profile, op.engine_id()) == id).unwrap();
56 1 : let scan = by_engine_id("2");
57 1 : let described = (scan.kind(), string(&profile, scan.label()), scan.estimated_rows());
58 1 : assert_eq!(described, (fb::ProfileOperatorKind::SCAN, "customer", 150000.0));
59 1 : assert_eq!(by_engine_id("7").kind(), fb::ProfileOperatorKind::JOIN);
60 1 : assert!(by_engine_id("7").estimated_rows().is_nan());
61 1 : assert_eq!(by_engine_id("8").kind(), fb::ProfileOperatorKind::REMOTE_EXCHANGE);
62 1 : assert_eq!(by_engine_id("3").kind(), fb::ProfileOperatorKind::VIRTUAL_TABLE);
63 : // The remote source reads both fragments
64 1 : assert_eq!(by_engine_id("6").children_count(), 2);
65 1 : }
66 :
67 : #[test]
68 1 : fn maps_the_node_names_of_the_plan_printer_to_kinds() {
69 : use fb::ProfileOperatorKind as K;
70 1 : let names = [
71 1 : ("FilterProject", "{}", K::FILTER),
72 1 : ("IndexSource", "{}", K::SCAN),
73 1 : ("Offset", "{}", K::LIMIT),
74 1 : ("PatternRecognition", "{}", K::WINDOW),
75 1 : ("SpatialInnerJoin", "{}", K::JOIN),
76 1 : ("InnerIndexJoin", "{}", K::JOIN),
77 1 : ("SpatialLeftJoin", "{}", K::LEFT_OUTER_JOIN),
78 1 : ("SourceOuterIndexJoin", "{}", K::LEFT_OUTER_JOIN),
79 1 : ("FullJoin", "{}", K::FULL_OUTER_JOIN),
80 1 : ("Intersect", r#"{"isDistinct": "true"}"#, K::INTERSECT),
81 1 : ("Except", r#"{"isDistinct": ""}"#, K::EXCEPT_ALL),
82 1 : ("TableCommit", "{}", K::INSERT),
83 1 : ("LocalExchange", "{}", K::EXCHANGE),
84 1 : ("LocalExchange", r#"{"partitioning": "SINGLE"}"#, K::GATHER_EXCHANGE),
85 1 : ("LocalExchange", r#"{"partitioning": "ROUND_ROBIN"}"#, K::REPARTITION_EXCHANGE),
86 1 : ("LocalExchange", r#"{"partitioning": "BROADCAST"}"#, K::BROADCAST_EXCHANGE),
87 1 : ("RemoteMerge", "{}", K::REMOTE_EXCHANGE),
88 1 : ("GroupId", "{}", K::OTHER),
89 1 : ];
90 18 : let node = |(i, (name, descriptor, _)): (usize, &(&str, &str, K))| {
91 18 : format!(r#"{{"id": "{i}", "name": "{name}", "descriptor": {descriptor}}}"#)
92 18 : };
93 1 : let children: Vec<String> = names.iter().enumerate().map(node).collect();
94 1 : let plan = format!(r#"{{"0": {{"id": "root", "name": "Union", "children": [{}]}}}}"#, children.join(","));
95 1 : let buffer = parse(&plan).unwrap();
96 1 : let profile = root(&buffer);
97 1 : let operators = profile.operators().unwrap();
98 18 : for (i, (name, descriptor, kind)) in names.into_iter().enumerate() {
99 171 : let op = operators.iter().find(|op| string(&profile, op.engine_id()) == i.to_string()).unwrap();
100 18 : assert_eq!((name, descriptor, op.kind()), (name, descriptor, kind));
101 : }
102 1 : }
103 :
104 : #[test]
105 1 : fn reads_the_stage_plans_of_a_query_info() {
106 2 : let stage = |id: &str, plan: &str| {
107 2 : serde_json::json!({"stageId": format!("q.{id}"), "plan": {"jsonRepresentation": plan}})
108 2 : };
109 1 : let info = serde_json::json!({"queryId": "q", "stages": {"stages": [
110 1 : stage("0", r#"{"id": "1", "name": "Output", "children": [
111 1 : {"id": "2", "name": "RemoteSource", "descriptor": {"sourceFragmentIds": "[1]"}, "children": []}]}"#),
112 1 : stage("1", r#"{"id": "3", "name": "TableScan", "children": []}"#),
113 : // A pruned stage has no plan
114 1 : serde_json::json!({"stageId": "q.2"}),
115 : ]}});
116 1 : let buffer = parse(&info.to_string()).unwrap();
117 1 : let profile = root(&buffer);
118 1 : assert_eq!(profile.operators().unwrap().len(), 3);
119 1 : assert_eq!(profile.fragments().unwrap().len(), 2);
120 1 : }
121 :
122 : #[test]
123 1 : fn rejects_documents_that_are_no_trino_plans() {
124 1 : let broken_stage = r#"{"stages":{"stages":[{"stageId":"q.0","plan":{"jsonRepresentation":"{"}}]}}"#;
125 4 : for plan in ["", "[]", r#"{"operator":"scan"}"#, broken_stage] {
126 4 : let mut profile = QueryProfile::new();
127 4 : let parsed = profile.parse_trino_plan(plan.as_bytes());
128 4 : assert_eq!(parsed, Err(StatusCode::VIEWMODEL_INPUT_JSON_PARSER_ERROR), "{plan}");
129 : }
130 1 : }
131 :
132 : #[test]
133 1 : fn stage_progress_finds_the_fragments_of_the_plan() {
134 1 : let mut profile = QueryProfile::new();
135 1 : profile.parse_trino_plan(JOIN.as_bytes()).unwrap();
136 1 : let version = profile.structure_version();
137 1 : let stats = br#"{"state":"RUNNING","rootStage":{"stageId":"0","state":"RUNNING","subStages":[
138 1 : {"stageId":"1","state":"RUNNING","subStages":[
139 1 : {"stageId":"2","state":"FINISHED"},{"stageId":"3","state":"FINISHED"}]}]}}"#;
140 1 : let updates = profile.apply_trino_statement_stats(stats, 1).unwrap();
141 1 : assert_eq!(profile.structure_version(), version);
142 7 : let fragment = |index: u32| updates.iter().find(|u| u.entity == fb::ProfileEntity::FRAGMENT && u.index == index);
143 1 : assert_eq!(fragment(2).unwrap().state.status, super::Status::Succeeded);
144 1 : assert_eq!(fragment(1).unwrap().state.status, super::Status::Running);
145 1 : }
146 :
147 : #[test]
148 1 : fn operator_summaries_register_pipelines_and_merge_operators_of_a_plan_node() {
149 1 : let mut profile = QueryProfile::new();
150 1 : profile.parse_trino_plan(JOIN.as_bytes()).unwrap();
151 4 : let summary = |pipeline: u32, operator: u32, node: &str, input: u64, output: u64| {
152 4 : serde_json::json!({
153 4 : "pipelineId": pipeline, "operatorId": operator, "planNodeId": node,
154 4 : "inputPositions": input, "outputPositions": output,
155 4 : "inputDataSize": "1.5kB", "outputDataSize": "10B",
156 4 : "addInputCpu": "1.50ms", "getOutputCpu": "500.00us", "finishCpu": "0.00ns",
157 4 : "blockedWall": "2.00s", "peakUserMemoryReservation": "1MB"
158 : })
159 4 : };
160 1 : let info = serde_json::json!({"stages": {"stages": [{
161 1 : "stageId": "q.1", "state": "FINISHED", "stageStats": {"operatorSummaries": [
162 : // Pipeline 0: the scan feeds the join, the join is two Trino operators (probe, output)
163 1 : summary(0, 2, "7", 40, 4),
164 1 : summary(0, 0, "5", 100, 50),
165 1 : summary(0, 1, "7", 50, 40),
166 : // An operator of a plan node that the plan does not know
167 1 : summary(0, 3, "999", 1, 1),
168 : ]}
169 : }]}})
170 1 : .to_string();
171 1 : let version = profile.structure_version();
172 1 : let updates = profile.apply_trino_query_info(info.as_bytes(), 5).unwrap();
173 1 : assert_eq!(profile.structure_version(), version + 1);
174 7 : let slots: Vec<_> = updates.iter().filter(|u| u.entity == fb::ProfileEntity::SLOT).collect();
175 1 : assert_eq!(slots.len(), 2);
176 1 : let join = slots[1].state;
177 : // Input of the first operator, output of the last, times summed in microseconds
178 1 : assert_eq!((join.stats.tuples_in, join.stats.tuples_out), (Some(50), Some(4)));
179 1 : assert_eq!(join.stats.cpu_time_us, Some(4000));
180 1 : assert_eq!(join.stats.blocked_time_us, Some(4_000_000));
181 1 : assert_eq!((join.stats.bytes_in, join.stats.memory_bytes), (Some(1536), Some(2 * 1024 * 1024)));
182 1 : assert_eq!(join.status, super::Status::Succeeded);
183 :
184 : // The next sample finds the pipeline
185 1 : assert!(profile.apply_trino_query_info(info.as_bytes(), 6).unwrap().is_empty());
186 1 : assert_eq!(profile.structure_version(), version + 1);
187 1 : let buffer = profile.pack_buffer();
188 1 : let packed = root(&buffer);
189 1 : let pipeline = packed.pipelines().unwrap().get(0);
190 1 : assert_eq!((string(&packed, pipeline.engine_id()), pipeline.fragment()), ("1.0", 1));
191 1 : }
|
