@vespermcp/mcp-server 1.1.3 → 1.2.1

package/build/python/export_engine.py

@@ -2,25 +2,102 @@ import sys
 import json
 import polars as pl
 import os
+import time
+
+# Optional imports for extra formats
+try:
+    import pyarrow as pa
+    import pyarrow.feather as pf
+    HAS_PYARROW = True
+except ImportError:
+    HAS_PYARROW = False
 
-# Optional TensorFlow import for TFRecord support
 try:
     import tensorflow as tf
     HAS_TENSORFLOW = True
 except ImportError:
     HAS_TENSORFLOW = False
 
-def export_data(file_path, output_path, format, options=None):
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _load(file_path: str, options: dict) -> pl.DataFrame:
+    """Load any supported input format into a Polars DataFrame."""
+    sample_rows = options.get("sample_rows")  # int | None
+    columns = options.get("columns")          # list[str] | None
+
+    ext = os.path.splitext(file_path)[1].lower()
+    if ext == ".csv":
+        df = pl.read_csv(file_path, ignore_errors=True)
+    elif ext in (".parquet", ".pq"):
+        df = pl.read_parquet(file_path)
+    elif ext in (".feather", ".ftr", ".arrow", ".ipc"):
+        df = pl.read_ipc(file_path)
+    elif ext == ".jsonl":
+        df = pl.read_ndjson(file_path)
+    else:
+        raise ValueError(f"Unsupported input format: {ext}")
+
+    # Column selection (before sampling for speed)
+    if columns:
+        valid = [c for c in columns if c in df.columns]
+        if valid:
+            df = df.select(valid)
+
+    # Optional sampling
+    if sample_rows and sample_rows < len(df):
+        seed = options.get("random_seed", 42)
+        df = df.sample(n=sample_rows, seed=seed)
+
+    return df
+
+
+def _safe_csv_df(df: pl.DataFrame) -> pl.DataFrame:
+    """Stringify complex columns so CSV doesn't choke."""
+    for col in df.columns:
+        dtype = df.schema[col]
+        is_simple = (
+            dtype.is_numeric()
+            or dtype.is_temporal()
+            or str(dtype).lower() in ["string", "utf8", "boolean", "bool"]
+        )
+        if not is_simple:
+            def safe_serialize(val):
+                try:
+                    if hasattr(val, "to_list"):
+                        return json.dumps(val.to_list())
+                    if hasattr(val, "to_dict"):
+                        return json.dumps(val.to_dict())
+                    return json.dumps(val)
+                except Exception:
+                    return str(val)
+            df = df.with_columns(
+                pl.col(col).map_elements(safe_serialize, return_dtype=pl.Utf8)
+            )
+    return df
+
+
+def _write_preview(df: pl.DataFrame, output_path: str, n: int = 500):
+    """Write a small CSV preview next to the exported file."""
+    preview_path = os.path.splitext(output_path)[0] + "_preview.csv"
+    preview_df = _safe_csv_df(df.head(min(n, len(df))))
+    preview_df.write_csv(preview_path)
+    return preview_path
+
+
+# ---------------------------------------------------------------------------
+# Main export function
+# ---------------------------------------------------------------------------
+
+def export_data(file_path: str, output_path: str, format: str, options: dict | None = None):
     options = options or {}
-    
-    # Load Data
+    t0 = time.perf_counter()
+
+    # ---- Load ----
     try:
-        if file_path.endswith(".csv"):
-            df = pl.read_csv(file_path, ignore_errors=True)
-        elif file_path.endswith(".parquet"):
-            df = pl.read_parquet(file_path)
-        else:
-            return {"error": f"Unsupported input format: {file_path}"}
+        df = _load(file_path, options)
     except Exception as e:
         return {"error": f"Failed to load input file: {str(e)}"}
 
@@ -28,104 +105,123 @@ def export_data(file_path, output_path, format, options=None):
     if output_dir and not os.path.exists(output_dir):
         os.makedirs(output_dir, exist_ok=True)
 
+    preview_path = None
+    generate_preview = options.get("preview", False)
+
     try:
