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: 5,
267 retry_backoff_ms: 5_000,
268 lock_timeout_s: 10,
269 memory_threshold_mb: 2_048,
270 max_batch_memory_mb: None,
271 on_batch_memory_exceeded: BatchMemoryPolicy::Warn,
272 adaptive: false,
273 min_parallel: None,
274 max_value_mb: Some(256),
275 configured_profile: TuningProfile::Safe,
276 },
277 }
278 }
279
280 pub fn profile_name(&self) -> &'static str {
281 match self.configured_profile {
282 TuningProfile::Fast => "fast",
283 TuningProfile::Balanced => "balanced",
284 TuningProfile::Safe => "safe",
285 }
286 }
287
288 pub fn effective_batch_size(&self, schema: Option<&SchemaRef>) -> usize {
291 if let (Some(mem_mb), Some(schema)) = (self.batch_size_memory_mb, schema) {
292 let computed = compute_batch_size_from_memory(mem_mb, schema);
293 log::info!(
294 "batch_size_memory_mb={}: estimated row ~{}B, computed batch_size={}",
295 mem_mb,
296 estimate_row_bytes(schema),
297 computed
298 );
299 computed
300 } else {
301 self.batch_size
302 }
303 }
304
305 pub fn batch_memory_bytes(batch: &arrow::record_batch::RecordBatch) -> usize {
311 batch
312 .columns()
313 .iter()
314 .map(|col| col.get_array_memory_size())
315 .sum()
316 }
317
318 pub fn resource_summary(&self) -> ResourceSummary {
325 const NARROW_BYTES: f64 = 200.0;
326 const WIDE_BYTES: f64 = 10_240.0;
327 let batch = self.batch_size as f64;
328 let batch_narrow_mb = batch * NARROW_BYTES / (1024.0 * 1024.0);
329 let batch_wide_mb = batch * WIDE_BYTES / (1024.0 * 1024.0);
330 ResourceSummary {
331 profile: self.profile_name().to_string(),
332 batch_size: self.batch_size,
333 batch_size_memory_mb: self.batch_size_memory_mb,
334 memory_threshold_mb: self.memory_threshold_mb,
335 throttle_ms: self.throttle_ms,
336 batch_narrow_mb,
337 batch_wide_mb,
338 wide_table_risk: batch_wide_mb > 128.0,
339 }
340 }
341}
342
343impl std::fmt::Display for SourceTuning {
344 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345 write!(
346 f,
347 "profile={}, batch_size={}, throttle={}ms, timeout={}s, retries={}, lock_timeout={}s",
348 self.profile_name(),
349 self.batch_size,
350 self.throttle_ms,
351 self.statement_timeout_s,
352 self.max_retries,
353 self.lock_timeout_s,
354 )
355 }
356}
357
358#[derive(Debug, Clone)]
367pub struct ResourceSummary {
368 #[allow(dead_code)]
369 pub profile: String,
370 pub batch_size: usize,
371 pub batch_size_memory_mb: Option<usize>,
372 pub memory_threshold_mb: usize,
373 pub throttle_ms: u64,
374 pub batch_narrow_mb: f64,
375 pub batch_wide_mb: f64,
376 pub wide_table_risk: bool,
377}
378
379#[cfg(test)]
380mod tests {
381 use super::*;
382
383 fn cfg_with_profile(profile: TuningProfile) -> TuningConfig {
384 TuningConfig {
385 profile: Some(profile),
386 ..Default::default()
387 }
388 }
389
390 #[test]
391 fn default_config_uses_balanced_profile() {
392 let t = SourceTuning::from_config(None);
393 assert_eq!(t.batch_size, 10_000);
394 assert_eq!(t.throttle_ms, 50);
395 assert_eq!(t.statement_timeout_s, 300);
396 assert_eq!(t.max_retries, 3);
397 assert_eq!(t.retry_backoff_ms, 2_000);
398 assert_eq!(t.lock_timeout_s, 30);
399 }
400
401 #[test]
402 fn fallback_profile_used_when_no_profile_in_config() {
403 let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Fast);
404 assert_eq!(t.batch_size, 50_000);
405 assert_eq!(t.throttle_ms, 0, "fallback to Fast must zero the throttle");
406 assert_eq!(t.profile_name(), "fast");
407
408 let t = SourceTuning::from_config_with_default_profile(None, TuningProfile::Safe);
409 assert_eq!(t.throttle_ms, 500);
410 assert_eq!(t.profile_name(), "safe");
411 }
412
413 #[test]
414 fn explicit_profile_wins_over_fallback() {
415 let cfg = cfg_with_profile(TuningProfile::Balanced);
416 let t = SourceTuning::from_config_with_default_profile(Some(&cfg), TuningProfile::Fast);
417 assert_eq!(
418 t.throttle_ms, 50,
419 "explicit balanced profile must keep its throttle"
420 );
421 assert_eq!(t.profile_name(), "balanced");
422 }
423
424 #[test]
425 fn fast_profile_favors_throughput() {
426 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
427 assert_eq!(t.batch_size, 50_000);
428 assert_eq!(t.throttle_ms, 0);
429 assert_eq!(t.statement_timeout_s, 0);
430 assert_eq!(t.max_retries, 1);
431 }
432
433 #[test]
434 fn safe_profile_limits_impact() {
435 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
436 assert_eq!(t.batch_size, 2_000);
437 assert_eq!(t.throttle_ms, 500);
438 assert_eq!(t.statement_timeout_s, 120);
439 assert_eq!(t.max_retries, 5);
440 assert_eq!(t.retry_backoff_ms, 5_000);
441 assert_eq!(t.lock_timeout_s, 10);
442 }
443
444 #[test]
445 fn explicit_fields_override_profile_defaults() {
446 let cfg = TuningConfig {
447 profile: Some(TuningProfile::Safe),
448 batch_size: Some(3_000),
449 throttle_ms: Some(250),
450 ..Default::default()
451 };
452 let t = SourceTuning::from_config(Some(&cfg));
453 assert_eq!(t.batch_size, 3_000, "explicit batch_size should win");
454 assert_eq!(t.throttle_ms, 250, "explicit throttle_ms should win");
455 assert_eq!(
456 t.statement_timeout_s, 120,
457 "non-overridden field stays at safe default"
458 );
459 assert_eq!(
460 t.max_retries, 5,
461 "non-overridden field stays at safe default"
462 );
463 }
464
465 #[test]
466 fn profile_name_fast() {
467 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Fast)));
468 assert_eq!(t.profile_name(), "fast");
469 }
470
471 #[test]
472 fn profile_name_balanced() {
473 let t = SourceTuning::from_config(None);
474 assert_eq!(t.profile_name(), "balanced");
475 }
476
477 #[test]
478 fn profile_name_safe() {
479 let t = SourceTuning::from_config(Some(&cfg_with_profile(TuningProfile::Safe)));
480 assert_eq!(t.profile_name(), "safe");
481 }
482
483 #[test]
484 fn display_contains_all_fields() {
485 let t = SourceTuning::from_config(None);
486 let s = t.to_string();
487 assert!(s.contains("profile=balanced"), "missing profile in: {s}");
488 assert!(s.contains("batch_size=10000"), "missing batch_size in: {s}");
489 assert!(s.contains("throttle=50ms"), "missing throttle in: {s}");
490 assert!(s.contains("timeout=300s"), "missing timeout in: {s}");
491 assert!(s.contains("retries=3"), "missing retries in: {s}");
492 assert!(
493 s.contains("lock_timeout=30s"),
494 "missing lock_timeout in: {s}"
495 );
496 }
497
498 #[test]
499 fn merge_tuning_export_overrides_source_fields() {
500 let source = TuningConfig {
501 profile: Some(TuningProfile::Fast),
502 batch_size: Some(1_000),
503 throttle_ms: Some(0),
504 ..Default::default()
505 };
506 let export = TuningConfig {
507 profile: Some(TuningProfile::Safe),
508 batch_size: None,
509 ..Default::default()
510 };
511 let m = merge_tuning_config(Some(&source), Some(&export)).expect("merged");
512 assert_eq!(m.profile, Some(TuningProfile::Safe));
513 assert_eq!(
514 m.batch_size,
515 Some(1_000),
516 "export omitted batch_size -> keep source"
517 );
518 assert_eq!(m.throttle_ms, Some(0));
519 }
520
521 #[test]
522 fn merge_tuning_export_only() {
523 let e = cfg_with_profile(TuningProfile::Fast);
524 let m = merge_tuning_config(None, Some(&e)).expect("merged");
525 assert_eq!(m.profile, Some(TuningProfile::Fast));
526 }
527
528 #[test]
529 fn effective_batch_size_without_memory() {
530 let t = SourceTuning::from_config(None);
531 assert_eq!(t.effective_batch_size(None), 10_000);
532 }
533
534 #[test]
535 fn effective_batch_size_with_memory() {
536 use arrow::datatypes::{DataType, Field, Schema};
537 use std::sync::Arc;
538 let cfg = TuningConfig {
539 batch_size_memory_mb: Some(256),
540 ..Default::default()
541 };
542 let t = SourceTuning::from_config(Some(&cfg));
543 let schema = Arc::new(Schema::new(vec![
544 Field::new("id", DataType::Int64, false),
545 Field::new("name", DataType::Utf8, true),
546 ]));
547 let bs = t.effective_batch_size(Some(&schema));
548 assert!((1_000..=150_000).contains(&bs), "got {bs}");
549 assert_eq!(bs, 150_000);
551 }
552
553 #[test]
554 fn a_tuning_block_omitting_batch_size_memory_mb_keeps_the_profile_default() {
555 let only_throttle = TuningConfig {
560 throttle_ms: Some(50),
561 ..Default::default()
562 };
563 assert_eq!(
564 SourceTuning::from_config(Some(&only_throttle)).batch_size_memory_mb,
565 Some(32),
566 "an unrelated tuning field must not disable Balanced's memory cap"
567 );
568 let fast = TuningConfig {
569 profile: Some(TuningProfile::Fast),
570 ..Default::default()
571 };
572 assert_eq!(
573 SourceTuning::from_config(Some(&fast)).batch_size_memory_mb,
574 Some(64),
575 "naming the Fast profile must keep its Some(64) memory cap"
576 );
577 let rowcap = TuningConfig {
580 batch_size: Some(5_000),
581 ..Default::default()
582 };
583 assert_eq!(
584 SourceTuning::from_config(Some(&rowcap)).batch_size_memory_mb,
585 None,
586 "an explicit batch_size drops the memory cap so the row cap wins"
587 );
588 let memcap = TuningConfig {
590 batch_size_memory_mb: Some(16),
591 ..Default::default()
592 };
593 assert_eq!(
594 SourceTuning::from_config(Some(&memcap)).batch_size_memory_mb,
595 Some(16)
596 );
597 }
598
599 #[test]
600 fn resource_summary_balanced_profile() {
601 let t = SourceTuning::from_config(None);
602 let r = t.resource_summary();
603 assert_eq!(r.profile, "balanced");
604 assert_eq!(r.batch_size, 10_000);
605 assert_eq!(r.batch_size_memory_mb, Some(32));
610 assert_eq!(r.memory_threshold_mb, 4_096);
611 assert_eq!(r.throttle_ms, 50);
612 assert!(
614 r.batch_narrow_mb < 5.0,
615 "narrow too high: {}",
616 r.batch_narrow_mb
617 );
618 assert!(
620 !r.wide_table_risk,
621 "balanced 10k should not trigger wide_table_risk"
622 );
623 }
624
625 #[test]
626 fn resource_summary_fast_profile_triggers_wide_table_risk() {
627 let t = SourceTuning::from_config(Some(&TuningConfig {
628 profile: Some(TuningProfile::Fast),
629 ..Default::default()
630 }));
631 let r = t.resource_summary();
632 assert_eq!(r.batch_size, 50_000);
633 assert!(r.wide_table_risk, "fast 50k should trigger wide_table_risk");
635 }
636
637 #[test]
638 fn resource_summary_with_adaptive_batch() {
639 let cfg = TuningConfig {
640 batch_size_memory_mb: Some(64),
641 ..Default::default()
642 };
643 let t = SourceTuning::from_config(Some(&cfg));
644 let r = t.resource_summary();
645 assert_eq!(r.batch_size_memory_mb, Some(64));
646 }
647}