1use arrow::datatypes::SchemaRef;
10use schemars::JsonSchema;
11use serde::{Deserialize, Serialize};
12
13use super::memory::{compute_batch_size_from_memory, estimate_row_bytes};
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct SourceTuning {
17 pub batch_size: usize,
18 pub batch_size_memory_mb: Option<usize>,
19 pub throttle_ms: u64,
20 pub statement_timeout_s: u64,
21 pub max_retries: u32,
22 pub retry_backoff_ms: u64,
23 pub lock_timeout_s: u64,
24 pub memory_threshold_mb: usize,
26 pub max_batch_memory_mb: Option<usize>,
28 pub on_batch_memory_exceeded: BatchMemoryPolicy,
29 pub adaptive: bool,
33 pub min_parallel: Option<usize>,
38 pub max_value_mb: Option<usize>,
44 configured_profile: TuningProfile,
45}
46
47#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq)]
48#[serde(rename_all = "lowercase")]
49pub enum TuningProfile {
50 Fast,
51 Balanced,
52 Safe,
53}
54
55#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq, Eq, Default)]
57#[serde(rename_all = "snake_case")]
58pub enum BatchMemoryPolicy {
59 #[default]
61 Warn,
62 Fail,
64 AutoShrink,
67}
68
69#[derive(Debug, Deserialize, Serialize, JsonSchema, Default, Clone)]
70#[serde(deny_unknown_fields)]
71pub struct TuningConfig {
72 pub profile: Option<TuningProfile>,
73 pub batch_size: Option<usize>,
74 pub batch_size_memory_mb: Option<usize>,
76 pub throttle_ms: Option<u64>,
77 pub statement_timeout_s: Option<u64>,
78 pub max_retries: Option<u32>,
79 pub retry_backoff_ms: Option<u64>,
80 pub lock_timeout_s: Option<u64>,
81 pub memory_threshold_mb: Option<usize>,
82 pub max_batch_memory_mb: Option<usize>,
85 pub on_batch_memory_exceeded: Option<BatchMemoryPolicy>,
87 pub adaptive: Option<bool>,
91 pub min_parallel: Option<usize>,
94 pub max_value_mb: Option<usize>,
98}
99
100pub fn merge_tuning_config(
103 source: Option<&TuningConfig>,
104 export: Option<&TuningConfig>,
105) -> Option<TuningConfig> {
106 match (source, export) {
107 (None, None) => None,
108 (Some(s), None) => Some(s.clone()),
109 (None, Some(e)) => Some(e.clone()),
110 (Some(s), Some(e)) => Some(TuningConfig {
111 profile: e.profile.or(s.profile),
112 batch_size: e.batch_size.or(s.batch_size),
113 batch_size_memory_mb: e.batch_size_memory_mb.or(s.batch_size_memory_mb),
114 throttle_ms: e.throttle_ms.or(s.throttle_ms),
115 statement_timeout_s: e.statement_timeout_s.or(s.statement_timeout_s),
116 max_retries: e.max_retries.or(s.max_retries),
117 retry_backoff_ms: e.retry_backoff_ms.or(s.retry_backoff_ms),
118 lock_timeout_s: e.lock_timeout_s.or(s.lock_timeout_s),
119 memory_threshold_mb: e.memory_threshold_mb.or(s.memory_threshold_mb),
120 max_batch_memory_mb: e.max_batch_memory_mb.or(s.max_batch_memory_mb),
121 on_batch_memory_exceeded: e.on_batch_memory_exceeded.or(s.on_batch_memory_exceeded),
122 adaptive: e.adaptive.or(s.adaptive),
123 min_parallel: e.min_parallel.or(s.min_parallel),
124 max_value_mb: e.max_value_mb.or(s.max_value_mb),
125 }),
126 }
127}
128
129impl SourceTuning {
130 pub(crate) fn max_value_bytes(&self) -> Option<usize> {
136 self.max_value_mb
137 .filter(|&mb| mb > 0)
138 .map(|mb| mb * 1024 * 1024)
139 }
140
141 #[allow(dead_code)]
146 pub fn from_config(config: Option<&TuningConfig>) -> Self {
147 Self::from_config_with_default_profile(config, TuningProfile::Balanced)
148 }
149
150 pub fn from_config_with_default_profile(
153 config: Option<&TuningConfig>,
154 fallback_profile: TuningProfile,
155 ) -> Self {
156 let profile = config.and_then(|c| c.profile).unwrap_or(fallback_profile);
157
158 let mut tuning = Self::from_profile(profile);
159 tuning.configured_profile = profile;
160
161 if let Some(cfg) = config {
162 if let Some(v) = cfg.batch_size {
163 tuning.batch_size = v;
164 }
165 if cfg.batch_size_memory_mb.is_some() {
176 tuning.batch_size_memory_mb = cfg.batch_size_memory_mb;
177 } else if cfg.batch_size.is_some() {
178 tuning.batch_size_memory_mb = None;
179 }
180 if let Some(v) = cfg.throttle_ms {
181 tuning.throttle_ms = v;
182 }
183 if let Some(v) = cfg.statement_timeout_s {
184 tuning.statement_timeout_s = v;
185 }
186 if let Some(v) = cfg.max_retries {
187 tuning.max_retries = v;
188 }
189 if let Some(v) = cfg.retry_backoff_ms {
190 tuning.retry_backoff_ms = v;
191 }
192 if let Some(v) = cfg.lock_timeout_s {
193 tuning.lock_timeout_s = v;
194 }
195 if let Some(v) = cfg.memory_threshold_mb {
196 tuning.memory_threshold_mb = v;
197 }
198 tuning.max_batch_memory_mb = cfg.max_batch_memory_mb;
199 if let Some(v) = cfg.on_batch_memory_exceeded {
200 tuning.on_batch_memory_exceeded = v;
201 }
202 if let Some(v) = cfg.adaptive {
203 tuning.adaptive = v;
204 }
205 if cfg.min_parallel.is_some() {
206 tuning.min_parallel = cfg.min_parallel;
207 }
208 if cfg.max_value_mb.is_some() {
209 tuning.max_value_mb = cfg.max_value_mb;
210 }
211 }
212
213 tuning
214 }
215
216 fn from_profile(profile: TuningProfile) -> Self {
217 match profile {
218 TuningProfile::Fast => Self {
219 batch_size: 50_000,
222 batch_size_memory_mb: Some(64),
223 throttle_ms: 0,
224 statement_timeout_s: 0,
225 max_retries: 1,
226 retry_backoff_ms: 1_000,
227 lock_timeout_s: 0,
228 memory_threshold_mb: 0,
229 max_batch_memory_mb: None,
230 on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
231 adaptive: false,
232 min_parallel: None,
233 max_value_mb: Some(256),
234 configured_profile: TuningProfile::Fast,
235 },
236 TuningProfile::Balanced => Self {
237 batch_size: 10_000,
247 batch_size_memory_mb: Some(32),
248 throttle_ms: 50,
249 statement_timeout_s: 300,
250 max_retries: 3,
251 retry_backoff_ms: 2_000,
252 lock_timeout_s: 30,
253 memory_threshold_mb: 4_096,
254 max_batch_memory_mb: None,
255 on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
256 adaptive: false,
257 min_parallel: None,
258 max_value_mb: Some(256),
259 configured_profile: TuningProfile::Balanced,
260 },
261 TuningProfile::Safe => Self {
262 batch_size: 2_000,
263 batch_size_memory_mb: None,
264 throttle_ms: 500,
265 statement_timeout_s: 120,
266 max_retries: 10,
273 retry_backoff_ms: 5_000,
274 lock_timeout_s: 10,
275 memory_threshold_mb: 2_048,
276 max_batch_memory_mb: None,
277 on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
278 adaptive: false,
279 min_parallel: None,
280 max_value_mb: Some(256),
281 configured_profile: TuningProfile::Safe,
282 },
283 }
284 }
285
286 pub fn profile_name(&self) -> &'static str {
287 match self.configured_profile {
288 TuningProfile::Fast => "fast",
289 TuningProfile::Balanced => "balanced",
290 TuningProfile::Safe => "safe",
291 }
292 }
293
294 pub fn effective_batch_size(&self, schema: Option<&SchemaRef>) -> usize {
297 if let (Some(mem_mb), Some(schema)) = (self.batch_size_memory_mb, schema) {
298 let computed = compute_batch_size_from_memory(mem_mb, schema);
299 log::info!(
300 "batch_size_memory_mb={}: estimated row ~{}B, computed batch_size={}",
301 mem_mb,
302 estimate_row_bytes(schema),
303 computed
304 );
305 computed
306 } else {
307 self.batch_size
308 }
309 }
310
311 pub fn batch_memory_bytes(batch: &arrow::record_batch::RecordBatch) -> usize {
317 batch
318 .columns()
319 .iter()
320 .map(|col| col.get_array_memory_size())
321 .sum()
322 }
323
324 pub fn resource_summary(&self) -> ResourceSummary {
331 const NARROW_BYTES: f64 = 200.0;
332 const WIDE_BYTES: f64 = 10_240.0;
333 let batch = self.batch_size as f64;
334 let batch_narrow_mb = batch * NARROW_BYTES / (1024.0 * 1024.0);
335 let batch_wide_mb = batch * WIDE_BYTES / (1024.0 * 1024.0);
336 ResourceSummary {
337 profile: self.profile_name().to_string(),
338 batch_size: self.batch_size,
339 batch_size_memory_mb: self.batch_size_memory_mb,
340 memory_threshold_mb: self.memory_threshold_mb,
341 throttle_ms: self.throttle_ms,
342 batch_narrow_mb,
343 batch_wide_mb,
344 wide_table_risk: batch_wide_mb > 128.0,
345 }
346 }
347}
348
349impl std::fmt::Display for SourceTuning {
350 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
351 write!(
352 f,
353 "profile={}, batch_size={}, throttle={}ms, timeout={}s, retries={}, lock_timeout={}s",
354 self.profile_name(),
355 self.batch_size,
356 self.throttle_ms,
357 self.statement_timeout_s,
358 self.max_retries,
359 self.lock_timeout_s,
360 )
361 }
362}
363
364#[derive(Debug, Clone)]
373pub struct ResourceSummary {
374 #[allow(dead_code)]
375 pub profile: String,
376 pub batch_size: usize,
377 pub batch_size_memory_mb: Option<usize>,
378 pub memory_threshold_mb: usize,
379 pub throttle_ms: u64,
380 pub batch_narrow_mb: f64,
381 pub batch_wide_mb: f64,
382 pub wide_table_risk: bool,
383}
384
385#[cfg(test)]
386mod tests {
387 use super::*;
388
389 fn cfg_with_profile(profile: TuningProfile) -> TuningConfig {
390 TuningConfig {
391 profile: Some(profile),
392 ..Default::default()
393 }
394 }
395
396 #[test]
397 fn default_config_uses_balanced_profile() {
398 let t = SourceTuning::from_config(None);
399 assert_eq!(t.batch_size, 10_000);
400 assert_eq!(t.throttle_ms, 50);
401 assert_eq!(t.statement_timeout_s, 300);
402 assert_eq!(t.max_retries, 3);
403 assert_eq!(t.retry_backoff_ms, 2_000);
404 assert_eq!(t.lock_timeout_s, 30);
405 }
406
407 #[test]
408 fn fallback_profile_used_when_no_profile_in_config() {
409 let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Fast);
410 assert_eq!(t.batch_size, 50_000);
411 assert_eq!(t.throttle_ms, 0, "fallback to Fast must zero the throttle");
412 assert_eq!(t.profile_name(), "fast");
413
414 let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Safe);
415 assert_eq!(t.throttle_ms, 500);
416 assert_eq!(t.profile_name(), "safe");
417 }
418
419 #[test]
420 fn explicit_profile_wins_over_fallback() {
421 let cfg = cfg_with_profile(TuningProfile::Balanced);
422 let t = SourceTuning::from_config_with_default_profile(Some(&cfg), TuningProfile::Fast);
423 assert_eq!(
424 t.throttle_ms, 50,
425 "explicit balanced profile must keep its throttle"
426 );
427 assert_eq!(t.profile_name(), "balanced");
428 }
429
430 #[test]
431 fn fast_profile_favors_throughput() {
432 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
433 assert_eq!(t.batch_size, 50_000);
434 assert_eq!(t.throttle_ms, 0);
435 assert_eq!(t.statement_timeout_s, 0);
436 assert_eq!(t.max_retries, 1);
437 }
438
439 #[test]
440 fn safe_profile_limits_impact() {
441 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
442 assert_eq!(t.batch_size, 2_000);
443 assert_eq!(t.throttle_ms, 500);
444 assert_eq!(t.statement_timeout_s, 120);
445 assert_eq!(t.max_retries, 10);
446 assert_eq!(t.retry_backoff_ms, 5_000);
447 assert_eq!(t.lock_timeout_s, 10);
448 }
449
450 #[test]
451 fn explicit_fields_override_profile_defaults() {
452 let cfg = TuningConfig {
453 profile: Some(TuningProfile::Safe),
454 batch_size: Some(3_000),
455 throttle_ms: Some(250),
456 ..Default::default()
457 };
458 let t = SourceTuning::from_config(Some(&cfg));
459 assert_eq!(t.batch_size, 3_000, "explicit batch_size should win");
460 assert_eq!(t.throttle_ms, 250, "explicit throttle_ms should win");
461 assert_eq!(
462 t.statement_timeout_s, 120,
463 "non-overridden field stays at safe default"
464 );
465 assert_eq!(
466 t.max_retries, 10,
467 "non-overridden field stays at safe default"
468 );
469 }
470
471 #[test]
472 fn profile_name_fast() {
473 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
474 assert_eq!(t.profile_name(), "fast");
475 }
476
477 #[test]
478 fn profile_name_balanced() {
479 let t = SourceTuning::from_config(None);
480 assert_eq!(t.profile_name(), "balanced");
481 }
482
483 #[test]
484 fn profile_name_safe() {
485 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
486 assert_eq!(t.profile_name(), "safe");
487 }
488
489 #[test]
490 fn display_contains_all_fields() {
491 let t = SourceTuning::from_config(None);
492 let s = t.to_string();
493 assert!(s.contains("profile=balanced"), "missing profile in: {s}");
494 assert!(s.contains("batch_size=10000"), "missing batch_size in: {s}");
495 assert!(s.contains("throttle=50ms"), "missing throttle in: {s}");
496 assert!(s.contains("timeout=300s"), "missing timeout in: {s}");
497 assert!(s.contains("retries=3"), "missing retries in: {s}");
498 assert!(
499 s.contains("lock_timeout=30s"),
500 "missing lock_timeout in: {s}"
501 );
502 }
503
504 #[test]
505 fn merge_tuning_export_overrides_source_fields() {
506 let source = TuningConfig {
507 profile: Some(TuningProfile::Fast),
508 batch_size: Some(1_000),
509 throttle_ms: Some(0),
510 ..Default::default()
511 };
512 let export = TuningConfig {
513 profile: Some(TuningProfile::Safe),
514 batch_size: None,
515 ..Default::default()
516 };
517 let m = merge_tuning_config(Some(&source), Some(&export)).expect("merged");
518 assert_eq!(m.profile, Some(TuningProfile::Safe));
519 assert_eq!(
520 m.batch_size,
521 Some(1_000),
522 "export omitted batch_size -> keep source"
523 );
524 assert_eq!(m.throttle_ms, Some(0));
525 }
526
527 #[test]
528 fn merge_tuning_export_only() {
529 let e = cfg_with_profile(TuningProfile::Fast);
530 let m = merge_tuning_config(None, Some(&e)).expect("merged");
531 assert_eq!(m.profile, Some(TuningProfile::Fast));
532 }
533
534 #[test]
535 fn effective_batch_size_without_memory() {
536 let t = SourceTuning::from_config(None);
537 assert_eq!(t.effective_batch_size(None), 10_000);
538 }
539
540 #[test]
541 fn effective_batch_size_with_memory() {
542 use arrow::datatypes::{DataType, Field, Schema};
543 use std::sync::Arc;
544 let cfg = TuningConfig {
545 batch_size_memory_mb: Some(256),
546 ..Default::default()
547 };
548 let t = SourceTuning::from_config(Some(&cfg));
549 let schema = Arc::new(Schema::new(vec![
550 Field::new("id", DataType::Int64, false),
551 Field::new("name", DataType::Utf8, true),
552 ]));
553 let bs = t.effective_batch_size(Some(&schema));
554 assert!((1_000..=150_000).contains(&bs), "got {bs}");
555 assert_eq!(bs, 150_000);
557 }
558
559 #[test]
560 fn a_tuning_block_omitting_batch_size_memory_mb_keeps_the_profile_default() {
561 let only_throttle = TuningConfig {
566 throttle_ms: Some(50),
567 ..Default::default()
568 };
569 assert_eq!(
570 SourceTuning::from_config(Some(&only_throttle)).batch_size_memory_mb,
571 Some(32),
572 "an unrelated tuning field must not disable Balanced's memory cap"
573 );
574 let fast = TuningConfig {
575 profile: Some(TuningProfile::Fast),
576 ..Default::default()
577 };
578 assert_eq!(
579 SourceTuning::from_config(Some(&fast)).batch_size_memory_mb,
580 Some(64),
581 "naming the Fast profile must keep its Some(64) memory cap"
582 );
583 let rowcap = TuningConfig {
586 batch_size: Some(5_000),
587 ..Default::default()
588 };
589 assert_eq!(
590 SourceTuning::from_config(Some(&rowcap)).batch_size_memory_mb,
591 None,
592 "an explicit batch_size drops the memory cap so the row cap wins"
593 );
594 let memcap = TuningConfig {
596 batch_size_memory_mb: Some(16),
597 ..Default::default()
598 };
599 assert_eq!(
600 SourceTuning::from_config(Some(&memcap)).batch_size_memory_mb,
601 Some(16)
602 );
603 }
604
605 #[test]
606 fn resource_summary_balanced_profile() {
607 let t = SourceTuning::from_config(None);
608 let r = t.resource_summary();
609 assert_eq!(r.profile, "balanced");
610 assert_eq!(r.batch_size, 10_000);
611 assert_eq!(r.batch_size_memory_mb, Some(32));
616 assert_eq!(r.memory_threshold_mb, 4_096);
617 assert_eq!(r.throttle_ms, 50);
618 assert!(
620 r.batch_narrow_mb < 5.0,
621 "narrow too high: {}",
622 r.batch_narrow_mb
623 );
624 assert!(
626 !r.wide_table_risk,
627 "balanced 10k should not trigger wide_table_risk"
628 );
629 }
630
631 #[test]
632 fn resource_summary_fast_profile_triggers_wide_table_risk() {
633 let t = SourceTuning::from_config(Some(&TuningConfig {
634 profile: Some(TuningProfile::Fast),
635 ..Default::default()
636 }));
637 let r = t.resource_summary();
638 assert_eq!(r.batch_size, 50_000);
639 assert!(r.wide_table_risk, "fast 50k should trigger wide_table_risk");
641 }
642
643 #[test]
644 fn resource_summary_with_adaptive_batch() {
645 let cfg = TuningConfig {
646 batch_size_memory_mb: Some(64),
647 ..Default::default()
648 };
649 let t = SourceTuning::from_config(Some(&cfg));
650 let r = t.resource_summary();
651 assert_eq!(r.batch_size_memory_mb, Some(64));
652 }
653}