-        # Export Logic
-        if format == "csv":
-            # Stringify ANY column that might not be CSV-safe (List, Struct, Object, etc.)
-            for col in df.columns:
-                dtype = df.schema[col]
-                is_simple = (
-                    dtype.is_numeric() or 
-                    dtype.is_temporal() or 
-                    str(dtype).lower() in ["string", "utf8", "boolean", "bool"]
-                )
-                if not is_simple:
-                    def safe_serialize(val):
-                        try:
-                            if hasattr(val, "to_list"):
-                                return json.dumps(val.to_list())
-                            if hasattr(val, "to_dict"):
-                                return json.dumps(val.to_dict())
-                            return json.dumps(val)
-                        except:
-                            return str(val)
-                    df = df.with_columns(pl.col(col).map_elements(safe_serialize, return_dtype=pl.Utf8))
-            df.write_csv(output_path)
-            
+        # ---- Feather (Arrow IPC) – fastest binary format ----
+        if format == "feather":
+            if not HAS_PYARROW:
+                return {"error": "pyarrow is not installed. Run: pip install pyarrow"}
+            compression = options.get("compression", "lz4")
+            if compression in ("uncompressed", "none", "None", None):
+                compression = "uncompressed"
+            # Polars write_ipc uses Arrow IPC (= Feather v2) under the hood
+            arrow_table = df.to_arrow()
+            pf.write_feather(arrow_table, output_path, compression=compression)
+            if generate_preview:
+                preview_path = _write_preview(df, output_path)
+
+        # ---- Parquet – best compression, big-data friendly ----
         elif format == "parquet":
             compression = options.get("compression", "snappy")
+            if compression in ("uncompressed", "none", "None", None):
+                compression = "uncompressed"
             df.write_parquet(output_path, compression=compression)
-            
+            if generate_preview:
+                preview_path = _write_preview(df, output_path)
+
+        # ---- CSV – human-readable fallback ----
+        elif format == "csv":
+            df = _safe_csv_df(df)
+            df.write_csv(output_path)
+
+        # ---- JSONL ----
         elif format == "jsonl":
             df.write_ndjson(output_path)
-            
-        elif format == "arrow" or format == "ipc":
+            if generate_preview:
+                preview_path = _write_preview(df, output_path)
+
+        # ---- Arrow IPC (legacy name kept for compat) ----
+        elif format in ("arrow", "ipc"):
             compression = options.get("compression", "uncompressed")
-            if compression == "uncompressed": compression = None
+            if compression == "uncompressed":
+                compression = None
             df.write_ipc(output_path, compression=compression)
+            if generate_preview:
+                preview_path = _write_preview(df, output_path)
 
+        # ---- TFRecord ----
         elif format == "tfrecord":
             if not HAS_TENSORFLOW:
                 return {"error": "TensorFlow is not installed. Cannot export to TFRecord."}
-            
-            # TFRecord Export Logic (using TensorFlow)
             with tf.io.TFRecordWriter(output_path) as writer:
-                # Convert Polars -> Pandas for iteration (simpler for now)
-                # TODO: Optimize with Arrow -> Tensor conversion if needed for massive data
                 pdf = df.to_pandas()
                 for _, row in pdf.iterrows():
                     feature = {}
                     for col, value in row.items():
                         if value is None:
                             continue
-                        
-                        # Type inference for TFRecord features
                         if isinstance(value, int):
                             feature[col] = tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
                         elif isinstance(value, float):
                             feature[col] = tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
                         elif isinstance(value, str):
-                            feature[col] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[value.encode('utf-8')]))
+                            feature[col] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[value.encode("utf-8")]))
                         elif isinstance(value, bytes):
                             feature[col] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
                         else:
-                            # Fallback to string for unknown types
-                            feature[col] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(value).encode('utf-8')]))
-                            
+                            feature[col] = tf.train.Feature(bytes_list=tf.train.BytesList(value=[str(value).encode("utf-8")]))
                     example = tf.train.Example(features=tf.train.Features(feature=feature))
                     writer.write(example.SerializeToString())
 
         else:
             return {"error": f"Unknown export format: {format}"}
 
