genomicframe_core/

plan.rs

//! Logical query plan representation for GenomicFrame
//!
//! This module provides the lazy query plan abstraction. Plans are built up
//! through method chaining and can be optimized before execution.
//!
//! # Examples
//!
//! ```no_run
//! use genomicframe_core::plan::{LogicalPlan, PlanNode};
//! use genomicframe_core::expression::{col, lit};
//! use genomicframe_core::schema::FileFormat;
//!
//! // Build a query plan
//! let plan = LogicalPlan::scan("variants.vcf", FileFormat::Vcf)
//!     .filter(col("qual").gt(lit(30.0)))
//!     .select(&["chrom", "pos", "ref", "alt"]);
//!
//! // Optimize before execution
//! let optimized = plan.optimize();
//! ```

use crate::expression::Expr;
use crate::schema::{FileFormat, GenomicSchema};
use std::path::PathBuf;

// ============================================================================
// Core Plan Types
// ============================================================================

/// A lazy query plan that can be optimized before execution
///
/// LogicalPlan represents a series of operations (scan, filter, select, etc.)
/// without actually performing them. This enables optimization passes that can
/// reorder operations, combine filters, and push predicates down to readers.
#[derive(Debug, Clone)]
pub struct LogicalPlan {
    /// The root node of the plan tree
    pub root: PlanNode,

    /// Schema information for this plan
    pub schema: GenomicSchema,
}

/// A node in the logical plan tree
///
/// Each node represents a single operation. Nodes form a tree where
/// operations flow from leaves (scans) to the root (final output).
#[derive(Debug, Clone)]
pub enum PlanNode {
    /// Scan a file from disk
    ///
    /// This is always a leaf node - the source of data
    Scan {
        /// Path to the file
        path: PathBuf,

        /// File format
        format: FileFormat,

        /// Optional projection (which columns to read)
        projection: Option<Vec<String>>,
    },

    /// Filter records by predicate
    ///
    /// Only records where the predicate evaluates to true are kept
    Filter {
        /// Input plan to filter
        input: Box<PlanNode>,

        /// Predicate expression
        predicate: Expr,
    },

    /// Select/project specific columns
    ///
    /// Reduces the output to only the specified columns
    Select {
        /// Input plan
        input: Box<PlanNode>,

        /// Column names to keep
        columns: Vec<String>,
    },

    /// Add a computed column
    WithColumn {
        /// Input plan
        input: Box<PlanNode>,

        /// New column name
        name: String,

        /// Expression to compute column value
        expr: Expr,
    },

    /// Limit number of records
    Limit {
        /// Input plan
        input: Box<PlanNode>,

        /// Maximum number of records
        count: usize,
    },

    /// Maximum number of records to scan (before filtering)
    ///
    /// This limits how many records are read from the source,
    /// regardless of how many pass filters. Useful for sampling
    /// large files without processing the entire dataset.
    MaxScan {
        /// Input plan
        input: Box<PlanNode>,

        /// Maximum number of records to scan
        count: usize,
    },

    /// Join two plans (future support)
    Join {
        /// Left input
        left: Box<PlanNode>,

        /// Right input
        right: Box<PlanNode>,

        /// Join type
        join_type: JoinType,

        /// Join keys
        on: Vec<String>,
    },
}

/// Type of join operation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JoinType {
    /// Inner join (only matching records)
    Inner,

    /// Left outer join (all left records, matching right)
    Left,

    /// Right outer join (all right records, matching left)
    Right,

    /// Full outer join (all records from both sides)
    Full,

    /// Overlap join (genomic interval overlap)
    Overlap,
}

// ============================================================================
// LogicalPlan Builder Methods
// ============================================================================

impl LogicalPlan {
    /// Create a plan to scan a file
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("data.vcf", FileFormat::Vcf);
    /// ```
    pub fn scan<P: Into<PathBuf>>(path: P, format: FileFormat) -> Self {
        let schema = format.schema();
        Self {
            root: PlanNode::Scan {
                path: path.into(),
                format,
                projection: None,
            },
            schema,
        }
    }

    /// Add a filter to the plan
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::expression::{col, lit};
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("variants.vcf", FileFormat::Vcf)
    ///     .filter(col("qual").gt(lit(30.0)));
    /// ```
    pub fn filter(self, predicate: Expr) -> Self {
        Self {
            root: PlanNode::Filter {
                input: Box::new(self.root),
                predicate,
            },
            schema: self.schema,
        }
    }

