kotoba_query_engine/

executor.rs

//! Query Executor
//!
//! Executes optimized query plans against the graph database.

use std::sync::Arc;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use anyhow::Result;
use futures::stream::{self, StreamExt};

use crate::ast::*;
use crate::types::*;
use crate::planner::*;
use crate::{ProjectionPort, IndexManagerPort};
use kotoba_storage::KeyValueStore;

/// Query executor with KeyValueStore backend
pub struct QueryExecutor<T: KeyValueStore> {
    storage: Arc<T>,
}

impl<T: KeyValueStore + 'static> QueryExecutor<T> {
    /// Convert Vertex to serde_json::Value
    fn vertex_to_json_value(&self, vertex: Vertex) -> serde_json::Value {
        serde_json::json!({
            "id": vertex.id,
            "labels": vertex.labels,
            "properties": vertex.properties
        })
    }

    /// Convert Edge to serde_json::Value
    fn edge_to_json_value(&self, edge: Edge) -> serde_json::Value {
        serde_json::json!({
            "id": edge.id,
            "label": edge.label,
            "from_vertex": edge.from_vertex,
            "to_vertex": edge.to_vertex,
            "properties": edge.properties
        })
    }

    pub fn new(storage: Arc<T>) -> Self {
        Self { storage }
    }

    /// Execute a query plan
    pub async fn execute(
        &self,
        plan: ExecutionPlan,
        context: crate::QueryContext,
    ) -> Result<QueryResult> {
        let mut current_result = ExecutionResult::Empty;

        // Execute each step in order
        for step in plan.steps {
            current_result = match step {
                ExecutionStep::Match(match_plan) => {
                    self.execute_match(match_plan, current_result).await?
                }
                ExecutionStep::Filter(filter_plan) => {
                    self.execute_filter(filter_plan, current_result).await?
                }
                ExecutionStep::GroupBy(group_by_plan) => {
                    self.execute_group_by(group_by_plan, current_result).await?
                }
                ExecutionStep::Sort(sort_plan) => {
                    self.execute_sort(sort_plan, current_result).await?
                }
                ExecutionStep::Limit(limit_clause) => {
                    self.execute_limit(limit_clause, current_result).await?
                }
                ExecutionStep::Return(return_plan) => {
                    return self.execute_return(return_plan, current_result).await;
                }
            };
        }

        // If no return step was executed, return the current result
        Ok(QueryResult::from(current_result))
    }

    async fn execute_match(
        &self,
        match_plan: MatchPlan,
        _previous_result: ExecutionResult,
    ) -> Result<ExecutionResult> {
        let mut results = Vec::new();

        // Execute vertex scans
        for vertex_scan in match_plan.vertex_scans {
            let vertices = self.execute_vertex_scan(vertex_scan).await?;
            results.extend(vertices);
        }

        // Execute edge scans and joins
        for edge_scan in match_plan.edge_scans {
            let edges = self.execute_edge_scan(edge_scan).await?;
            // TODO: Implement proper joining logic
            results.extend(edges.into_iter().map(|e| vec![e]));
        }

        Ok(ExecutionResult::Rows(results))
    }

    async fn execute_vertex_scan(&self, scan_plan: VertexScanPlan) -> Result<Vec<Vec<Value>>> {
        let mut results = Vec::new();

        // For now, implement basic vertex scanning using KeyValueStore
        // TODO: Implement more sophisticated scanning with filters and indices

        let prefix = "vertex:".to_string();
        let vertex_keys = self.storage.scan(prefix.as_bytes()).await?;

        for key_bytes in vertex_keys {
            if let Ok(key_str) = std::str::from_utf8(&key_bytes.0) {
                if key_str.starts_with("vertex:") {
                    if let Some(vertex_data) = self.storage.get(&key_bytes.0).await? {
                        if let Ok(vertex_json) = serde_json::from_slice::<Value>(&vertex_data) {
                            results.push(vec![vertex_json]);
                        }
                    }
                }
            }
        }

        Ok(results)
    }

    async fn execute_edge_scan(&self, scan_plan: EdgeScanPlan) -> Result<Vec<Value>> {
        let mut results = Vec::new();

        // For now, implement basic edge scanning using KeyValueStore
        // TODO: Implement more sophisticated scanning with filters and indices

        let prefix = "edge:".to_string();
        let edge_keys = self.storage.scan(prefix.as_bytes()).await?;

        for key_bytes in edge_keys {
            if let Ok(key_str) = std::str::from_utf8(&key_bytes.0) {
                if key_str.starts_with("edge:") {
                    if let Some(edge_data) = self.storage.get(&key_bytes.0).await? {
                        if let Ok(edge_json) = serde_json::from_slice::<Value>(&edge_data) {
                            results.push(edge_json);
                        }
                    }
                }
            }
        }

        Ok(results)
    }

    async fn execute_filter(
        &self,
        filter_plan: FilterPlan,
        input: ExecutionResult,
    ) -> Result<ExecutionResult> {
        match input {
            ExecutionResult::Rows(rows) => {
                let mut filtered_rows = Vec::new();

                for row in rows {
                    if self.evaluate_filter(&filter_plan, &row).await? {
                        filtered_rows.push(row);
                    }
                }

                Ok(ExecutionResult::Rows(filtered_rows))
            }
            _ => Ok(input),
        }
    }

    async fn evaluate_filter(&self, _filter_plan: &FilterPlan, _row: &[Value]) -> Result<bool> {
        // TODO: Implement filter evaluation
        // For now, return true for all rows
        Ok(true)
    }

    async fn execute_group_by(
        &self,
        group_by_plan: GroupByPlan,
        input: ExecutionResult,
    ) -> Result<ExecutionResult> {
        match input {
            ExecutionResult::Rows(rows) => {
                let mut groups = std::collections::HashMap::new();

                for row in rows {
                    let key = self.compute_group_key(&group_by_plan.keys, &row).await?;
                    groups.entry(key).or_insert_with(Vec::new).push(row);
                }

                Ok(ExecutionResult::Grouped(groups))
            }
            _ => Ok(input),
        }
    }

    async fn compute_group_key(&self, _keys: &[ValueExpression], _row: &[Value]) -> Result<String> {
        // TODO: Implement group key computation
        Ok("default_group".to_string())
    }

    async fn execute_sort(
        &self,
        sort_plan: SortPlan,
        input: ExecutionResult,
    ) -> Result<ExecutionResult> {
        match input {
            ExecutionResult::Rows(mut rows) => {
                // TODO: Implement sorting based on sort keys
                // For now, just return as-is
                Ok(ExecutionResult::Rows(rows))
            }
            _ => Ok(input),
        }
    }

    async fn execute_limit(
        &self,
        limit_clause: LimitClause,
        input: ExecutionResult,
    ) -> Result<ExecutionResult> {
        match input {
            ExecutionResult::Rows(rows) => {
                let start = limit_clause.offset.unwrap_or(0) as usize;
                let end = start + limit_clause.count as usize;
                let limited_rows = rows.into_iter()
                    .skip(start)
                    .take(limit_clause.count as usize)
                    .collect();

                Ok(ExecutionResult::Rows(limited_rows))
            }
            _ => Ok(input),
        }
    }

    async fn execute_return(
        &self,
        return_plan: ReturnPlan,
        input: ExecutionResult,
    ) -> Result<QueryResult> {
        let mut results = Vec::new();

        match input {
            ExecutionResult::Rows(rows) => {
                for row in rows {
                    let mut result_row = Vec::new();

                    for item in &return_plan.items {
                        let value = self.evaluate_expression(&item.expression, &row).await?;
                        result_row.push(value);
                    }

                    results.push(result_row);
                }
            }
            ExecutionResult::Grouped(groups) => {
                // Handle grouped results
                for (_key, rows) in groups {
                    // TODO: Implement aggregation
                    if let Some(row) = rows.first() {
                        let mut result_row = Vec::new();
                        for item in &return_plan.items {
                            let value = self.evaluate_expression(&item.expression, row).await?;
                            result_row.push(value);
                        }
                        results.push(result_row);
                    }
                }
            }
            ExecutionResult::Empty => {}
        }

        // Apply DISTINCT if requested
        if return_plan.distinct {
            // TODO: Implement distinct logic
        }

        let rows_returned = results.len() as u64;
        Ok(QueryResult {
            columns: return_plan.items.iter()
                .map(|item| item.alias.clone().unwrap_or_else(|| "column".to_string()))
                .collect(),
            rows: results,
            statistics: crate::QueryStatistics {
                total_time_ms: 0,
                planning_time_ms: 0,
                execution_time_ms: 0,
                rows_scanned: 0,
                rows_returned,
                indices_used: vec![],
            },
        })
    }

    async fn evaluate_expression(&self, _expression: &ValueExpression, _row: &[serde_json::Value]) -> Result<serde_json::Value> {
        // TODO: Implement expression evaluation
        // For now, return a placeholder
        Ok(serde_json::Value::String("placeholder".to_string()))
    }
}

