Skip to content

Commit f715a98

Browse files
sundy-lisundy-li
authored
feat(query): support send spill file stats to client (#17186)
1 parent c17767b commit f715a98

File tree

10 files changed

+84
-33
lines changed

10 files changed

+84
-33
lines changed

‎src/common/base/src/base/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ pub use ordered_float::OrderedFloat;
4141
pub use profiling::Profiling;
4242
pub use progress::Progress;
4343
pub use progress::ProgressValues;
44+
pub use progress::SpillProgress;
4445
pub use select::select3;
4546
pub use select::Select3Output;
4647
pub use semaphore::Semaphore;

‎src/common/base/src/base/progress.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,20 @@ impl Progress {
6262
ProgressValues { rows, bytes }
6363
}
6464
}
65+
66+
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
67+
pub struct SpillProgress {
68+
pub file_nums: usize,
69+
pub bytes: usize,
70+
}
71+
72+
impl SpillProgress {
73+
pub fn new(file_nums: usize, bytes: usize) -> Self {
74+
Self { file_nums, bytes }
75+
}
76+
77+
pub fn incr(&mut self, other: &SpillProgress) {
78+
self.file_nums += other.file_nums;
79+
self.bytes += other.bytes;
80+
}
81+
}

‎src/query/service/src/interpreters/hook/vacuum_hook.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,10 +48,10 @@ pub fn hook_vacuum_temp_files(query_ctx: &Arc<QueryContext>) -> Result<()> {
4848

4949
let mut node_files = HashMap::new();
5050
for node in cluster.nodes.iter() {
51-
let num = query_ctx.get_spill_file_nums(Some(node.id.clone()));
52-
if num != 0 {
51+
let stats = query_ctx.get_spill_file_stats(Some(node.id.clone()));
52+
if stats.file_nums != 0 {
5353
if let Some(index) = cluster.index_of_nodeid(&node.id) {
54-
node_files.insert(index, num);
54+
node_files.insert(index, stats.file_nums);
5555
}
5656
}
5757
}

‎src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -230,9 +230,9 @@ impl StatisticsSender {
230230
progress_info.push(ProgressInfo::ResultProgress(result_progress_values));
231231
}
232232

233-
let spill_file_nums = ctx.get_spill_file_nums(None);
234-
if spill_file_nums != 0 {
235-
progress_info.push(ProgressInfo::SpillTotalFileNums(spill_file_nums));
233+
let stats = ctx.get_spill_file_stats(None);
234+
if stats.file_nums != 0 {
235+
progress_info.push(ProgressInfo::SpillTotalStats(stats))
236236
}
237237
progress_info
238238
}

‎src/query/service/src/servers/flight/v1/packets/packet_data_progressinfo.rs‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use byteorder::BigEndian;
2121
use byteorder::ReadBytesExt;
2222
use byteorder::WriteBytesExt;
2323
use databend_common_base::base::ProgressValues;
24+
use databend_common_base::base::SpillProgress;
2425
use databend_common_exception::ErrorCode;
2526
use databend_common_exception::Result;
2627

@@ -33,7 +34,7 @@ pub enum ProgressInfo {
3334
ScanProgress(ProgressValues),
3435
WriteProgress(ProgressValues),
3536
ResultProgress(ProgressValues),
36-
SpillTotalFileNums(usize),
37+
SpillTotalStats(SpillProgress),
3738
}
3839

3940
impl ProgressInfo {
@@ -42,8 +43,8 @@ impl ProgressInfo {
4243
ProgressInfo::ScanProgress(values) => ctx.get_scan_progress().incr(values),
4344
ProgressInfo::WriteProgress(values) => ctx.get_write_progress().incr(values),
4445
ProgressInfo::ResultProgress(values) => ctx.get_result_progress().incr(values),
45-
ProgressInfo::SpillTotalFileNums(values) => {
46-
ctx.set_cluster_spill_file_nums(source_target, *values)
46+
ProgressInfo::SpillTotalStats(values) => {
47+
ctx.set_cluster_spill_progress(source_target, values.clone())
4748
}
4849
};
4950
}
@@ -53,9 +54,10 @@ impl ProgressInfo {
5354
ProgressInfo::ScanProgress(values) => (1_u8, values),
5455
ProgressInfo::WriteProgress(values) => (2_u8, values),
5556
ProgressInfo::ResultProgress(values) => (3_u8, values),
56-
ProgressInfo::SpillTotalFileNums(values) => {
57+
ProgressInfo::SpillTotalStats(values) => {
5758
bytes.write_u8(4)?;
58-
bytes.write_u64::<BigEndian>(values as u64)?;
59+
bytes.write_u64::<BigEndian>(values.file_nums as u64)?;
60+
bytes.write_u64::<BigEndian>(values.bytes as u64)?;
5961
return Ok(());
6062
}
6163
};
@@ -70,8 +72,11 @@ impl ProgressInfo {
7072
let info_type = bytes.read_u8()?;
7173

7274
if info_type == 4 {
73-
let values = bytes.read_u64::<BigEndian>()? as usize;
74-
return Ok(ProgressInfo::SpillTotalFileNums(values));
75+
let nums = bytes.read_u64::<BigEndian>()? as usize;
76+
let bytes = bytes.read_u64::<BigEndian>()? as usize;
77+
return Ok(ProgressInfo::SpillTotalStats(SpillProgress::new(
78+
nums, bytes,
79+
)));
7580
}
7681

7782
let rows = bytes.read_u64::<BigEndian>()? as usize;

‎src/query/service/src/servers/http/v1/query/execute_state.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ use std::time::SystemTime;
1818

1919
use databend_common_base::base::tokio::sync::RwLock;
2020
use databend_common_base::base::ProgressValues;
21+
use databend_common_base::base::SpillProgress;
2122
use databend_common_base::runtime::CatchUnwindFuture;
2223
use databend_common_exception::ErrorCode;
2324
use databend_common_exception::Result;
@@ -78,6 +79,7 @@ pub struct Progresses {
7879
pub write_progress: ProgressValues,
7980
pub result_progress: ProgressValues,
8081
pub total_scan: ProgressValues,
82+
pub spill_progress: SpillProgress,
8183
}
8284

8385
impl Progresses {
@@ -87,6 +89,7 @@ impl Progresses {
8789
write_progress: ctx.get_write_progress_value(),
8890
result_progress: ctx.get_result_progress_value(),
8991
total_scan: ctx.get_total_scan_value(),
92+
spill_progress: ctx.get_total_spill_progress(),
9093
}
9194
}
9295
}

‎src/query/service/src/sessions/query_ctx.rs‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ use dashmap::mapref::multiple::RefMulti;
3333
use dashmap::DashMap;
3434
use databend_common_base::base::Progress;
3535
use databend_common_base::base::ProgressValues;
36+
use databend_common_base::base::SpillProgress;
3637
use databend_common_base::runtime::profile::Profile;
3738
use databend_common_base::runtime::profile::ProfileStatisticsName;
3839
use databend_common_base::runtime::GlobalIORuntime;
@@ -318,7 +319,7 @@ impl QueryContext {
318319

319320
pub fn update_init_query_id(&self, id: String) {
320321
self.shared.spilled_files.write().clear();
321-
self.shared.cluster_spill_file_nums.write().clear();
322+
self.shared.cluster_spill_progress.write().clear();
322323
*self.shared.init_query_id.write() = id;
323324
}
324325

@@ -363,32 +364,47 @@ impl QueryContext {
363364
&self,
364365
location: crate::spillers::Location,
365366
layout: crate::spillers::Layout,
367+
data_size: usize,
366368
) {
367369
if matches!(location, crate::spillers::Location::Remote(_)) {
368370
let current_id = self.get_cluster().local_id();
369-
let mut w = self.shared.cluster_spill_file_nums.write();
370-
w.entry(current_id).and_modify(|e| *e += 1).or_insert(1);
371+
let mut w = self.shared.cluster_spill_progress.write();
372+
let p = SpillProgress::new(1, data_size);
373+
w.entry(current_id)
374+
.and_modify(|stats| {
375+
stats.incr(&p);
376+
})
377+
.or_insert(p);
371378
}
372379
{
373380
let mut w = self.shared.spilled_files.write();
374381
w.insert(location, layout);
375382
}
376383
}
377384

378-
pub fn set_cluster_spill_file_nums(&self, source_target: &str, num: usize) {
379-
if num != 0 {
385+
pub fn set_cluster_spill_progress(&self, source_target: &str, stats: SpillProgress) {
386+
if stats.file_nums != 0 {
380387
let _ = self
381388
.shared
382-
.cluster_spill_file_nums
389+
.cluster_spill_progress
383390
.write()
384-
.insert(source_target.to_string(), num);
391+
.insert(source_target.to_string(), stats);
385392
}
386393
}
387394

388-
pub fn get_spill_file_nums(&self, node_id: Option<String>) -> usize {
389-
let r = self.shared.cluster_spill_file_nums.read();
395+
pub fn get_spill_file_stats(&self, node_id: Option<String>) -> SpillProgress {
396+
let r = self.shared.cluster_spill_progress.read();
390397
let node_id = node_id.unwrap_or(self.get_cluster().local_id());
391-
r.get(&node_id).cloned().unwrap_or(0)
398+
r.get(&node_id).cloned().unwrap_or(SpillProgress::default())
399+
}
400+
401+
pub fn get_total_spill_progress(&self) -> SpillProgress {
402+
let r = self.shared.cluster_spill_progress.read();
403+
let mut total = SpillProgress::default();
404+
for (_, stats) in r.iter() {
405+
total.incr(stats);
406+
}
407+
total
392408
}
393409

394410
pub fn get_spill_layout(

‎src/query/service/src/sessions/query_ctx_shared.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ use std::time::SystemTime;
2424
use dashmap::DashMap;
2525
use databend_common_base::base::short_sql;
2626
use databend_common_base::base::Progress;
27+
use databend_common_base::base::SpillProgress;
2728
use databend_common_base::runtime::drop_guard;
2829
use databend_common_base::runtime::Runtime;
2930
use databend_common_catalog::catalog::Catalog;
@@ -143,7 +144,7 @@ pub struct QueryContextShared {
143144

144145
pub(in crate::sessions) query_queued_duration: Arc<RwLock<Duration>>,
145146

146-
pub(in crate::sessions) cluster_spill_file_nums: Arc<RwLock<HashMap<String, usize>>>,
147+
pub(in crate::sessions) cluster_spill_progress: Arc<RwLock<HashMap<String, SpillProgress>>>,
147148
pub(in crate::sessions) spilled_files:
148149
Arc<RwLock<HashMap<crate::spillers::Location, crate::spillers::Layout>>>,
149150
}
@@ -203,7 +204,7 @@ impl QueryContextShared {
203204
multi_table_insert_status: Default::default(),
204205
query_queued_duration: Arc::new(RwLock::new(Duration::from_secs(0))),
205206

206-
cluster_spill_file_nums: Default::default(),
207+
cluster_spill_progress: Default::default(),
207208
spilled_files: Default::default(),
208209
}))
209210
}

‎src/query/service/src/spillers/spiller.rs‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -147,17 +147,21 @@ impl Spiller {
147147

148148
/// Spill some [`DataBlock`] to storage. These blocks will be concat into one.
149149
pub async fn spill(&self, data_block: Vec<DataBlock>) -> Result<Location> {
150-
let (location, layout) = self.spill_unmanage(data_block).await?;
150+
let (location, layout, data_size) = self.spill_unmanage(data_block).await?;
151151

152152
// Record columns layout for spilled data.
153-
self.ctx.add_spill_file(location.clone(), layout.clone());
153+
self.ctx
154+
.add_spill_file(location.clone(), layout.clone(), data_size);
154155
self.private_spilled_files
155156
.write()
156157
.insert(location.clone(), layout);
157158
Ok(location)
158159
}
159160

160-
async fn spill_unmanage(&self, data_block: Vec<DataBlock>) -> Result<(Location, Layout)> {
161+
async fn spill_unmanage(
162+
&self,
163+
data_block: Vec<DataBlock>,
164+
) -> Result<(Location, Layout, usize)> {
161165
debug_assert!(!data_block.is_empty());
162166
let instant = Instant::now();
163167

@@ -176,7 +180,7 @@ impl Spiller {
176180
// Record statistics.
177181
record_write_profile(&location, &instant, data_size);
178182
let layout = columns_layout.pop().unwrap();
179-
Ok((location, layout))
183+
Ok((location, layout, data_size))
180184
}
181185

182186
pub fn create_unique_location(&self) -> String {
@@ -205,8 +209,11 @@ impl Spiller {
205209
}
206210

207211
writer.close().await?;
208-
self.ctx
209-
.add_spill_file(Location::Remote(location.clone()), Layout::Aggregate);
212+
self.ctx.add_spill_file(
213+
Location::Remote(location.clone()),
214+
Layout::Aggregate,
215+
write_bytes,
216+
);
210217

211218
self.private_spilled_files
212219
.write()
@@ -289,7 +296,8 @@ impl Spiller {
289296
// Record statistics.
290297
record_write_profile(&location, &instant, write_bytes);
291298

292-
self.ctx.add_spill_file(location.clone(), layout.clone());
299+
self.ctx
300+
.add_spill_file(location.clone(), layout.clone(), write_bytes);
293301
self.private_spilled_files
294302
.write()
295303
.insert(location.clone(), layout);

‎tests/sqllogictests/src/main.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ static HYBRID_CONFIGS: LazyLock<Vec<(Box<ClientType>, usize)>> = LazyLock::new(|
6363
(Box::new(ClientType::MySQL), 3),
6464
(
6565
Box::new(ClientType::Ttc(
66-
"sundyli/ttc-rust:latest".to_string(),
66+
"datafuselabs/ttc-rust:latest".to_string(),
6767
TTC_PORT_START,
6868
)),
6969
7,

0 commit comments

Comments
 (0)