    /// Select specific columns
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("data.vcf", FileFormat::Vcf)
    ///     .select(&["chrom", "pos", "ref", "alt"]);
    /// ```
    pub fn select(self, columns: &[&str]) -> Self {
        Self {
            root: PlanNode::Select {
                input: Box::new(self.root),
                columns: columns.iter().map(|s| s.to_string()).collect(),
            },
            schema: self.schema, // TODO: Update schema to reflect selected columns
        }
    }

    /// Add a computed column
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::expression::{col, Expr};
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("data.vcf", FileFormat::Vcf)
    ///     .with_column("is_high_qual", col("qual").gt(lit(30.0)));
    /// # use genomicframe_core::expression::lit;
    /// ```
    pub fn with_column(self, name: &str, expr: Expr) -> Self {
        Self {
            root: PlanNode::WithColumn {
                input: Box::new(self.root),
                name: name.to_string(),
                expr,
            },
            schema: self.schema, // TODO: Add new column to schema
        }
    }

    /// Limit the number of records
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("data.vcf", FileFormat::Vcf)
    ///     .limit(1000);
    /// ```
    pub fn limit(self, count: usize) -> Self {
        Self {
            root: PlanNode::Limit {
                input: Box::new(self.root),
                count,
            },
            schema: self.schema,
        }
    }

    /// Limit the maximum number of records to scan from source (before filtering)
    ///
    /// This is useful for sampling large files or preventing long-running queries
    /// on huge datasets. The scan limit is applied before any filters, so you may
    /// get fewer results than the scan limit if filters are selective.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::schema::FileFormat;
    /// use genomicframe_core::expression::{col, lit};
    ///
    /// // Scan only first 10,000 records, then filter
    /// let plan = LogicalPlan::scan("huge.vcf", FileFormat::Vcf)
    ///     .max_scan(10_000)
    ///     .filter(col("qual").gte(lit(30.0)));
    /// ```
    pub fn max_scan(self, count: usize) -> Self {
        Self {
            root: PlanNode::MaxScan {
                input: Box::new(self.root),
                count,
            },
            schema: self.schema,
        }
    }

    /// Get the file format of this plan's source
    pub fn format(&self) -> Option<FileFormat> {
        self.root.format()
    }
}

// ============================================================================
// PlanNode Helper Methods
// ============================================================================

impl PlanNode {
    /// Get the file format if this is or contains a Scan node
    pub fn format(&self) -> Option<FileFormat> {
        match self {
            PlanNode::Scan { format, .. } => Some(*format),
            PlanNode::Filter { input, .. } => input.format(),
            PlanNode::Select { input, .. } => input.format(),
            PlanNode::WithColumn { input, .. } => input.format(),
            PlanNode::Limit { input, .. } => input.format(),
            PlanNode::MaxScan { input, .. } => input.format(),
            PlanNode::Join { left, .. } => left.format(), // Use left side format
        }
    }

    /// Get all filter predicates in this plan
    pub fn filters(&self) -> Vec<&Expr> {
        let mut result = Vec::new();
        self.collect_filters(&mut result);
        result
    }

    fn collect_filters<'a>(&'a self, acc: &mut Vec<&'a Expr>) {
        match self {
            PlanNode::Filter { input, predicate } => {
                acc.push(predicate);
                input.collect_filters(acc);
            }
            PlanNode::Select { input, .. } => input.collect_filters(acc),
            PlanNode::WithColumn { input, .. } => input.collect_filters(acc),
            PlanNode::Limit { input, .. } => input.collect_filters(acc),
            PlanNode::MaxScan { input, .. } => input.collect_filters(acc),
            PlanNode::Join { left, right, .. } => {
                left.collect_filters(acc);
                right.collect_filters(acc);
            }
            PlanNode::Scan { .. } => {}
        }
    }

    /// Check if this plan contains any filters
    pub fn has_filters(&self) -> bool {
        !self.filters().is_empty()
    }
}

// ============================================================================
// Plan Optimization
// ============================================================================

