Skip to main content

radixdb_executor/optimizer/
workload.rs

1// Copyright 2026 RadixDB Contributors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Workload Learning and Edge-Aware Query Optimization
16//!
17//! This module implements RadixDB's unique optimization features:
18//!
19//! 1. **Workload Learning**: Learns from historical query patterns to predict future behavior
20//!    - Query pattern fingerprinting and frequency tracking
21//!    - Automatic index recommendation based on access patterns
22//!    - Hot column detection for pre-materialization hints
23//!    - Temporal pattern detection (batch vs interactive workloads)
24//!
25//! 2. **Edge-Aware Planning**: Special optimizations for edge computing environments
26//!    - Memory-constrained execution strategies
27//!    - Network partition tolerance (graceful degradation)
28//!    - Battery-aware query scheduling (for IoT/mobile)
29//!    - Incremental result computation for slow connections
30//!
31//! These features make RadixDB unique in that it learns from your specific workload
32//! patterns rather than relying solely on static cost models.
33
34#![allow(clippy::too_many_arguments)]
35
36use radixdb_core::time_compat::Instant;
37use rustc_hash::FxHashMap;
38use std::sync::atomic::{AtomicU64, Ordering};
39use std::sync::RwLock;
40use std::time::Duration;
41
42/// Maximum number of query fingerprints to store
43/// This prevents unbounded memory growth in workload learner
44const MAX_FINGERPRINTS: usize = 50000;
45
46/// Global workload learner instance
47static WORKLOAD_LEARNER: std::sync::OnceLock<WorkloadLearner> = std::sync::OnceLock::new();
48
49/// Get the global workload learner instance
50pub fn global_workload_learner() -> &'static WorkloadLearner {
51    WORKLOAD_LEARNER.get_or_init(WorkloadLearner::new)
52}
53
54/// Configuration for workload-aware optimization
55#[derive(Debug, Clone)]
56pub struct WorkloadConfig {
57    /// Enable workload learning
58    pub learning_enabled: bool,
59    /// Edge computing mode (enables memory-optimized execution)
60    pub edge_mode: EdgeMode,
61    /// Maximum memory for query execution (0 = unlimited)
62    pub memory_limit_mb: u64,
63    /// Enable incremental result streaming
64    pub incremental_results: bool,
65}
66
67impl Default for WorkloadConfig {
68    fn default() -> Self {
69        Self {
70            learning_enabled: true,
71            edge_mode: EdgeMode::Standard,
72            memory_limit_mb: 0,
73            incremental_results: false,
74        }
75    }
76}
77
78/// Edge computing mode
79#[derive(Debug, Clone, Copy, PartialEq)]
80pub enum EdgeMode {
81    /// Standard mode - no special constraints
82    Standard,
83    /// Constrained mode - limited memory, prefer streaming operators
84    Constrained,
85    /// Ultra-low mode - extreme memory limits, single-pass only
86    UltraLow,
87    /// Mobile mode - battery aware, network resilient
88    Mobile,
89}
90
91impl EdgeMode {
92    /// Get the memory multiplier for cost estimation
93    /// Lower means we penalize memory-heavy operations more
94    pub fn memory_cost_multiplier(&self) -> f64 {
95        match self {
96            EdgeMode::Standard => 1.0,
97            EdgeMode::Constrained => 5.0,
98            EdgeMode::UltraLow => 20.0,
99            EdgeMode::Mobile => 3.0,
100        }
101    }
102
103    /// Get the preferred batch size for this mode
104    pub fn preferred_batch_size(&self) -> usize {
105        match self {
106            EdgeMode::Standard => 10000,
107            EdgeMode::Constrained => 1000,
108            EdgeMode::UltraLow => 100,
109            EdgeMode::Mobile => 500,
110        }
111    }
112}
113
114/// Query pattern classification
115#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
116pub enum QueryPattern {
117    /// Point lookup by primary key
118    PointLookup,
119    /// Range scan
120    RangeScan,
121    /// Full table scan
122    FullScan,
123    /// Aggregation query
124    Aggregation,
125    /// Join-heavy query
126    JoinHeavy,
127    /// Complex analytical query
128    Analytical,
129    /// Insert-heavy workload
130    InsertHeavy,
131    /// Update-heavy workload
132    UpdateHeavy,
133    /// Mixed OLTP workload
134    MixedOLTP,
135    /// Unknown/other
136    Unknown,
137}
138
139/// Learned statistics for a query pattern
140#[derive(Debug, Clone)]
141pub struct PatternStats {
142    /// Number of times this pattern was observed
143    pub frequency: u64,
144    /// Average execution time in microseconds
145    pub avg_execution_time_us: f64,
146    /// Peak memory usage in bytes
147    pub peak_memory_bytes: u64,
148    /// Average rows scanned
149    pub avg_rows_scanned: u64,
150    /// Average rows returned
151    pub avg_rows_returned: u64,
152    /// Tables most commonly accessed with this pattern
153    pub hot_tables: Vec<String>,
154    /// Columns most commonly filtered on
155    pub hot_filter_columns: Vec<String>,
156    /// Columns most commonly in ORDER BY
157    pub hot_sort_columns: Vec<String>,
158    /// Last observed timestamp
159    pub last_seen: Instant,
160}
161
162impl PatternStats {
163    fn new() -> Self {
164        Self {
165            frequency: 0,
166            avg_execution_time_us: 0.0,
167            peak_memory_bytes: 0,
168            avg_rows_scanned: 0,
169            avg_rows_returned: 0,
170            hot_tables: Vec::new(),
171            hot_filter_columns: Vec::new(),
172            hot_sort_columns: Vec::new(),
173            last_seen: Instant::now(),
174        }
175    }
176
177    /// Update stats with new observation
178    fn observe(
179        &mut self,
180        execution_time_us: u64,
181        memory_bytes: u64,
182        rows_scanned: u64,
183        rows_returned: u64,
184        tables: Vec<String>,
185        filter_columns: Vec<String>,
186        sort_columns: Vec<String>,
187    ) {
188        self.frequency += 1;
189
190        // Exponential moving average for execution time
191        let alpha = 0.3;
192        self.avg_execution_time_us =
193            alpha * execution_time_us as f64 + (1.0 - alpha) * self.avg_execution_time_us;
194
195        // Keep max memory
196        self.peak_memory_bytes = self.peak_memory_bytes.max(memory_bytes);
197
198        // Exponential moving average for row counts
199        self.avg_rows_scanned =
200            ((alpha * rows_scanned as f64 + (1.0 - alpha) * self.avg_rows_scanned as f64) as u64)
201                .max(1);
202        self.avg_rows_returned =
203            ((alpha * rows_returned as f64 + (1.0 - alpha) * self.avg_rows_returned as f64) as u64)
204                .max(1);
205
206        // Update hot lists (keep top 10)
207        for table in tables {
208            Self::update_hot_list(&mut self.hot_tables, table);
209        }
210        for col in filter_columns {
211            Self::update_hot_list(&mut self.hot_filter_columns, col);
212        }
213        for col in sort_columns {
214            Self::update_hot_list(&mut self.hot_sort_columns, col);
215        }
216
217        self.last_seen = Instant::now();
218    }
219
220    fn update_hot_list(list: &mut Vec<String>, item: String) {
221        if !list.contains(&item) && list.len() < 10 {
222            list.push(item);
223        }
224    }
225}
226
227/// Index recommendation from workload analysis
228#[derive(Debug, Clone)]
229pub struct IndexRecommendation {
230    /// Table name
231    pub table: String,
232    /// Column(s) to index
233    pub columns: Vec<String>,
234    /// Expected benefit score (higher = more beneficial)
235    pub benefit_score: f64,
236    /// Reason for recommendation
237    pub reason: String,
238}
239
240/// Temporal workload pattern
241#[derive(Debug, Clone, Copy, PartialEq)]
242pub enum TemporalPattern {
243    /// Mostly interactive queries (low latency)
244    Interactive,
245    /// Mostly batch processing (throughput focused)
246    Batch,
247    /// Mixed workload
248    Mixed,
249    /// No clear pattern
250    Unknown,
251}
252
253/// Workload learner - learns from query patterns to optimize future queries
254pub struct WorkloadLearner {
255    /// Pattern statistics
256    patterns: RwLock<FxHashMap<QueryPattern, PatternStats>>,
257    /// Query fingerprint to pattern mapping
258    fingerprints: RwLock<FxHashMap<u64, QueryPattern>>,
259    /// Table access frequency
260    table_access_counts: RwLock<FxHashMap<String, AtomicU64>>,
261    /// Column filter frequency (table.column -> count)
262    filter_column_counts: RwLock<FxHashMap<String, AtomicU64>>,
263    /// Total queries observed
264    total_queries: AtomicU64,
265    /// Short queries (< 10ms)
266    short_queries: AtomicU64,
267    /// Long queries (> 1s)
268    long_queries: AtomicU64,
269    /// Current workload config
270    config: RwLock<WorkloadConfig>,
271}
272
273impl WorkloadLearner {
274    /// Create a new workload learner
275    pub fn new() -> Self {
276        Self {
277            patterns: RwLock::new(FxHashMap::default()),
278            fingerprints: RwLock::new(FxHashMap::default()),
279            table_access_counts: RwLock::new(FxHashMap::default()),
280            filter_column_counts: RwLock::new(FxHashMap::default()),
281            total_queries: AtomicU64::new(0),
282            short_queries: AtomicU64::new(0),
283            long_queries: AtomicU64::new(0),
284            config: RwLock::new(WorkloadConfig::default()),
285        }
286    }
287
288    /// Update configuration
289    pub fn set_config(&self, config: WorkloadConfig) {
290        if let Ok(mut cfg) = self.config.write() {
291            *cfg = config;
292        }
293    }
294
295    /// Get current configuration
296    pub fn config(&self) -> WorkloadConfig {
297        self.config.read().map(|c| c.clone()).unwrap_or_default()
298    }
299
300    /// Classify a query into a pattern based on its characteristics
301    pub fn classify_query(
302        &self,
303        has_pk_lookup: bool,
304        has_range_predicate: bool,
305        has_full_scan: bool,
306        has_aggregation: bool,
307        join_count: usize,
308        is_insert: bool,
309        is_update: bool,
310    ) -> QueryPattern {
311        if is_insert {
312            return QueryPattern::InsertHeavy;
313        }
314        if is_update {
315            return QueryPattern::UpdateHeavy;
316        }
317
318        if has_pk_lookup && join_count == 0 && !has_aggregation {
319            return QueryPattern::PointLookup;
320        }
321
322        if join_count >= 3 || (join_count >= 2 && has_aggregation) {
323            return QueryPattern::Analytical;
324        }
325
326        if join_count >= 2 {
327            return QueryPattern::JoinHeavy;
328        }
329
330        if has_aggregation {
331            return QueryPattern::Aggregation;
332        }
333
334        if has_range_predicate && !has_full_scan {
335            return QueryPattern::RangeScan;
336        }
337
338        if has_full_scan {
339            return QueryPattern::FullScan;
340        }
341
342        QueryPattern::Unknown
343    }
344
345    /// Record a query execution for learning
346    pub fn record_query(
347        &self,
348        query_fingerprint: u64,
349        pattern: QueryPattern,
350        execution_time: Duration,
351        memory_bytes: u64,
352        rows_scanned: u64,
353        rows_returned: u64,
354        tables: Vec<String>,
355        filter_columns: Vec<String>,
356        sort_columns: Vec<String>,
357    ) {
358        if !self.is_learning_enabled() {
359            return;
360        }
361
362        let execution_time_us = execution_time.as_micros() as u64;
363
364        // Update query counters
365        self.total_queries.fetch_add(1, Ordering::Relaxed);
366        if execution_time < Duration::from_millis(10) {
367            self.short_queries.fetch_add(1, Ordering::Relaxed);
368        } else if execution_time > Duration::from_secs(1) {
369            self.long_queries.fetch_add(1, Ordering::Relaxed);
370        }
371
372        // Update fingerprint mapping with size cap
373        if let Ok(mut fingerprints) = self.fingerprints.write() {
374            // Evict oldest entries if at capacity (simple approach: clear half when full)
375            if fingerprints.len() >= MAX_FINGERPRINTS
376                && !fingerprints.contains_key(&query_fingerprint)
377            {
378                // Clear half the entries to make room and amortize the eviction cost
379                let target_size = MAX_FINGERPRINTS / 2;
380                let keys_to_remove: Vec<u64> = fingerprints
381                    .keys()
382                    .take(fingerprints.len() - target_size)
383                    .copied()
384                    .collect();
385                for key in keys_to_remove {
386                    fingerprints.remove(&key);
387                }
388            }
389            fingerprints.insert(query_fingerprint, pattern);
390        }
391
392        // Update pattern stats
393        if let Ok(mut patterns) = self.patterns.write() {
394            let stats = patterns.entry(pattern).or_insert_with(PatternStats::new);
395            stats.observe(
396                execution_time_us,
397                memory_bytes,
398                rows_scanned,
399                rows_returned,
400                tables.clone(),
401                filter_columns.clone(),
402                sort_columns,
403            );
404        }
405
406        // Update table access counts
407        if let Ok(table_counts) = self.table_access_counts.read() {
408            for table in &tables {
409                if let Some(count) = table_counts.get(table) {
410                    count.fetch_add(1, Ordering::Relaxed);
411                }
412            }
413        }
414        // Add new tables
415        if let Ok(mut table_counts) = self.table_access_counts.write() {
416            for table in tables {
417                table_counts
418                    .entry(table)
419                    .or_insert_with(|| AtomicU64::new(1));
420            }
421        }
422
423        // Update filter column counts
424        if let Ok(mut filter_counts) = self.filter_column_counts.write() {
425            for col in filter_columns {
426                filter_counts
427                    .entry(col)
428                    .or_insert_with(|| AtomicU64::new(0))
429                    .fetch_add(1, Ordering::Relaxed);
430            }
431        }
432    }
433
434    /// Get pattern for a known query fingerprint
435    pub fn get_pattern(&self, fingerprint: u64) -> Option<QueryPattern> {
436        self.fingerprints
437            .read()
438            .ok()
439            .and_then(|f| f.get(&fingerprint).copied())
440    }
441
442    /// Get statistics for a pattern
443    pub fn get_pattern_stats(&self, pattern: QueryPattern) -> Option<PatternStats> {
444        self.patterns
445            .read()
446            .ok()
447            .and_then(|p| p.get(&pattern).cloned())
448    }
449
450    /// Generate index recommendations based on learned workload
451    pub fn recommend_indexes(&self) -> Vec<IndexRecommendation> {
452        let mut recommendations = Vec::new();
453
454        let filter_counts = match self.filter_column_counts.read() {
455            Ok(c) => c,
456            Err(_) => return recommendations,
457        };
458
459        let total = self.total_queries.load(Ordering::Relaxed) as f64;
460        if total < 100.0 {
461            // Need at least 100 queries before making recommendations
462            return recommendations;
463        }
464
465        // Find frequently filtered columns
466        for (column, count) in filter_counts.iter() {
467            let freq = count.load(Ordering::Relaxed) as f64 / total;
468            if freq > 0.1 {
469                // Filtered in more than 10% of queries
470                let parts: Vec<&str> = column.split('.').collect();
471                if parts.len() == 2 {
472                    let benefit = freq * 100.0; // Simple benefit score
473                    recommendations.push(IndexRecommendation {
474                        table: parts[0].to_string(),
475                        columns: vec![parts[1].to_string()],
476                        benefit_score: benefit,
477                        reason: format!("Column filtered in {:.1}% of queries", freq * 100.0),
478                    });
479                }
480            }
481        }
482
483        // Sort by benefit score
484        recommendations.sort_by(|a, b| {
485            b.benefit_score
486                .partial_cmp(&a.benefit_score)
487                .unwrap_or(std::cmp::Ordering::Equal)
488        });
489
490        // Return top 5
491        recommendations.truncate(5);
492        recommendations
493    }
494
495    /// Detect the temporal workload pattern
496    pub fn detect_temporal_pattern(&self) -> TemporalPattern {
497        let total = self.total_queries.load(Ordering::Relaxed);
498        if total < 100 {
499            return TemporalPattern::Unknown;
500        }
501
502        let short = self.short_queries.load(Ordering::Relaxed);
503        let long = self.long_queries.load(Ordering::Relaxed);
504
505        let short_ratio = short as f64 / total as f64;
506        let long_ratio = long as f64 / total as f64;
507
508        if short_ratio > 0.8 {
509            TemporalPattern::Interactive
510        } else if long_ratio > 0.3 {
511            TemporalPattern::Batch
512        } else if short_ratio > 0.5 && long_ratio > 0.1 {
513            TemporalPattern::Mixed
514        } else {
515            TemporalPattern::Unknown
516        }
517    }
518
519    /// Get hot tables (most frequently accessed)
520    pub fn hot_tables(&self, limit: usize) -> Vec<(String, u64)> {
521        let table_counts = match self.table_access_counts.read() {
522            Ok(c) => c,
523            Err(_) => return Vec::new(),
524        };
525
526        let mut tables: Vec<_> = table_counts
527            .iter()
528            .map(|(k, v)| (k.clone(), v.load(Ordering::Relaxed)))
529            .collect();
530
531        tables.sort_unstable_by_key(|entry| std::cmp::Reverse(entry.1));
532        tables.truncate(limit);
533        tables
534    }
535
536    /// Get optimization hints based on learned workload
537    pub fn get_optimization_hints(&self) -> WorkloadHints {
538        let pattern = self.detect_temporal_pattern();
539        let config = self.config();
540
541        WorkloadHints {
542            prefer_nested_loop: pattern == TemporalPattern::Interactive,
543            prefer_hash_join: pattern == TemporalPattern::Batch,
544            enable_bloom_filters: self.total_queries.load(Ordering::Relaxed) > 1000,
545            target_batch_size: config.edge_mode.preferred_batch_size(),
546            memory_constrained: config.edge_mode != EdgeMode::Standard,
547            incremental_results: config.incremental_results,
548        }
549    }
550
551    fn is_learning_enabled(&self) -> bool {
552        self.config
553            .read()
554            .map(|config| config.learning_enabled)
555            .unwrap_or(false)
556    }
557
558    /// Enable or disable learning
559    pub fn set_learning_enabled(&self, enabled: bool) {
560        if let Ok(mut config) = self.config.write() {
561            config.learning_enabled = enabled;
562        }
563    }
564
565    /// Get total queries observed
566    pub fn total_queries(&self) -> u64 {
567        self.total_queries.load(Ordering::Relaxed)
568    }
569
570    /// Clear all learned data
571    pub fn clear(&self) {
572        if let Ok(mut p) = self.patterns.write() {
573            p.clear();
574        }
575        if let Ok(mut f) = self.fingerprints.write() {
576            f.clear();
577        }
578        if let Ok(mut t) = self.table_access_counts.write() {
579            t.clear();
580        }
581        if let Ok(mut f) = self.filter_column_counts.write() {
582            f.clear();
583        }
584        self.total_queries.store(0, Ordering::Relaxed);
585        self.short_queries.store(0, Ordering::Relaxed);
586        self.long_queries.store(0, Ordering::Relaxed);
587    }
588}
589
590impl Default for WorkloadLearner {
591    fn default() -> Self {
592        Self::new()
593    }
594}
595
596/// Optimization hints derived from workload learning
597#[derive(Debug, Clone)]
598pub struct WorkloadHints {
599    /// Prefer nested loop joins for low-latency
600    pub prefer_nested_loop: bool,
601    /// Prefer hash joins for throughput
602    pub prefer_hash_join: bool,
603    /// Enable bloom filter optimizations
604    pub enable_bloom_filters: bool,
605    /// Target batch size for operators
606    pub target_batch_size: usize,
607    /// Memory is constrained
608    pub memory_constrained: bool,
609    /// Enable incremental result streaming
610    pub incremental_results: bool,
611}
612
613impl Default for WorkloadHints {
614    fn default() -> Self {
615        Self {
616            prefer_nested_loop: false,
617            prefer_hash_join: false,
618            enable_bloom_filters: false,
619            target_batch_size: 10000,
620            memory_constrained: false,
621            incremental_results: false,
622        }
623    }
624}
625
626/// Edge-aware query planner enhancements
627pub struct EdgeAwarePlanner {
628    /// Workload hints
629    hints: WorkloadHints,
630    /// Memory limit in bytes (0 = unlimited)
631    memory_limit: u64,
632}
633
634impl EdgeAwarePlanner {
635    /// Create a new edge-aware planner
636    pub fn new(hints: WorkloadHints, memory_limit: u64) -> Self {
637        Self {
638            hints,
639            memory_limit,
640        }
641    }
642
643    /// Create from global workload learner
644    pub fn from_global() -> Self {
645        let learner = global_workload_learner();
646        let hints = learner.get_optimization_hints();
647        let config = learner.config();
648        Self {
649            hints,
650            memory_limit: config.memory_limit_mb.saturating_mul(1024 * 1024),
651        }
652    }
653
654    /// Adjust cost for edge computing constraints
655    pub fn adjust_cost(&self, base_cost: f64, memory_estimate: u64) -> f64 {
656        let mut cost = base_cost;
657
658        // Penalize high memory usage in constrained mode
659        if self.hints.memory_constrained && self.memory_limit > 0 {
660            if memory_estimate > self.memory_limit {
661                // Heavy penalty for exceeding limit
662                cost *= 100.0;
663            } else if memory_estimate > self.memory_limit / 2 {
664                // Moderate penalty for high usage
665                cost *= 2.0;
666            }
667        }
668
669        cost
670    }
671
672    /// Decide if we should use streaming execution
673    pub fn should_stream(&self, estimated_rows: u64) -> bool {
674        if self.hints.incremental_results {
675            return true;
676        }
677
678        if self.hints.memory_constrained {
679            if self.memory_limit == 0 {
680                return false;
681            }
682            // Stream if estimated rows might cause memory issues
683            let row_size_estimate = 100; // bytes per row estimate
684            estimated_rows.saturating_mul(row_size_estimate) > self.memory_limit / 2
685        } else {
686            false
687        }
688    }
689
690    /// Get preferred batch size
691    pub fn batch_size(&self) -> usize {
692        self.hints.target_batch_size
693    }
694
695    /// Check if bloom filters should be used
696    pub fn use_bloom_filters(&self) -> bool {
697        self.hints.enable_bloom_filters
698    }
699
700    /// Recommend join algorithm for edge constraints
701    pub fn recommend_join_for_edge(
702        &self,
703        build_rows: u64,
704        probe_rows: u64,
705        memory_per_build_row: u64,
706    ) -> EdgeJoinRecommendation {
707        let build_memory = build_rows.saturating_mul(memory_per_build_row);
708
709        if self.memory_limit > 0 && build_memory > self.memory_limit {
710            // Cannot use hash join - would exceed memory
711            EdgeJoinRecommendation::ForceNestedLoop {
712                reason: "Hash join would exceed memory limit".to_string(),
713            }
714        } else if self.hints.prefer_nested_loop && probe_rows < 1000 {
715            // Interactive mode prefers nested loop for small probes
716            EdgeJoinRecommendation::PreferNestedLoop {
717                reason: "Interactive workload with small probe set".to_string(),
718            }
719        } else if self.hints.prefer_hash_join {
720            EdgeJoinRecommendation::PreferHashJoin {
721                reason: "Batch workload optimized for throughput".to_string(),
722            }
723        } else {
724            EdgeJoinRecommendation::UseDefault
725        }
726    }
727}
728
729/// Join recommendation for edge computing
730#[derive(Debug, Clone)]
731pub enum EdgeJoinRecommendation {
732    /// Force nested loop due to constraints
733    ForceNestedLoop { reason: String },
734    /// Prefer nested loop (but not mandatory)
735    PreferNestedLoop { reason: String },
736    /// Prefer hash join (but not mandatory)
737    PreferHashJoin { reason: String },
738    /// Use default algorithm selection
739    UseDefault,
740}
741
742#[cfg(test)]
743mod tests {
744    use super::*;
745
746    #[test]
747    fn test_workload_learner_basic() {
748        let learner = WorkloadLearner::new();
749
750        // Record some queries
751        for i in 0..10 {
752            learner.record_query(
753                i,
754                QueryPattern::PointLookup,
755                Duration::from_micros(100),
756                1024,
757                1,
758                1,
759                vec!["users".to_string()],
760                vec!["users.id".to_string()],
761                vec![],
762            );
763        }
764
765        assert_eq!(learner.total_queries(), 10);
766
767        let stats = learner.get_pattern_stats(QueryPattern::PointLookup);
768        assert!(stats.is_some());
769        let stats = stats.unwrap();
770        assert_eq!(stats.frequency, 10);
771    }
772
773    #[test]
774    fn test_query_classification() {
775        let learner = WorkloadLearner::new();
776
777        assert_eq!(
778            learner.classify_query(true, false, false, false, 0, false, false),
779            QueryPattern::PointLookup
780        );
781
782        assert_eq!(
783            learner.classify_query(false, false, false, true, 0, false, false),
784            QueryPattern::Aggregation
785        );
786
787        assert_eq!(
788            learner.classify_query(false, false, false, true, 3, false, false),
789            QueryPattern::Analytical
790        );
791
792        assert_eq!(
793            learner.classify_query(false, false, false, false, 2, false, false),
794            QueryPattern::JoinHeavy
795        );
796
797        assert_eq!(
798            learner.classify_query(false, false, false, false, 0, true, false),
799            QueryPattern::InsertHeavy
800        );
801    }
802
803    #[test]
804    fn test_temporal_pattern_detection() {
805        let learner = WorkloadLearner::new();
806
807        // Record mostly short queries
808        for i in 0..100 {
809            learner.record_query(
810                i,
811                QueryPattern::PointLookup,
812                Duration::from_micros(500), // < 10ms = short
813                1024,
814                1,
815                1,
816                vec!["users".to_string()],
817                vec![],
818                vec![],
819            );
820        }
821
822        assert_eq!(
823            learner.detect_temporal_pattern(),
824            TemporalPattern::Interactive
825        );
826    }
827
828    #[test]
829    fn test_hot_tables() {
830        let learner = WorkloadLearner::new();
831
832        // Access 'orders' 5 times
833        for i in 0..5 {
834            learner.record_query(
835                i,
836                QueryPattern::PointLookup,
837                Duration::from_micros(100),
838                1024,
839                1,
840                1,
841                vec!["orders".to_string()],
842                vec![],
843                vec![],
844            );
845        }
846
847        // Access 'users' 3 times
848        for i in 5..8 {
849            learner.record_query(
850                i,
851                QueryPattern::PointLookup,
852                Duration::from_micros(100),
853                1024,
854                1,
855                1,
856                vec!["users".to_string()],
857                vec![],
858                vec![],
859            );
860        }
861
862        let hot = learner.hot_tables(2);
863        assert_eq!(hot.len(), 2);
864        assert_eq!(hot[0].0, "orders");
865        assert_eq!(hot[1].0, "users");
866    }
867
868    #[test]
869    fn test_edge_mode_settings() {
870        assert_eq!(EdgeMode::Standard.memory_cost_multiplier(), 1.0);
871        assert_eq!(EdgeMode::Constrained.memory_cost_multiplier(), 5.0);
872        assert_eq!(EdgeMode::UltraLow.memory_cost_multiplier(), 20.0);
873
874        assert_eq!(EdgeMode::Standard.preferred_batch_size(), 10000);
875        assert_eq!(EdgeMode::UltraLow.preferred_batch_size(), 100);
876    }
877
878    #[test]
879    fn test_edge_aware_planner() {
880        let hints = WorkloadHints {
881            prefer_nested_loop: false,
882            prefer_hash_join: true,
883            enable_bloom_filters: true,
884            target_batch_size: 1000,
885            memory_constrained: true,
886            incremental_results: false,
887        };
888
889        let planner = EdgeAwarePlanner::new(hints, 1024 * 1024); // 1MB limit
890
891        // Test cost adjustment
892        let base_cost = 100.0;
893        let adjusted = planner.adjust_cost(base_cost, 512 * 1024); // 512KB - under half
894        assert_eq!(adjusted, base_cost);
895
896        let adjusted = planner.adjust_cost(base_cost, 768 * 1024); // 768KB - over half
897        assert_eq!(adjusted, base_cost * 2.0);
898
899        let adjusted = planner.adjust_cost(base_cost, 2 * 1024 * 1024); // 2MB - over limit
900        assert_eq!(adjusted, base_cost * 100.0);
901    }
902
903    #[test]
904    fn test_edge_join_recommendation() {
905        let hints = WorkloadHints {
906            prefer_nested_loop: false,
907            prefer_hash_join: false,
908            enable_bloom_filters: false,
909            target_batch_size: 1000,
910            memory_constrained: true,
911            incremental_results: false,
912        };
913
914        let planner = EdgeAwarePlanner::new(hints, 1024 * 1024); // 1MB limit
915
916        // Hash join would need 2MB - should force nested loop
917        let rec = planner.recommend_join_for_edge(20000, 100000, 100);
918        assert!(matches!(
919            rec,
920            EdgeJoinRecommendation::ForceNestedLoop { .. }
921        ));
922
923        // Hash join fits in memory
924        let rec = planner.recommend_join_for_edge(1000, 100000, 100);
925        assert!(matches!(rec, EdgeJoinRecommendation::UseDefault));
926    }
927
928    #[test]
929    fn test_workload_config() {
930        let learner = WorkloadLearner::new();
931
932        let config = WorkloadConfig {
933            learning_enabled: true,
934            edge_mode: EdgeMode::Constrained,
935            memory_limit_mb: 512,
936            incremental_results: true,
937        };
938
939        learner.set_config(config.clone());
940        let retrieved = learner.config();
941
942        assert_eq!(retrieved.edge_mode, EdgeMode::Constrained);
943        assert_eq!(retrieved.memory_limit_mb, 512);
944        assert!(retrieved.incremental_results);
945    }
946
947    #[test]
948    fn v2_r5_config_is_authoritative_and_memory_math_is_conservative() {
949        let learner = WorkloadLearner::new();
950        let mut config = learner.config();
951        config.learning_enabled = false;
952        config.edge_mode = EdgeMode::Constrained;
953        config.memory_limit_mb = u64::MAX;
954        learner.set_config(config);
955        learner.record_query(
956            1,
957            QueryPattern::PointLookup,
958            Duration::from_millis(1),
959            1,
960            1,
961            1,
962            vec!["t".to_string()],
963            vec![],
964            vec![],
965        );
966        assert_eq!(learner.total_queries(), 0);
967
968        let planner = EdgeAwarePlanner::new(
969            WorkloadHints {
970                memory_constrained: true,
971                ..WorkloadHints::default()
972            },
973            1024,
974        );
975        assert!(planner.should_stream(u64::MAX));
976        assert!(matches!(
977            planner.recommend_join_for_edge(u64::MAX, 1, u64::MAX),
978            EdgeJoinRecommendation::ForceNestedLoop { .. }
979        ));
980    }
981
982    #[test]
983    fn test_clear() {
984        let learner = WorkloadLearner::new();
985
986        // Record some queries
987        for i in 0..10 {
988            learner.record_query(
989                i,
990                QueryPattern::PointLookup,
991                Duration::from_micros(100),
992                1024,
993                1,
994                1,
995                vec!["users".to_string()],
996                vec![],
997                vec![],
998            );
999        }
1000
1001        assert_eq!(learner.total_queries(), 10);
1002
1003        learner.clear();
1004
1005        assert_eq!(learner.total_queries(), 0);
1006        assert!(learner
1007            .get_pattern_stats(QueryPattern::PointLookup)
1008            .is_none());
1009    }
1010}