""" MVCC Executor. Performs lock-free versioned operations on Parquet files using Polars LazyFrames. """ from __future__ import annotations import os import polars as pl from config import DATA_DIR, LAZY_THRESHOLD_BYTES from services.session_manager import session_manager from core.column_registry import column_registry from core.execution.polars_builder import build_polars_expression def _build_predicate(column: str, condition: str, value) -> pl.Expr: """Build a Polars expression from a comparison condition string.""" col = pl.col(column) if condition == ">": return col > value elif condition == "<": return col < value elif condition == ">=": return col >= value elif condition == "<=": return col <= value elif condition == "!=": return col != value elif condition == "==" or condition == "=": if isinstance(value, (int, float)): return col == value try: return col == float(value) except (ValueError, TypeError): return col.cast(pl.String).str.to_lowercase() == str(value).lower() elif condition in ("contains", "~"): return col.cast(pl.String).str.contains(str(value), literal=True) else: return col.cast(pl.String).str.to_lowercase() == str(value).lower() def _format_number(val) -> str: """Pretty-print numbers with commas.""" if val is None: return "N/A" if isinstance(val, float): if val == int(val): return f"{int(val):,}" return f"{val:,.2f}" if isinstance(val, int): return f"{val:,}" return str(val) def execute_mvcc_command(session_id: str, intent: dict) -> dict: """Execute the command and return a result dict. Updates metadata and creates a new parquet version on write. """ meta = session_manager.get(session_id) if meta is None: return { "status": "error", "message": "Session nahi mila. Pehle file upload karo.", "diff": {} } version = getattr(meta, "current_version", 0) # Resolve read path # If version is 0, we can fall back to the initial upload path (data/sessions/{session_id}.parquet). v_path = os.path.join(DATA_DIR, f"{session_id}_v{version}.parquet") if not os.path.exists(v_path): if version == 0: v_path = os.path.join(DATA_DIR, f"{session_id}.parquet") if not os.path.exists(v_path): raise FileNotFoundError(f"Initial Parquet file not found for session {session_id}") else: raise FileNotFoundError(f"Parquet file not found for version {version} of session {session_id}") op = intent.get("operation") if not op: return { "message": "Operation specifies nahi kiya gaya", "diff": {} } # ── Read-only operations (Aggregates) ── if op in ("sum", "average", "count", "min", "max"): lf = pl.scan_parquet(v_path) if op == "sum": col = intent["column"] result = lf.select(pl.col(col).sum()).collect().item() return { "message": f"{col} ka sum = {_format_number(result)}", "diff": {"operation": "sum", "column": col, "result": result}, } elif op == "average": col = intent["column"] result = lf.select(pl.col(col).mean()).collect().item() return { "message": f"{col} ka average = {_format_number(result)}", "diff": {"operation": "average", "column": col, "result": result}, } elif op == "count": col = intent.get("column") if col: result = lf.select(pl.col(col).count()).collect().item() return { "message": f"{col} mein {_format_number(result)} non-null values hain", "diff": {"operation": "count", "column": col, "result": result}, } else: result = lf.select(pl.len()).collect().item() return { "message": f"Total {_format_number(result)} rows hain", "diff": {"operation": "count", "column": None, "result": result}, } elif op == "min": col = intent["column"] result = lf.select(pl.col(col).min()).collect().item() return { "message": f"{col} ka minimum = {_format_number(result)}", "diff": {"operation": "min", "column": col, "result": result}, } elif op == "max": col = intent["column"] result = lf.select(pl.col(col).max()).collect().item() return { "message": f"{col} ka maximum = {_format_number(result)}", "diff": {"operation": "max", "column": col, "result": result}, } # ── Write operations ── lf = pl.scan_parquet(v_path) new_version = version + 1 new_path = os.path.join(DATA_DIR, f"{session_id}_v{new_version}.parquet") before_count = lf.select(pl.len()).collect().item() if op in ("increase", "decrease"): col = intent["column"] value = intent.get("value") if value is None: value = 10.0 is_percent = intent.get("is_percent", True) expr = build_polars_expression(op, col, {"value": value, "is_percent": is_percent}) lf = lf.with_columns(expr) elif op == "cast_type": col = intent["column"] target = intent["target_dtype"] expr = build_polars_expression("cast_type", col, {"target_dtype": target}) lf = lf.with_columns(expr) elif op == "find_replace": col = intent["column"] old_val = intent["old_value"] new_val = intent["new_value"] schema = lf.collect_schema() dtype = schema[col] is_numeric = dtype in (pl.Float64, pl.Int64, pl.Int32, pl.Float32) if is_numeric: try: float(old_val) float(new_val) except (ValueError, TypeError): return { "message": f"'{col}' numeric hai, replace values bhi numbers hone chahiye", "diff": {}, } expr = build_polars_expression("find_replace", col, { "old_value": old_val, "new_value": new_val, "is_numeric": is_numeric }) lf = lf.with_columns(expr) elif op == "delete_column": col = intent["column"] lf = lf.drop(col) elif op == "rename_column": old = intent["column"] new = intent["new_name"] lf = lf.rename({old: new}) elif op == "add_column": new_col = intent.get("new_column_name") or intent.get("column", "new_col") default_val = intent.get("value", 0) lf = lf.with_columns(pl.lit(default_val).alias(new_col)) elif op == "remove_duplicates": col = intent.get("column") if col: lf = lf.unique(subset=[col], keep="first") else: lf = lf.unique(keep="first") elif op == "filter": col = intent["column"] condition = intent["condition"] filter_value = intent["filter_value"] predicate = _build_predicate(col, condition, filter_value) lf = lf.filter(predicate) elif op in ("sort_asc", "sort_desc"): col = intent["column"] descending = (op == "sort_desc") lf = lf.sort(col, descending=descending) else: return { "message": f"'{op}' operation supported nahi hai", "diff": {}, } # Lock-free versioned write prev_size = os.path.getsize(v_path) try: if prev_size > LAZY_THRESHOLD_BYTES: lf.sink_parquet(new_path) else: lf.collect().write_parquet(new_path) except Exception: # Fallback to eager write if streaming sink is not supported for this query plan (e.g. sort/unique) lf.collect().write_parquet(new_path) # Read updated metadata lf_new = pl.scan_parquet(new_path) new_schema = lf_new.collect_schema() after_count = lf_new.select(pl.len()).collect().item() columns_meta = [{"name": name, "dtype": str(dtype)} for name, dtype in new_schema.items()] # Format result message and diff message = "" diff = {} if op in ("increase", "decrease"): col = intent["column"] value = intent.get("value") if value is None: value = 10.0 is_percent = intent.get("is_percent", True) op_label = f"{value}% badha diya" if op == "increase" else f"{value}% ghata diya" if not is_percent: op_label = f"{value} joda" if op == "increase" else f"{value} ghata diya" message = f"{col} ko {op_label} ({after_count:,} rows updated)" diff = { "operation": op, "column": col, "value": value, "is_percent": is_percent, "affected_rows": after_count, } elif op == "cast_type": col = intent["column"] target = intent["target_dtype"] message = f"'{col}' ka type {target} mein change ho gaya" diff = {"operation": "cast_type", "column": col, "target_dtype": target} elif op == "find_replace": col = intent["column"] old_val = intent["old_value"] new_val = intent["new_value"] dtype = new_schema[col] is_numeric = dtype in (pl.Float64, pl.Int64, pl.Int32, pl.Float32) if is_numeric: try: count = lf_new.filter(pl.col(col) == float(new_val)).select(pl.len()).collect().item() except (ValueError, TypeError): count = "N/A" else: count = lf_new.filter(pl.col(col).cast(pl.String) == str(new_val)).select(pl.len()).collect().item() message = f"'{col}' mein '{old_val}' ko '{new_val}' se replace kiya ({count} rows changed)" diff = { "operation": "find_replace", "column": col, "old_value": old_val, "new_value": new_val, "affected_rows": count, } elif op == "delete_column": col = intent["column"] message = f"'{col}' column delete ho gaya" diff = {"operation": "delete_column", "column": col} elif op == "rename_column": old = intent["column"] new = intent["new_name"] message = f"'{old}' ka naam badal ke '{new}' ho gaya" diff = {"operation": "rename_column", "old_name": old, "new_name": new} elif op == "add_column": new_col = intent.get("new_column_name") or intent.get("column", "new_col") default_val = intent.get("value", 0) message = f"Naya column '{new_col}' add ho gaya (default: {default_val})" diff = {"operation": "add_column", "column": new_col} elif op == "remove_duplicates": col = intent.get("column") removed = before_count - after_count if col: message = f"'{col}' ke duplicate rows hata diye ({removed} rows removed, {after_count:,} remaining)" else: message = f"Duplicate rows hata diye ({removed} rows removed, {after_count:,} remaining)" diff = { "operation": "remove_duplicates", "column": col, "rows_removed": removed, "rows_remaining": after_count, } elif op == "filter": col = intent["column"] condition = intent["condition"] filter_value = intent["filter_value"] message = f"{col} {condition} {filter_value} → {after_count:,} rows bach gaye" diff = { "operation": "filter", "column": col, "condition": condition, "filter_value": filter_value, "rows_after": after_count, } elif op in ("sort_asc", "sort_desc"): col = intent["column"] descending = (op == "sort_desc") order = "descending (bada se chhota)" if descending else "ascending (chhota se bada)" message = f"{col} ko {order} sort kiya ({after_count:,} rows)" diff = { "operation": "sort", "column": col, "descending": descending, "affected_rows": after_count, } # Update session metadata in session_manager meta.current_version = new_version meta.row_count = after_count meta.columns = columns_meta meta.touch() # Update column registry column_registry.register(session_id, [c["name"] for c in columns_meta]) # Clean up the previous intermediate versioned file if version > 0: prev_version_path = os.path.join(DATA_DIR, f"{session_id}_v{version}.parquet") if os.path.exists(prev_version_path): try: os.remove(prev_version_path) except Exception as e: print(f"[MVCC Executor] Warning cleaning up version {version}: {e}") return { "message": message, "diff": diff, }