impl LogicalPlan {
    /// Optimize this plan
    ///
    /// Applies multiple optimization passes:
    /// 1. Combine adjacent filters into single AND filter
    /// 2. Push filters down closer to scans
    /// 3. Prune unnecessary columns early
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use genomicframe_core::plan::LogicalPlan;
    /// use genomicframe_core::expression::{col, lit};
    /// use genomicframe_core::schema::FileFormat;
    ///
    /// let plan = LogicalPlan::scan("data.vcf", FileFormat::Vcf)
    ///     .filter(col("qual").gt(lit(30.0)))
    ///     .filter(col("chrom").eq(lit("chr1")));
    ///
    /// // Before: Scan -> Filter(qual) -> Filter(chrom)
    /// // After:  Scan -> Filter(qual AND chrom)
    /// let optimized = plan.optimize();
    /// ```
    pub fn optimize(self) -> Self {
        let mut plan = self;

        // Pass 1: Combine adjacent filters
        plan = plan.combine_filters();

        // Pass 2: Push filters down (closer to scan)
        plan = plan.push_down_filters();

        // Pass 3: Prune columns (push select down)
        plan = plan.prune_columns();

        plan
    }

    /// Combine adjacent filter operations into single filter with AND
    pub fn combine_filters(self) -> Self {
        Self {
            root: self.root.combine_filters_node(),
            schema: self.schema,
        }
    }

    /// Push filter operations down toward the scan
    pub fn push_down_filters(self) -> Self {
        Self {
            root: self.root.push_down_filters_node(),
            schema: self.schema,
        }
    }

    /// Push column selection down (prune early)
    pub fn prune_columns(self) -> Self {
        Self {
            root: self.root.prune_columns_node(),
            schema: self.schema,
        }
    }

    /// Get a human-readable representation of the plan
    pub fn explain(&self) -> String {
        self.root.explain(0)
    }
}

impl PlanNode {
    fn combine_filters_node(self) -> Self {
        match self {
            // Two adjacent filters -> combine with AND
            PlanNode::Filter { input, predicate } => {
                match *input {
                    PlanNode::Filter {
                        input: inner_input,
                        predicate: inner_predicate,
                    } => {
                        // Combine predicates
                        let combined = predicate.and(inner_predicate);
                        PlanNode::Filter {
                            input: Box::new(inner_input.combine_filters_node()),
                            predicate: combined,
                        }
                    }
                    other => PlanNode::Filter {
                        input: Box::new(other.combine_filters_node()),
                        predicate,
                    },
                }
            }
            PlanNode::Select { input, columns } => PlanNode::Select {
                input: Box::new(input.combine_filters_node()),
                columns,
            },
            PlanNode::WithColumn { input, name, expr } => PlanNode::WithColumn {
                input: Box::new(input.combine_filters_node()),
                name,
                expr,
            },
            PlanNode::Limit { input, count } => PlanNode::Limit {
                input: Box::new(input.combine_filters_node()),
                count,
            },
            PlanNode::MaxScan { input, count } => PlanNode::MaxScan {
                input: Box::new(input.combine_filters_node()),
                count,
            },
            // Scan and Join don't need changes
            other => other,
        }
    }

    fn push_down_filters_node(self) -> Self {
        match self {
            // Filter before Select -> try to push filter down
            PlanNode::Select { input, columns } => {
                match *input {
                    PlanNode::Filter {
                        input: filter_input,
                        predicate,
                    } => {
                        // Push select down, keep filter above
                        let select_below = PlanNode::Select {
                            input: filter_input,
                            columns,
                        };
                        PlanNode::Filter {
                            input: Box::new(select_below.push_down_filters_node()),
                            predicate,
                        }
                    }
                    other => PlanNode::Select {
                        input: Box::new(other.push_down_filters_node()),
                        columns,
                    },
                }
            }
            PlanNode::Filter { input, predicate } => PlanNode::Filter {
                input: Box::new(input.push_down_filters_node()),
                predicate,
            },
            PlanNode::WithColumn { input, name, expr } => PlanNode::WithColumn {
                input: Box::new(input.push_down_filters_node()),
                name,
                expr,
            },
            PlanNode::Limit { input, count } => PlanNode::Limit {
                input: Box::new(input.push_down_filters_node()),
                count,
            },
            PlanNode::MaxScan { input, count } => PlanNode::MaxScan {
                input: Box::new(input.push_down_filters_node()),
                count,
            },
            other => other,
        }
    }