/// Statement executor for DDL/DML operations
pub struct StatementExecutor<T: KeyValueStore> {
    storage: Arc<T>,
}

impl<T: KeyValueStore + 'static> StatementExecutor<T> {
    pub fn new(storage: Arc<T>) -> Self {
        Self { storage }
    }

    pub async fn execute(
        &self,
        _statement: GqlStatement,
        _context: crate::QueryContext,
    ) -> Result<StatementResult> {
        // TODO: Implement statement execution
        Ok(StatementResult {
            success: true,
            message: "Statement executed successfully".to_string(),
            affected_rows: None,
            execution_time_ms: 0,
        })
    }
}

/// Execution result types
#[derive(Debug, Clone)]
pub enum ExecutionResult {
    Empty,
    Rows(Vec<Vec<serde_json::Value>>),
    Grouped(std::collections::HashMap<String, Vec<Vec<serde_json::Value>>>),
}


impl From<ExecutionResult> for QueryResult {
    fn from(result: ExecutionResult) -> Self {
        match result {
            ExecutionResult::Rows(rows) => {
                let rows_returned = rows.len() as u64;
                QueryResult {
                    columns: vec!["result".to_string()], // Placeholder
                    rows,
                    statistics: crate::QueryStatistics {
                        total_time_ms: 0,
                        planning_time_ms: 0,
                        execution_time_ms: 0,
                        rows_scanned: 0,
                        rows_returned,
                        indices_used: vec![],
                    },
                }
            },
            _ => QueryResult {
                columns: Vec::new(),
                rows: Vec::new(),
                statistics: crate::QueryStatistics {
                    total_time_ms: 0,
                    planning_time_ms: 0,
                    execution_time_ms: 0,
                    rows_scanned: 0,
                    rows_returned: 0u64,
                    indices_used: vec![],
                },
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_query_executor_creation() {
        // Test that executor can be created
        // This will be expanded with actual execution tests
    }
}