-        return {
+        elapsed = round(time.perf_counter() - t0, 3)
+        file_size_mb = round(os.path.getsize(output_path) / (1024 * 1024), 2)
+
+        result = {
             "success": True,
             "output_path": output_path,
             "rows": len(df),
-            "format": format
+            "columns": len(df.columns),
+            "format": format,
+            "compression": options.get("compression", "default"),
+            "file_size_mb": file_size_mb,
+            "elapsed_seconds": elapsed,
         }
+        if preview_path:
+            result["preview_path"] = preview_path
+
+        return result
 
     except Exception as e:
         return {"error": f"Export failed: {str(e)}"}
 
+
 def main():
     if len(sys.argv) < 4:
-        print(json.dumps({"error": "Usage: export_engine.py <input_file> <output_file> <format> [options_json]"}), file=sys.stderr)
+        print(
+            json.dumps({"error": "Usage: export_engine.py <input_file> <output_file> <format> [options_json]"}),
+            file=sys.stderr,
+        )
         sys.exit(1)
 
     input_file = sys.argv[1]
     output_file = sys.argv[2]
     fmt = sys.argv[3]
-    
+
     options = {}
     if len(sys.argv) > 4:
         try:
             options = json.loads(sys.argv[4])
-        except:
+        except Exception:
             pass
 
     result = export_data(input_file, output_file, fmt, options)
     print(json.dumps(result))
 
+
 if __name__ == "__main__":
     main()

package/build/python/fusion_engine.py

@@ -0,0 +1,368 @@
+import sys
+import json
+import os
+import hashlib
+import subprocess
+from typing import List, Optional, Union, Dict, Any
+
+try:
+    import polars as pl
+    HAS_POLARS = True
+except ImportError:
+    HAS_POLARS = False
+
+
+def _safe_suffix(source_path: str, idx: int) -> str:
+    base = os.path.basename(source_path)
+    base = os.path.splitext(base)[0].replace(" ", "_")
+    if not base:
+        base = f"source{idx+1}"
+    return base
+
+
+def _load_with_polars(path: str):
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".csv":
+        return pl.read_csv(path, ignore_errors=True)
+    if ext in [".parquet", ".pq"]:
+        return pl.read_parquet(path)
+    if ext in [".feather", ".ftr", ".arrow", ".ipc"]:
+        return pl.read_ipc(path)
+    if ext in [".jsonl", ".ndjson"]:
+        return pl.read_ndjson(path)
+    if ext == ".json":
+        try:
+            return pl.read_json(path)
+        except Exception:
+            return pl.read_ndjson(path)
+    raise ValueError(f"Unsupported source format: {ext} ({path})")
+
+
+def _write_with_polars(df, output_path: str, fmt: str, compression: Optional[str]):
+    fmt = fmt.lower()
+    comp = compression
+    if comp in ["none", "None", "uncompressed"]:
+        comp = None
+
+    if fmt == "csv":
+        df.write_csv(output_path)
+    elif fmt == "parquet":
+        df.write_parquet(output_path, compression=(compression or "snappy"))
+    elif fmt == "feather":
+        ipc_comp = "lz4" if compression is None else compression
+        if ipc_comp == "uncompressed":
+            ipc_comp = None
+        df.write_ipc(output_path, compression=ipc_comp)
+    elif fmt in ["arrow", "ipc"]:
+        df.write_ipc(output_path, compression=comp)
+    elif fmt == "jsonl":
+        df.write_ndjson(output_path)
+    else:
+        raise ValueError(f"Unsupported output format: {fmt}")
+
+
+def _ensure_parent(path: str):
+    parent = os.path.dirname(path)
+    if parent and not os.path.exists(parent):
+        os.makedirs(parent, exist_ok=True)
+
+
+def _compute_null_ratio(df) -> float:
+    if len(df) == 0 or len(df.columns) == 0:
+        return 0.0
+    nulls = 0
+    for col in df.columns:
+        nulls += int(df[col].null_count())
+    total_cells = len(df) * len(df.columns)
+    return (nulls / total_cells) * 100 if total_cells else 0.0
+
+
+def _leakage_report(df, source_col: str = "_vesper_source", id_col: Optional[str] = None) -> Dict[str, Any]:
+    report = {
+        "leakage_detected": False,
+        "leakage_count": 0,
+        "id_column": id_col,
+        "warnings": []
+    }
+
+    if source_col not in df.columns:
+        report["warnings"].append("Source marker column missing; leakage check skipped.")
+        return report
+
+    if id_col and id_col in df.columns:
+        overlap = (
+            df.group_by(id_col)
+            .agg(pl.col(source_col).n_unique().alias("source_count"))
+            .filter(pl.col("source_count") > 1)
+        )
+        overlap_count = len(overlap)
+        if overlap_count > 0:
+            report["leakage_detected"] = True
+            report["leakage_count"] = overlap_count
+            report["warnings"].append(f"Found {overlap_count} IDs appearing across multiple sources")
+    else:
+        # Fallback: hash rows (excluding source marker) and check if same row appears in multiple sources
+        compare_cols = [c for c in df.columns if c != source_col]
+        if not compare_cols:
+            return report
+
+        row_sig = df.select(compare_cols).with_columns(
+            pl.concat_str([pl.col(c).cast(pl.Utf8, strict=False) for c in compare_cols], separator="||").alias("_row_sig")
+        )
+        tmp = row_sig.with_columns(df[source_col]).select(["_row_sig", source_col])
+
+        overlap = (
+            tmp.group_by("_row_sig")
+            .agg(pl.col(source_col).n_unique().alias("source_count"))
+            .filter(pl.col("source_count") > 1)
+        )
+        overlap_count = len(overlap)
+        if overlap_count > 0:
+            report["leakage_detected"] = True
+            report["leakage_count"] = overlap_count
+            report["warnings"].append(f"Found {overlap_count} duplicate rows across multiple sources")
+
+    return report
+
+
+def _run_quality_engine(output_path: str) -> Dict[str, Any]:
+    # Reuse existing quality engine script (same folder)
+    try:
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        quality_script = os.path.join(script_dir, "quality_engine.py")
+        cmd = [sys.executable, quality_script, output_path]
+        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
+        if proc.returncode != 0:
+            return {"error": f"quality_engine failed: {proc.stderr.strip()}"}
+        return json.loads(proc.stdout)
+    except Exception as e:
+        return {"error": str(e)}
+
+
+def _concat_polars(dfs, source_names: List[str]):
+    # Schema alignment by union columns
+    all_cols = []
+    col_set = set()
+    for df in dfs:
+        for col in df.columns:
+            if col not in col_set:
+                all_cols.append(col)
+                col_set.add(col)
+    if "_vesper_source" not in col_set:
+        all_cols.append("_vesper_source")
+
+    aligned = []
+    for i, df in enumerate(dfs):
+        tmp = df
+        if "_vesper_source" not in tmp.columns:
+            tmp = tmp.with_columns(pl.lit(source_names[i]).alias("_vesper_source"))
+        missing = [c for c in all_cols if c not in tmp.columns]
+        for m in missing:
+            tmp = tmp.with_columns(pl.lit(None).alias(m))
+        tmp = tmp.select(all_cols)
+        aligned.append(tmp)
+
+    return pl.concat(aligned, how="vertical_relaxed")
+
+
+def _join_polars(dfs, source_names: List[str], join_on: List[str], how: str):
+    if len(dfs) < 2:
+        raise ValueError("Join strategy requires at least 2 sources")
+
+    conflict_renames = []
+    left = dfs[0]
+    if "_vesper_source" not in left.columns:
+        left = left.with_columns(pl.lit(source_names[0]).alias("_vesper_source"))
+
+    for i in range(1, len(dfs)):
+        right = dfs[i]
+        if "_vesper_source" not in right.columns:
+            right = right.with_columns(pl.lit(source_names[i]).alias("_vesper_source"))
+
+        overlaps = [
+            c for c in right.columns
+            if c in left.columns and c not in join_on
+        ]
+        rename_map = {}
+        suffix = _safe_suffix(source_names[i], i)
+        for c in overlaps:
+            new_name = f"{c}_{suffix}"
+            rename_map[c] = new_name
+            conflict_renames.append({"source": source_names[i], "from": c, "to": new_name})
+
+        if rename_map:
+            right = right.rename(rename_map)
+
+        left = left.join(right, on=join_on, how=how, coalesce=True)
+
+    return left, conflict_renames
+
+
+def fuse_datasets(
+    sources: List[Union[str, Any]],
+    strategy: str = "concat",
+    join_on: Optional[Union[str, List[str]]] = None,
+    how: str = "inner",
+    dedup: bool = True,
+    run_quality_after: bool = True,
+    leakage_check: bool = True,
+    output_path: Optional[str] = None,
+    output_format: str = "feather",
+    compression: Optional[str] = "lz4",
+    preview: bool = True,
+    preview_rows: int = 500,
+    id_column: Optional[str] = None,
+    **kwargs,
+):
+    if not HAS_POLARS:
+        return {"error": "Polars is required for dataset fusion. Install with: pip install polars"}
+
+    if not sources or len(sources) < 2:
+        return {"error": "Need at least 2 sources to fuse"}
+
+    source_paths: List[str] = []
+    source_names: List[str] = []
+
+    for i, src in enumerate(sources):
+        if isinstance(src, str):
+            source_paths.append(src)
+            source_names.append(src)
+        elif isinstance(src, dict):
+            p = src.get("path") or src.get("local_path")
+            if not p:
+                return {"error": f"Source {i} missing path"}
+            source_paths.append(p)
+            source_names.append(src.get("name") or p)
+        else:
+            return {"error": f"Unsupported source type at index {i}"}
+
+    for p in source_paths:
+        if not os.path.exists(p):
+            return {"error": f"Source not found: {p}"}
+
+    if output_path is None:
+        ext_map = {
+            "feather": ".feather",
+            "parquet": ".parquet",
+            "csv": ".csv",
+            "jsonl": ".jsonl",
+            "arrow": ".arrow",
+            "ipc": ".arrow"
+        }
+        ext = ext_map.get(output_format, ".feather")
+        output_path = os.path.abspath(f"fused_dataset{ext}")
+
+    _ensure_parent(output_path)
+
+    try:
+        dfs = [_load_with_polars(p) for p in source_paths]
+    except Exception as e:
+        return {"error": f"Failed to load sources: {str(e)}"}
+
+    rows_before = sum(len(df) for df in dfs)
+    null_before = sum(_compute_null_ratio(df) for df in dfs) / len(dfs)
+
+    strategy = (strategy or "concat").lower()
+    how = (how or "inner").lower()
+    conflict_renames = []
+
+    try:
+        if strategy == "concat":
+            fused = _concat_polars(dfs, source_names)
+        elif strategy == "join":
+            if not join_on:
+                return {"error": "join_on is required when strategy='join'"}
+            join_keys = [join_on] if isinstance(join_on, str) else list(join_on)
+            for key in join_keys:
+                for idx, df in enumerate(dfs):
+                    if key not in df.columns:
+                        return {"error": f"Join key '{key}' missing in source {source_paths[idx]}"}
+            fused, conflict_renames = _join_polars(dfs, source_names, join_keys, how)
+        else:
+            return {"error": f"Unknown strategy: {strategy}. Use concat or join."}
+
+        duplicates_removed = 0
+        if dedup:
+            before = len(fused)
+            fused = fused.unique(maintain_order=True)
+            duplicates_removed = before - len(fused)
+
+        leakage = None
+        if leakage_check:
+            leakage = _leakage_report(fused, source_col="_vesper_source", id_col=id_column)
+
+        _write_with_polars(fused, output_path, output_format, compression)
+
+        preview_path = None
+        if preview:
+            preview_path = os.path.splitext(output_path)[0] + "_preview.csv"
+            fused.head(min(preview_rows, len(fused))).write_csv(preview_path)
+
+        quality_report = None
+        if run_quality_after:
+            quality_report = _run_quality_engine(output_path)
+
+        rows_after = len(fused)
+        null_after = _compute_null_ratio(fused)
+
+        return {
+            "success": True,
+            "output_path": output_path,
+            "preview_path": preview_path,
+            "stats": {
+                "sources_count": len(source_paths),
+                "rows_before": rows_before,
+                "rows_after": rows_after,
+                "columns_after": len(fused.columns),
+                "duplicates_removed": duplicates_removed,
+                "null_ratio_before": round(null_before, 3),
+                "null_ratio_after": round(null_after, 3),
+                "null_delta": round(null_after - null_before, 3),
+                "conflict_renames": conflict_renames,
+            },
+            "quality_report": quality_report,
+            "leakage_report": leakage,
+        }
+
+    except Exception as e:
+        return {"error": f"Fusion failed: {str(e)}"}
+
+
+def main():
+    if len(sys.argv) < 3:
+        print(json.dumps({
+            "error": "Usage: fusion_engine.py <sources_json> <output_path> [config_json]"
+        }))
+        sys.exit(1)
+
+    try:
+        sources = json.loads(sys.argv[1])
+        output_path = sys.argv[2]
+        config = {}
+        if len(sys.argv) > 3:
+            config = json.loads(sys.argv[3])
+
+        result = fuse_datasets(
+            sources=sources,
+            output_path=output_path,
+            strategy=config.get("strategy", "concat"),
+            join_on=config.get("join_on"),
+            how=config.get("how", "inner"),
+            dedup=config.get("dedup", True),
+            run_quality_after=config.get("run_quality_after", True),
+            leakage_check=config.get("leakage_check", True),
+            output_format=config.get("output_format", "feather"),
+            compression=config.get("compression", "lz4"),
+            preview=config.get("preview", True),
+            preview_rows=config.get("preview_rows", 500),
+            id_column=config.get("id_column"),
+        )
+
+        print(json.dumps(result))
+    except Exception as e:
+        print(json.dumps({"error": str(e)}))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()