    fn prune_columns_node(self) -> Self {
        // TODO: Track which columns are actually used and push projection down
        // For now, just recursively process children
        match self {
            PlanNode::Filter { input, predicate } => PlanNode::Filter {
                input: Box::new(input.prune_columns_node()),
                predicate,
            },
            PlanNode::Select { input, columns } => PlanNode::Select {
                input: Box::new(input.prune_columns_node()),
                columns,
            },
            PlanNode::WithColumn { input, name, expr } => PlanNode::WithColumn {
                input: Box::new(input.prune_columns_node()),
                name,
                expr,
            },
            PlanNode::Limit { input, count } => PlanNode::Limit {
                input: Box::new(input.prune_columns_node()),
                count,
            },
            other => other,
        }
    }

    fn explain(&self, indent: usize) -> String {
        let prefix = "  ".repeat(indent);
        match self {
            PlanNode::Scan {
                path,
                format,
                projection,
            } => {
                let proj = if let Some(cols) = projection {
                    format!(" [{}]", cols.join(", "))
                } else {
                    String::new()
                };
                format!("{}Scan: {:?} ({:?}){}", prefix, path, format, proj)
            }
            PlanNode::Filter { input, predicate } => {
                format!(
                    "{}Filter: {}\n{}",
                    prefix,
                    predicate,
                    input.explain(indent + 1)
                )
            }
            PlanNode::Select { input, columns } => {
                format!(
                    "{}Select: [{}]\n{}",
                    prefix,
                    columns.join(", "),
                    input.explain(indent + 1)
                )
            }
            PlanNode::WithColumn { input, name, expr } => {
                format!(
                    "{}WithColumn: {} = {}\n{}",
                    prefix,
                    name,
                    expr,
                    input.explain(indent + 1)
                )
            }
            PlanNode::Limit { input, count } => {
                format!("{}Limit: {}\n{}", prefix, count, input.explain(indent + 1))
            }
            PlanNode::MaxScan { input, count } => {
                format!(
                    "{}MaxScan: {}\n{}",
                    prefix,
                    count,
                    input.explain(indent + 1)
                )
            }
            PlanNode::Join {
                left,
                right,
                join_type,
                on,
            } => {
                format!(
                    "{}Join: {:?} ON [{}]\n{}{}",
                    prefix,
                    join_type,
                    on.join(", "),
                    left.explain(indent + 1),
                    right.explain(indent + 1)
                )
            }
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::expression::{col, lit};

    #[test]
    fn test_scan_plan() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf);
        assert_eq!(plan.format(), Some(FileFormat::Vcf));
    }

    #[test]
    fn test_filter_plan() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf).filter(col("qual").gt(lit(30.0)));

        assert!(plan.root.has_filters());
        assert_eq!(plan.root.filters().len(), 1);
    }

    #[test]
    fn test_chained_filters() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf)
            .filter(col("qual").gt(lit(30.0)))
            .filter(Expr::IsSnp);

        assert_eq!(plan.root.filters().len(), 2);
    }

    #[test]
    fn test_combine_filters() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf)
            .filter(col("qual").gt(lit(30.0)))
            .filter(Expr::IsSnp);

        let optimized = plan.combine_filters();

        // Should combine into single filter
        assert_eq!(optimized.root.filters().len(), 1);
    }

    #[test]
    fn test_select() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf).select(&["chrom", "pos"]);

        match plan.root {
            PlanNode::Select { columns, .. } => {
                assert_eq!(columns, vec!["chrom", "pos"]);
            }
            _ => panic!("Expected Select node"),
        }
    }

    #[test]
    fn test_limit() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf).limit(100);

        match plan.root {
            PlanNode::Limit { count, .. } => {
                assert_eq!(count, 100);
            }
            _ => panic!("Expected Limit node"),
        }
    }

    #[test]
    fn test_complex_plan() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf)
            .filter(col("qual").gt(lit(30.0)))
            .filter(Expr::IsSnp)
            .select(&["chrom", "pos", "ref", "alt"])
            .limit(1000);

        let optimized = plan.optimize();

        // Check that filters were combined
        assert_eq!(optimized.root.filters().len(), 1);
    }

    #[test]
    fn test_explain() {
        let plan = LogicalPlan::scan("test.vcf", FileFormat::Vcf)
            .filter(col("qual").gt(lit(30.0)))
            .select(&["chrom", "pos"]);

        let explanation = plan.explain();
        assert!(explanation.contains("Scan"));
        assert!(explanation.contains("Filter"));
        assert!(explanation.contains("Select"));
    }
}