pgdrift_core/
index.rs

1use crate::stats::FieldStats;
2use crate::types::JsonType;
3use serde::{Deserialize, Serialize};
4
5/// Type of index to recommend
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
7pub enum IndexType {
8    /// GIN Index for high density fiels with good query support
9    Gin,
10    /// Partial index for sparse fields
11    Partial,
12    /// B-tree index on extcted scalar values
13    BTreeExtracted,
14}
15
16impl IndexType {
17    pub fn to_name(&self) -> &str {
18        match self {
19            IndexType::Gin => "GIN",
20            IndexType::Partial => "Partial GIN",
21            IndexType::BTreeExtracted => "B-tree (extracted)",
22        }
23    }
24}
25
26/// Priority of the index recommendation
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
28pub enum IndexPriority {
29    /// High priority index recommendation
30    High,
31    /// Medium priority index recommendation
32    Medium,
33    /// Low priority index recommendation
34    Low,
35}
36
37impl IndexPriority {
38    pub fn to_name(&self) -> &str {
39        match self {
40            IndexPriority::High => "High",
41            IndexPriority::Medium => "Medium",
42            IndexPriority::Low => "Low",
43        }
44    }
45}
46
47/// Index recommendation for sepecific fields
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct IndexRecommendation {
50    pub field_path: String,
51    pub index_type: IndexType,
52    pub priority: IndexPriority,
53    pub reason: String,
54    pub sql: String,
55    pub estimated_benefit: String,
56}
57
58/// Configuration for index recommendations
59#[derive(Debug, Clone)]
60pub struct IndexConfig {
61    /// Density threshold for high density fields (default: 0.8)
62    pub high_density_threshold: f64,
63    /// Density threshold for medium density fields (default: 0.2)
64    pub medium_density_threshold: f64,
65    /// Minimum occurences for index recommendation (default: 100)
66    pub min_occurences: u64,
67}
68
69impl Default for IndexConfig {
70    fn default() -> Self {
71        IndexConfig {
72            high_density_threshold: 0.8,
73            medium_density_threshold: 0.2,
74            min_occurences: 100,
75        }
76    }
77}
78
79/// Analyze field stats and generate an appropriate index recommendation if needed
80pub fn recommend_index(
81    table: &str,
82    column: &str,
83    field_stats: &[FieldStats],
84    config: &IndexConfig,
85) -> Vec<IndexRecommendation> {
86    let mut recommendations = Vec::new();
87
88    // Collect all high-density fields first to create a single consolidated GIN index
89    let high_density_fields: Vec<&FieldStats> = field_stats
90        .iter()
91        .filter(|s| {
92            s.occurrences >= config.min_occurences
93                && s.density >= config.high_density_threshold
94                && !matches!(
95                    get_dominant_type(s),
96                    Some(JsonType::Object) | Some(JsonType::Array)
97                )
98        })
99        .collect();
100
101    // If there are high-density fields, create a single GIN index recommendation
102    if !high_density_fields.is_empty() {
103        // Use the field with highest density as the primary field for the recommendation
104        let primary_field = high_density_fields
105            .iter()
106            .max_by(|a, b| {
107                a.density
108                    .partial_cmp(&b.density)
109                    .unwrap_or(std::cmp::Ordering::Equal)
110            })
111            .unwrap();
112
113        recommendations.push(create_consolidated_gin_recommendation(
114            table,
115            column,
116            primary_field,
117            &high_density_fields,
118            IndexPriority::Medium,
119        ));
120    }
121
122    // Process other recommendations (partial GIN, B-tree)
123    for stats in field_stats {
124        if stats.occurrences < config.min_occurences {
125            continue;
126        }
127
128        let dominant_type = get_dominant_type(stats);
129        if dominant_type == Some(JsonType::Object) || dominant_type == Some(JsonType::Array) {
130            continue;
131        }
132
133        // Skip high-density fields (already handled above)
134        if stats.density >= config.high_density_threshold {
135            continue;
136        }
137
138        if stats.density > 0.0 && stats.density <= config.medium_density_threshold {
139            recommendations.push(create_partial_gin_recommendation(
140                table,
141                column,
142                stats,
143                IndexPriority::Medium,
144            ));
145        } else if stats.density > config.medium_density_threshold
146            && stats.density < config.high_density_threshold
147            && is_scalar_type(dominant_type)
148        {
149            recommendations.push(create_btree_extracted_recommendation(
150                table,
151                column,
152                stats,
153                dominant_type.unwrap(),
154                IndexPriority::Medium,
155            ));
156        }
157    }
158
159    recommendations.sort_by(|a, b| {
160        let priority_order = |p: &IndexPriority| match p {
161            IndexPriority::High => 0,
162            IndexPriority::Medium => 1,
163            IndexPriority::Low => 2,
164        };
165        priority_order(&a.priority).cmp(&priority_order(&b.priority))
166    });
167
168    recommendations
169}
170
171fn get_dominant_type(stats: &FieldStats) -> Option<JsonType> {
172    stats
173        .types
174        .iter()
175        .max_by_key(|(_, count)| *count)
176        .map(|(json_type, _)| *json_type)
177}
178
179fn is_scalar_type(json_type: Option<JsonType>) -> bool {
180    matches!(
181        json_type,
182        Some(JsonType::String) | Some(JsonType::Number) | Some(JsonType::Boolean)
183    )
184}
185
186fn create_consolidated_gin_recommendation(
187    table: &str,
188    column: &str,
189    primary_stats: &FieldStats,
190    all_high_density: &[&FieldStats],
191    priority: IndexPriority,
192) -> IndexRecommendation {
193    let index_name = generate_index_name(table, column, "gin", "gin");
194
195    // Create a list of all high-density fields with their densities
196    let field_list = all_high_density
197        .iter()
198        .map(|s| format!("{} ({:.1}%)", s.path, s.density * 100.0))
199        .collect::<Vec<_>>()
200        .join(", ");
201
202    let sql = format!(
203        "-- GIN index for high-density fields: {}\n\
204        CREATE INDEX {} ON {} USING GIN ({});",
205        field_list, index_name, table, column
206    );
207
208    let reason = if all_high_density.len() == 1 {
209        format!(
210            "High density ({:.1}%) - present in {}/{} samples. \
211             GIN index enables fast JSONB queries (@>, ?, ?&, ?|)",
212            primary_stats.density * 100.0,
213            primary_stats.occurrences,
214            primary_stats.total_samples
215        )
216    } else {
217        format!(
218            "{} high-density fields ({}). \
219             Single GIN index supports fast JSONB queries (@>, ?, ?&, ?|) for all fields.",
220            all_high_density.len(),
221            field_list
222        )
223    };
224
225    IndexRecommendation {
226        field_path: primary_stats.path.clone(),
227        index_type: IndexType::Gin,
228        priority,
229        reason,
230        sql,
231        estimated_benefit: "Improved query performance for existence checks and containment queries across all high-density fields.".to_string(),
232    }
233}
234
235fn create_partial_gin_recommendation(
236    table: &str,
237    column: &str,
238    stats: &FieldStats,
239    priority: IndexPriority,
240) -> IndexRecommendation {
241    let index_name = generate_index_name(table, column, &stats.path, "partial_gin");
242    let path_condition = json_path_to_sql_conditions(&stats.path);
243    let sql = format!(
244        "-- Partial GIN index for sparse field: {:.1}% of rows contain this field\n\
245        CREATE INDEX {} ON {} USING GIN ({}) WHERE {};",
246        stats.density * 100.0,
247        index_name,
248        table,
249        column,
250        path_condition
251    );
252
253    IndexRecommendation {
254        field_path: stats.path.clone(),
255        index_type: IndexType::Partial,
256        priority,
257        reason: format!(
258            "Sparce field ({:.1}%) - only {}/{} sample have this field. \
259                Partial index reduces index size and maintenance cost",
260            stats.density * 100.0,
261            stats.occurrences,
262            stats.total_samples
263        ),
264        sql,
265        estimated_benefit: format!(
266            "Smaller index (~{:.1}% of full GIN), faster updates, same query performance for matching rows",
267            stats.density * 100.0
268        ),
269    }
270}
271
272fn create_btree_extracted_recommendation(
273    table: &str,
274    column: &str,
275    stats: &FieldStats,
276    json_type: JsonType,
277    priority: IndexPriority,
278) -> IndexRecommendation {
279    let index_name = generate_index_name(table, column, &stats.path, "btree_ext");
280    let (extraction_expr, pg_type) = match json_type {
281        JsonType::String => (
282            format!("({} #>> '{{{}}}')", column, escape_json_path(&stats.path)),
283            "TEXT",
284        ),
285        JsonType::Number => (
286            format!(
287                "(({} #>> '{{{}}}')::NUMERIC)",
288                column,
289                escape_json_path(&stats.path)
290            ),
291            "NUMERIC",
292        ),
293        JsonType::Boolean => (
294            format!(
295                "(({} #>> '{{{}}}')::BOOLEAN)",
296                column,
297                escape_json_path(&stats.path)
298            ),
299            "BOOLEAN",
300        ),
301        _ => unreachable!(),
302    };
303
304    let sql = format!(
305        "-- B-tree index on extracted {} value: {:.1}% density\n\
306        CREATE INDEX {} ON {} ({}) WHERE {} IS NOT NULL;",
307        pg_type,
308        stats.density * 100.0,
309        index_name,
310        table,
311        extraction_expr,
312        extraction_expr
313    );
314
315    IndexRecommendation {
316        field_path: stats.path.clone(),
317        index_type: IndexType::BTreeExtracted,
318        priority,
319        reason: format!(
320            "Medium density {} field ({:.1}%) - {}/{} samples. \
321             B-tree index on extracted value range query and sorting.",
322            json_type,
323            stats.density * 100.0,
324            stats.occurrences,
325            stats.total_samples
326        ),
327        sql,
328        estimated_benefit:
329            "Improved query performance for lookups and range queries on scalar values.".to_string(),
330    }
331}
332
333fn generate_index_name(table: &str, column: &str, path: &str, index_type: &str) -> String {
334    let clean_path = path
335        .replace("[]", "_arr")
336        .replace(".", "_")
337        .chars()
338        .filter(|c| c.is_alphanumeric() || *c == '_')
339        .collect::<String>();
340
341    let max_len = 63;
342    let prefix = format!("idx_{}_{}_{}_{}", table, column, clean_path, index_type);
343
344    if prefix.len() <= max_len {
345        prefix
346    } else {
347        let hash = format!("{:x}", calculate_simple_hash(&prefix));
348        let truncate_len = max_len - hash.len() - 1;
349        format!("{}_{}", &prefix[..truncate_len], hash)
350    }
351}
352
353fn calculate_simple_hash(s: &str) -> u32 {
354    s.bytes()
355        .fold(0u32, |hash, b| hash.wrapping_mul(31).wrapping_add(b as u32))
356}
357
358fn json_path_to_sql_conditions(path: &str) -> String {
359    let parts: Vec<&str> = path.split('.').collect();
360    if parts.len() == 1 {
361        let clean_part = parts[0].replace("[]", "");
362        format!("metadata ? '{}'", clean_part)
363    } else {
364        let parent_path = parts[..parts.len() - 1]
365            .iter()
366            .map(|p| p.replace("[]", ""))
367            .collect::<Vec<_>>()
368            .join(",");
369        let last = parts.last().unwrap().replace("[]", "");
370        format!("metadata #> '{{{}}}' ? '{}'", parent_path, last)
371    }
372}
373
374fn escape_json_path(path: &str) -> String {
375    path.replace("[]", "") // Remove array notation
376        .replace('\'', "''") // Escape single quotes for SQL
377        .replace('.', ",") // Convert dots to commas for PostgreSQL {a,b,c} syntax
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    fn create_test_stats(path: &str, density: f64, occurrences: u64, total: u64) -> FieldStats {
385        let mut stats = FieldStats::new(path.to_string(), 1);
386        stats.occurrences = occurrences;
387        stats.total_samples = total;
388        stats.density = density;
389        stats
390    }
391
392    #[test]
393    fn test_high_density_recommends_gin() {
394        let mut stats = create_test_stats("user.email", 0.95, 9500, 10000);
395        stats.types.insert(JsonType::String, 9500);
396
397        let config = IndexConfig::default();
398        let recommendations = recommend_index("users", "metadata", &[stats], &config);
399
400        assert_eq!(recommendations.len(), 1);
401        assert_eq!(recommendations[0].index_type, IndexType::Gin);
402        assert_eq!(recommendations[0].priority, IndexPriority::Medium);
403        assert!(recommendations[0].sql.contains("CREATE INDEX"));
404        assert!(recommendations[0].sql.contains("USING GIN"));
405        assert!(recommendations[0].sql.contains("95.0%"));
406    }
407
408    #[test]
409    fn test_sparse_field_recommends_partial_gin() {
410        let mut stats = create_test_stats("billing.legacy_plan", 0.05, 100, 2000);
411        stats.types.insert(JsonType::String, 100);
412
413        let config = IndexConfig::default();
414        let recommendations = recommend_index("users", "metadata", &[stats], &config);
415
416        assert_eq!(recommendations.len(), 1);
417        assert_eq!(recommendations[0].index_type, IndexType::Partial);
418        assert_eq!(recommendations[0].priority, IndexPriority::Medium);
419        assert!(recommendations[0].sql.contains("WHERE"));
420        assert!(recommendations[0].sql.contains("5.0%"));
421        assert!(recommendations[0].estimated_benefit.contains("5.0%"));
422    }
423
424    #[test]
425    fn test_medium_density_string_recommends_btree() {
426        let mut stats = create_test_stats("username", 0.5, 500, 1000);
427        stats.types.insert(JsonType::String, 500);
428
429        let config = IndexConfig::default();
430        let recommendations = recommend_index("users", "metadata", &[stats], &config);
431
432        assert_eq!(recommendations.len(), 1);
433        assert_eq!(recommendations[0].index_type, IndexType::BTreeExtracted);
434        assert_eq!(recommendations[0].priority, IndexPriority::Medium);
435        assert!(recommendations[0].sql.contains("#>>"));
436        assert!(recommendations[0].sql.contains("TEXT"));
437    }
438
439    #[test]
440    fn test_medium_density_number_recommends_btree() {
441        let mut stats = create_test_stats("age", 0.6, 600, 1000);
442        stats.types.insert(JsonType::Number, 600);
443
444        let config = IndexConfig::default();
445        let recommendations = recommend_index("users", "metadata", &[stats], &config);
446
447        assert_eq!(recommendations.len(), 1);
448        assert_eq!(recommendations[0].index_type, IndexType::BTreeExtracted);
449        assert!(recommendations[0].sql.contains("::NUMERIC"));
450        assert!(recommendations[0].sql.contains("NUMERIC"));
451    }
452
453    #[test]
454    fn test_medium_density_boolean_recommends_btree() {
455        let mut stats = create_test_stats("is_active", 0.45, 450, 1000);
456        stats.types.insert(JsonType::Boolean, 450);
457
458        let config = IndexConfig::default();
459        let recommendations = recommend_index("users", "metadata", &[stats], &config);
460
461        assert_eq!(recommendations.len(), 1);
462        assert_eq!(recommendations[0].index_type, IndexType::BTreeExtracted);
463        assert!(recommendations[0].sql.contains("::BOOLEAN"));
464        assert!(recommendations[0].sql.contains("BOOLEAN"));
465    }
466
467    #[test]
468    fn test_skips_object_types() {
469        let mut stats = create_test_stats("user", 0.9, 900, 1000);
470        stats.types.insert(JsonType::Object, 900);
471
472        let config = IndexConfig::default();
473        let recommendations = recommend_index("users", "metadata", &[stats], &config);
474
475        assert_eq!(recommendations.len(), 0);
476    }
477
478    #[test]
479    fn test_skips_array_types() {
480        let mut stats = create_test_stats("tags", 0.9, 900, 1000);
481        stats.types.insert(JsonType::Array, 900);
482
483        let config = IndexConfig::default();
484        let recommendations = recommend_index("users", "metadata", &[stats], &config);
485
486        assert_eq!(recommendations.len(), 0);
487    }
488
489    #[test]
490    fn test_skips_low_occurrence_fields() {
491        let mut stats = create_test_stats("rare_field", 0.9, 50, 55);
492        stats.types.insert(JsonType::String, 50);
493
494        let config = IndexConfig::default();
495        let recommendations = recommend_index("users", "metadata", &[stats], &config);
496
497        assert_eq!(recommendations.len(), 0);
498    }
499
500    #[test]
501    fn test_respects_min_occurrences_threshold() {
502        let mut stats = create_test_stats("field", 0.9, 99, 110);
503        stats.types.insert(JsonType::String, 99);
504
505        let config = IndexConfig::default();
506        let recommendations = recommend_index("users", "metadata", &[stats], &config);
507
508        // Should skip because 99 < 100 (default min)
509        assert_eq!(recommendations.len(), 0);
510
511        let mut stats2 = create_test_stats("field2", 0.9, 100, 111);
512        stats2.types.insert(JsonType::String, 100);
513
514        let recommendations2 = recommend_index("users", "metadata", &[stats2], &config);
515
516        // Should recommend because 100 >= 100
517        assert_eq!(recommendations2.len(), 1);
518    }
519
520    #[test]
521    fn test_custom_high_density_threshold() {
522        let mut stats = create_test_stats("field", 0.7, 700, 1000);
523        stats.types.insert(JsonType::String, 700);
524
525        let config = IndexConfig {
526            high_density_threshold: 0.6,
527            medium_density_threshold: 0.2,
528            min_occurences: 100,
529        };
530
531        let recommendations = recommend_index("users", "metadata", &[stats], &config);
532
533        assert_eq!(recommendations.len(), 1);
534        assert_eq!(recommendations[0].index_type, IndexType::Gin);
535    }
536
537    #[test]
538    fn test_custom_medium_density_threshold() {
539        let mut stats = create_test_stats("field", 0.15, 150, 1000);
540        stats.types.insert(JsonType::String, 150);
541
542        let config = IndexConfig {
543            high_density_threshold: 0.8,
544            medium_density_threshold: 0.1,
545            min_occurences: 100,
546        };
547
548        let recommendations = recommend_index("users", "metadata", &[stats], &config);
549
550        // Should get BTree because 0.15 > 0.1 (medium threshold) and < 0.8 (high threshold)
551        assert_eq!(recommendations.len(), 1);
552        assert_eq!(recommendations[0].index_type, IndexType::BTreeExtracted);
553    }
554
555    #[test]
556    fn test_index_name_generation_basic() {
557        let name = generate_index_name("users", "metadata", "user.email", "gin");
558        assert_eq!(name, "idx_users_metadata_user_email_gin");
559        assert!(name.len() <= 63);
560    }
561
562    #[test]
563    fn test_index_name_generation_with_arrays() {
564        let name = generate_index_name("orders", "data", "items[].sku", "btree_ext");
565        assert!(name.contains("items_arr_sku"));
566        assert!(name.len() <= 63);
567    }
568
569    #[test]
570    fn test_index_name_generation_truncation() {
571        let long_path = "very.long.deeply.nested.path.that.will.definitely.exceed.the.postgresql.limit.for.index.names";
572        let name = generate_index_name("table_with_very_long_name", "column", long_path, "gin");
573        assert!(name.len() <= 63);
574        assert!(name.starts_with("idx_"));
575    }
576
577    #[test]
578    fn test_index_name_special_chars_removed() {
579        let name = generate_index_name("users", "data", "user-email@domain", "gin");
580        // Special chars should be filtered out
581        assert!(!name.contains('@'));
582        assert!(!name.contains('-'));
583    }
584
585    #[test]
586    fn test_json_path_escaping() {
587        // Dots are converted to commas for PostgreSQL path syntax
588        assert_eq!(escape_json_path("user.email"), "user,email");
589        assert_eq!(escape_json_path("tags[]"), "tags");
590        // Single quotes are escaped, dots become commas
591        assert_eq!(escape_json_path("user's.name"), "user''s,name");
592        // Array notation removed, dots become commas
593        assert_eq!(escape_json_path("items[].price"), "items,price");
594    }
595
596    #[test]
597    fn test_json_path_to_sql_conditions_simple() {
598        let condition = json_path_to_sql_conditions("email");
599        assert_eq!(condition, "metadata ? 'email'");
600    }
601
602    #[test]
603    fn test_json_path_to_sql_conditions_nested() {
604        let condition = json_path_to_sql_conditions("user.profile.email");
605        assert_eq!(condition, "metadata #> '{user,profile}' ? 'email'");
606    }
607
608    #[test]
609    fn test_json_path_to_sql_conditions_with_array() {
610        let condition = json_path_to_sql_conditions("tags[]");
611        assert_eq!(condition, "metadata ? 'tags'");
612    }
613
614    #[test]
615    fn test_priority_sorting() {
616        let mut stats1 = create_test_stats("low_priority", 0.5, 500, 1000);
617        stats1.types.insert(JsonType::String, 500);
618
619        let mut stats2 = create_test_stats("high_priority", 0.95, 950, 1000);
620        stats2.types.insert(JsonType::String, 950);
621
622        let mut stats3 = create_test_stats("medium_priority", 0.1, 100, 1000);
623        stats3.types.insert(JsonType::String, 100);
624
625        let config = IndexConfig::default();
626        let recommendations =
627            recommend_index("users", "metadata", &[stats1, stats2, stats3], &config);
628
629        assert_eq!(recommendations.len(), 3);
630        // All should be Medium priority in this implementation
631        assert!(recommendations[0].priority == IndexPriority::Medium);
632        assert!(recommendations[1].priority == IndexPriority::Medium);
633        assert!(recommendations[2].priority == IndexPriority::Medium);
634    }
635
636    #[test]
637    fn test_multiple_recommendations() {
638        let mut stats1 = create_test_stats("email", 0.95, 950, 1000);
639        stats1.types.insert(JsonType::String, 950);
640
641        let mut stats2 = create_test_stats("age", 0.6, 600, 1000);
642        stats2.types.insert(JsonType::Number, 600);
643
644        let mut stats3 = create_test_stats("legacy_id", 0.05, 100, 2000);
645        stats3.types.insert(JsonType::String, 100);
646
647        let config = IndexConfig::default();
648        let recommendations =
649            recommend_index("users", "metadata", &[stats1, stats2, stats3], &config);
650
651        assert_eq!(recommendations.len(), 3);
652
653        // Find each type
654        let gin = recommendations
655            .iter()
656            .find(|r| r.index_type == IndexType::Gin);
657        let btree = recommendations
658            .iter()
659            .find(|r| r.index_type == IndexType::BTreeExtracted);
660        let partial = recommendations
661            .iter()
662            .find(|r| r.index_type == IndexType::Partial);
663
664        assert!(gin.is_some());
665        assert!(btree.is_some());
666        assert!(partial.is_some());
667    }
668
669    #[test]
670    fn test_multiple_high_density_creates_single_gin() {
671        let mut stats1 = create_test_stats("email", 0.95, 950, 1000);
672        stats1.types.insert(JsonType::String, 950);
673
674        let mut stats2 = create_test_stats("name", 0.92, 920, 1000);
675        stats2.types.insert(JsonType::String, 920);
676
677        let mut stats3 = create_test_stats("status", 0.88, 880, 1000);
678        stats3.types.insert(JsonType::String, 880);
679
680        let mut stats4 = create_test_stats("user_id", 1.0, 1000, 1000);
681        stats4.types.insert(JsonType::String, 1000);
682
683        let config = IndexConfig::default();
684        let recommendations = recommend_index(
685            "users",
686            "metadata",
687            &[stats1, stats2, stats3, stats4],
688            &config,
689        );
690
691        // Should generate single GIN index, not four
692        let gin_count = recommendations
693            .iter()
694            .filter(|r| r.index_type == IndexType::Gin)
695            .count();
696        assert_eq!(
697            gin_count, 1,
698            "Should generate single GIN index for multiple high-density fields"
699        );
700
701        // Verify the recommendation mentions all fields
702        let gin_rec = recommendations
703            .iter()
704            .find(|r| r.index_type == IndexType::Gin)
705            .unwrap();
706        assert!(gin_rec.reason.contains("email") || gin_rec.sql.contains("email"));
707        assert!(gin_rec.reason.contains("name") || gin_rec.sql.contains("name"));
708        assert!(gin_rec.reason.contains("status") || gin_rec.sql.contains("status"));
709        assert!(gin_rec.reason.contains("user_id") || gin_rec.sql.contains("user_id"));
710        assert!(gin_rec.reason.contains("4 high-density fields"));
711    }
712
713    #[test]
714    fn test_get_dominant_type() {
715        let mut stats = create_test_stats("mixed_field", 0.5, 500, 1000);
716        stats.types.insert(JsonType::String, 450);
717        stats.types.insert(JsonType::Number, 50);
718
719        let dominant = get_dominant_type(&stats);
720        assert_eq!(dominant, Some(JsonType::String));
721    }
722
723    #[test]
724    fn test_is_scalar_type() {
725        assert!(is_scalar_type(Some(JsonType::String)));
726        assert!(is_scalar_type(Some(JsonType::Number)));
727        assert!(is_scalar_type(Some(JsonType::Boolean)));
728        assert!(!is_scalar_type(Some(JsonType::Object)));
729        assert!(!is_scalar_type(Some(JsonType::Array)));
730        assert!(!is_scalar_type(Some(JsonType::Null)));
731        assert!(!is_scalar_type(None));
732    }
733
734    #[test]
735    fn test_index_type_to_name() {
736        assert_eq!(IndexType::Gin.to_name(), "GIN");
737        assert_eq!(IndexType::Partial.to_name(), "Partial GIN");
738        assert_eq!(IndexType::BTreeExtracted.to_name(), "B-tree (extracted)");
739    }
740
741    #[test]
742    fn test_index_priority_to_name() {
743        assert_eq!(IndexPriority::High.to_name(), "High");
744        assert_eq!(IndexPriority::Medium.to_name(), "Medium");
745        assert_eq!(IndexPriority::Low.to_name(), "Low");
746    }
747
748    #[test]
749    fn test_calculate_simple_hash() {
750        let hash1 = calculate_simple_hash("test");
751        let hash2 = calculate_simple_hash("test");
752        let hash3 = calculate_simple_hash("different");
753
754        // Same input should produce same hash
755        assert_eq!(hash1, hash2);
756        // Different input should (likely) produce different hash
757        assert_ne!(hash1, hash3);
758    }
759
760    #[test]
761    fn test_recommendation_contains_field_path() {
762        let mut stats = create_test_stats("user.profile.email", 0.9, 900, 1000);
763        stats.types.insert(JsonType::String, 900);
764
765        let config = IndexConfig::default();
766        let recommendations = recommend_index("users", "metadata", &[stats], &config);
767
768        assert_eq!(recommendations.len(), 1);
769        assert_eq!(recommendations[0].field_path, "user.profile.email");
770    }
771
772    #[test]
773    fn test_gin_recommendation_has_complete_sql() {
774        let mut stats = create_test_stats("email", 0.95, 950, 1000);
775        stats.types.insert(JsonType::String, 950);
776
777        let config = IndexConfig::default();
778        let recommendations = recommend_index("users", "metadata", &[stats], &config);
779
780        let sql = &recommendations[0].sql;
781        assert!(sql.contains("CREATE INDEX"));
782        assert!(sql.contains("USING GIN"));
783        assert!(sql.contains("users"));
784        assert!(sql.contains("metadata"));
785        assert!(sql.starts_with("--")); // Has comment
786    }
787
788    #[test]
789    fn test_partial_gin_has_where_clause() {
790        let mut stats = create_test_stats("rare.field", 0.05, 100, 2000);
791        stats.types.insert(JsonType::String, 100);
792
793        let config = IndexConfig::default();
794        let recommendations = recommend_index("users", "metadata", &[stats], &config);
795
796        let sql = &recommendations[0].sql;
797        assert!(sql.contains("WHERE"));
798        assert!(sql.contains("metadata"));
799    }
800
801    #[test]
802    fn test_btree_has_extraction_and_where() {
803        let mut stats = create_test_stats("score", 0.5, 500, 1000);
804        stats.types.insert(JsonType::Number, 500);
805
806        let config = IndexConfig::default();
807        let recommendations = recommend_index("users", "metadata", &[stats], &config);
808
809        let sql = &recommendations[0].sql;
810        assert!(sql.contains("#>>"));
811        assert!(sql.contains("::NUMERIC"));
812        assert!(sql.contains("WHERE"));
813        assert!(sql.contains("IS NOT NULL"));
814    }
815}
pgdrift_core/index.rs

pgdrift_core/
index.rs
