mangleframes 0.1.3__tar.gz → 0.1.6__tar.gz

{mangleframes-0.1.3 → mangleframes-0.1.6}/PKG-INFO

@@ -1,9 +1,12 @@
 Metadata-Version: 2.4
 Name: mangleframes
-Version: 0.1.3
+Version: 0.1.6
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Rust
 Classifier: License :: OSI Approved :: MIT License
+Requires-Dist: databricks-connect>=16.1.7
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: maturin>=1.11.2
 Requires-Dist: pyarrow>=11.0.0
 Requires-Dist: pytest>=7.0 ; extra == 'dev'
 Requires-Dist: maturin>=1.4 ; extra == 'dev'
@@ -13,4 +16,4 @@ Provides-Extra: spark
 License-File: LICENSE
 Summary: PySpark DataFrame viewer with modern web UI
 License: MIT
-Requires-Python: >=3.9
+Requires-Python: >=3.11

{mangleframes-0.1.3 → mangleframes-0.1.6}/pyproject.toml

@@ -4,9 +4,9 @@ build-backend = "maturin"
 
 [project]
 name = "mangleframes"
-version = "0.1.3"
+version = "0.1.6"
 description = "PySpark DataFrame viewer with modern web UI"
-requires-python = ">=3.9"
+requires-python = ">=3.11"
 license = { text = "MIT" }
 classifiers = [
     "Programming Language :: Python :: 3",
@@ -14,6 +14,9 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
 ]
 dependencies = [
+    "databricks-connect>=16.1.7",
+    "loguru>=0.7.3",
+    "maturin>=1.11.2",
     "pyarrow>=11.0.0",
 ]
 

mangleframes-0.1.6/python/mangleframes/__init__.py

@@ -0,0 +1,128 @@
+"""MangleFrames - PySpark DataFrame viewer with modern web UI."""
+from __future__ import annotations
+
+import atexit
+import os
+import subprocess
+import threading
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from .launcher import launch_viewer
+from .protocol import clear_arrow_cache, clear_stats_cache, prefetch_frame
+from .server import DataFrameServer
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+__version__ = "0.1.6"
+__all__ = ["register", "unregister", "show", "cleanup"]
+
+_registry: dict[str, DataFrame] = {}
+_server: DataFrameServer | None = None
+_viewer_process: subprocess.Popen | None = None
+_cleanup_registered = False
+
+
+def _clean_stale_sockets() -> None:
+    """Remove socket files from PIDs that no longer exist."""
+    socket_dir = Path("/tmp")
+    for sock in socket_dir.glob("mangleframes-*.sock"):
+        try:
+            pid_str = sock.stem.split("-")[1]
+            pid = int(pid_str)
+            if not _pid_exists(pid):
+                sock.unlink()
+        except (IndexError, ValueError, OSError):
+            pass
+
+
+def _pid_exists(pid: int) -> bool:
+    """Check if a process with given PID exists."""
+    try:
+        os.kill(pid, 0)
+        return True
+    except OSError:
+        return False
+
+
+def cleanup() -> None:
+    """Clean up all MangleFrames resources."""
+    global _server, _viewer_process, _registry
+
+    if _viewer_process is not None:
+        _viewer_process.terminate()
+        try:
+            _viewer_process.wait(timeout=2.0)
+        except subprocess.TimeoutExpired:
+            _viewer_process.kill()
+        _viewer_process = None
+
+    if _server is not None:
+        _server.stop()
+        _server = None
+
+    _registry.clear()
+    clear_stats_cache()
+    clear_arrow_cache()
+
+
+def register(name: str, df: DataFrame) -> None:
+    """Register a DataFrame for viewing."""
+    global _server, _cleanup_registered
+
+    if not _cleanup_registered:
+        atexit.register(cleanup)
+        _cleanup_registered = True
+        _clean_stale_sockets()
+
+    _registry[name] = df
+    clear_stats_cache(name)
+    clear_arrow_cache(name)
+
+    def do_prefetch() -> None:
+        prefetch_frame(_registry, name, limit=10000)
+
+    threading.Thread(target=do_prefetch, daemon=True).start()
+
+    if _server is None:
+        _server = DataFrameServer(_registry)
+        _server.start()
+
+
+def unregister(name: str) -> None:
+    """Remove a DataFrame from the viewer."""
+    if name in _registry:
+        del _registry[name]
+        clear_stats_cache(name)
+        clear_arrow_cache(name)
+
+
+def show(port: int = 8765, block: bool = True) -> None:
+    """Open the viewer in a browser, launching the viewer if needed.
+
+    Args:
+        port: Port for the viewer web server.
+        block: If True, block until Ctrl+C (keeps server alive).
+    """
+    global _viewer_process, _server
+
+    if _server is None or not _server.is_running:
+        if not _registry:
+            raise RuntimeError("No DataFrames registered. Call register() first.")
+        _server = DataFrameServer(_registry)
+        _server.start()
+
+    if _viewer_process is None or _viewer_process.poll() is not None:
+        _viewer_process = launch_viewer(_server.socket_path, port)
+
+    if block:
+        try:
+            print(f"MangleFrames viewer running at http://localhost:{port}")
+            print("Press Ctrl+C to stop...")
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            print("\nStopping MangleFrames...")
+            cleanup()

{mangleframes-0.1.3 → mangleframes-0.1.6}/python/mangleframes/launcher.py

@@ -4,17 +4,22 @@ from __future__ import annotations
 import os
 import shutil
 import subprocess
-import webbrowser
 from pathlib import Path
 
 
 def find_viewer_binary() -> Path | None:
     """Find the mangleframes-viewer binary."""
+    import sys
+
     pkg_dir = Path(__file__).parent
     pkg_binary = pkg_dir / "bin" / "mangleframes-viewer"
     if pkg_binary.exists():
         return pkg_binary
 
+    venv_binary = Path(sys.executable).parent / "mangleframes-viewer"
+    if venv_binary.exists():
+        return venv_binary
+
     path_binary = shutil.which("mangleframes-viewer")
     if path_binary:
         return Path(path_binary)
@@ -35,13 +40,10 @@ def launch_viewer(socket_path: Path, port: int = 8765) -> subprocess.Popen:
     env["RUST_LOG"] = env.get("RUST_LOG", "info")
 
     return subprocess.Popen(
-        [str(binary), "--socket", str(socket_path), "--port", str(port)],
+        [str(binary), "--socket", str(socket_path), "--port", str(port), "--no-browser"],
         env=env,
         stdout=subprocess.DEVNULL,
         stderr=subprocess.DEVNULL,
     )
 
 
-def open_browser(port: int = 8765) -> None:
-    """Open the viewer in the default browser."""
-    webbrowser.open(f"http://localhost:{port}")

mangleframes-0.1.6/python/mangleframes/protocol.py

@@ -0,0 +1,421 @@
+"""Protocol handlers for DataFrame server commands."""
+from __future__ import annotations
+
+import json
+import struct
+import threading
+import time
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+from pyspark.sql import functions as F
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame
+
+STATUS_OK = 0
+STATUS_ERROR = 1
+
+# Cache for computed stats (cleared when DataFrame is re-registered)
+_stats_cache: dict[str, dict] = {}
+
+# Cache for prefetched Arrow data: name -> (limit, payload_bytes)
+_arrow_cache: dict[str, tuple[int, bytes]] = {}
+_arrow_cache_lock = threading.Lock()
+
+
+def clear_stats_cache(name: str | None = None) -> None:
+    """Clear cached stats for a DataFrame or all DataFrames."""
+    if name is None:
+        _stats_cache.clear()
+    elif name in _stats_cache:
+        del _stats_cache[name]
+
+
+def clear_arrow_cache(name: str | None = None) -> None:
+    """Clear cached Arrow data for a DataFrame or all DataFrames."""
+    with _arrow_cache_lock:
+        if name is None:
+            _arrow_cache.clear()
+        elif name in _arrow_cache:
+            del _arrow_cache[name]
+
+
+def _serialize_arrow_ipc(table: pa.Table) -> tuple[bytes, int]:
+    """Serialize Arrow table to IPC format, returning bytes and timing in ms."""
+    start = time.perf_counter()
+    sink = pa.BufferOutputStream()
+    with pa.ipc.RecordBatchStreamWriter(sink, table.schema) as writer:
+        for batch in table.to_batches():
+            writer.write_batch(batch)
+    ipc_ms = int((time.perf_counter() - start) * 1000)
+    return sink.getvalue().to_pybytes(), ipc_ms
+
+
+def prefetch_frame(registry: dict[str, DataFrame], name: str, limit: int = 10000) -> bool:
+    """Prefetch DataFrame as Arrow IPC bytes in background.
+
+    Returns True if prefetch succeeded, False otherwise.
+    """
+    if name not in registry:
+        return False
+
+    try:
+        df = registry[name]
+        start = time.perf_counter()
+        limited_df = df.limit(limit) if limit > 0 else df
+        table = limited_df.toArrow()
+        spark_ms = int((time.perf_counter() - start) * 1000)
+        total_rows = table.num_rows
+
+        arrow_bytes, ipc_ms = _serialize_arrow_ipc(table)
+        # 24-byte header: spark_ms, ipc_ms, total_rows (all little-endian u64)
+        payload = struct.pack("<QQQ", spark_ms, ipc_ms, total_rows) + arrow_bytes
+
+        with _arrow_cache_lock:
+            _arrow_cache[name] = (limit, payload)
+
+        return True
+    except Exception:
+        return False
+
+
+def encode_response(status: int, payload: bytes) -> bytes:
+    """Encode response with status and length prefix."""
+    return struct.pack(">II", status, len(payload)) + payload
+
+
+def encode_json_response(data: Any) -> bytes:
+    """Encode JSON data as successful response."""
+    return encode_response(STATUS_OK, json.dumps(data).encode("utf-8"))
+
+
+def encode_error(message: str) -> bytes:
+    """Encode error message response."""
+    return encode_response(STATUS_ERROR, message.encode("utf-8"))
+
+
+def handle_list(registry: dict[str, DataFrame]) -> bytes:
+    """Return list of registered DataFrame names."""
+    return encode_json_response(list(registry.keys()))
+
+
+def handle_schema(registry: dict[str, DataFrame], name: str) -> bytes:
+    """Return schema of a DataFrame as JSON."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    df = registry[name]
+    columns = [
+        {"name": field.name, "type": str(field.dataType), "nullable": field.nullable}
+        for field in df.schema.fields
+    ]
+    return encode_json_response({"name": name, "columns": columns})
+
+
+def handle_get(registry: dict[str, DataFrame], name: str, limit: int) -> bytes:
+    """Return DataFrame data as Arrow IPC stream with timing info."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    # Check cache first - return cached data if limit is sufficient
+    with _arrow_cache_lock:
+        if name in _arrow_cache:
+            cached_limit, cached_payload = _arrow_cache[name]
+            if cached_limit >= limit:
+                return encode_response(STATUS_OK, cached_payload)
+
+    # Cache miss or insufficient limit - materialize from Spark
+    df = registry[name]
+
+    start = time.perf_counter()
+    limited_df = df.limit(limit) if limit > 0 else df
+    table = limited_df.toArrow()
+    spark_ms = int((time.perf_counter() - start) * 1000)
+    total_rows = table.num_rows
+
+    arrow_bytes, ipc_ms = _serialize_arrow_ipc(table)
+    # 24-byte header: spark_ms, ipc_ms, total_rows (all little-endian u64)
+    payload = struct.pack("<QQQ", spark_ms, ipc_ms, total_rows) + arrow_bytes
+
+    # Cache result for future requests
+    with _arrow_cache_lock:
+        _arrow_cache[name] = (limit, payload)
+
+    return encode_response(STATUS_OK, payload)
+
+
+def _is_numeric_type(dtype_str: str) -> bool:
+    """Check if a Spark type string represents a numeric type."""
+    dtype_lower = dtype_str.lower()
+    return any(t in dtype_lower for t in ["int", "long", "double", "float", "decimal"])
+
+
+def _is_temporal_type(dtype_str: str) -> bool:
+    """Check if a Spark type string represents a temporal type."""
+    dtype_lower = dtype_str.lower()
+    return any(t in dtype_lower for t in ["date", "timestamp"])
+
+
+def handle_stats(registry: dict[str, DataFrame], name: str) -> bytes:
+    """Return basic statistics for a DataFrame using single aggregation."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    # Return cached stats if available
+    if name in _stats_cache:
+        return encode_json_response(_stats_cache[name])
+
+    df = registry[name]
+    fields = df.schema.fields
+
+    # Build all aggregation expressions in one pass
+    agg_exprs = [F.count(F.lit(1)).alias("__total")]
+    for field in fields:
+        col_name = field.name
+        agg_exprs.append(
+            F.sum(F.when(F.col(col_name).isNull(), 1).otherwise(0)).alias(f"{col_name}__nulls")
+        )
+        if _is_numeric_type(str(field.dataType)):
+            agg_exprs.append(F.min(col_name).alias(f"{col_name}__min"))
+            agg_exprs.append(F.max(col_name).alias(f"{col_name}__max"))
+
+    # Single Spark action
+    result = df.agg(*agg_exprs).collect()[0]
+    row_count = result["__total"]
+
+    # Extract stats from result
+    column_stats = []
+    for field in fields:
+        col_name = field.name
+        dtype_str = str(field.dataType)
+        stats = {"name": col_name, "type": dtype_str, "nullable": field.nullable}
+        stats["null_count"] = result[f"{col_name}__nulls"] or 0
+
+        if _is_numeric_type(dtype_str):
+            min_val = result[f"{col_name}__min"]
+            max_val = result[f"{col_name}__max"]
+            stats["min"] = str(min_val) if min_val is not None else None
+            stats["max"] = str(max_val) if max_val is not None else None
+
+        column_stats.append(stats)
+
+    stats_data = {"name": name, "row_count": row_count, "columns": column_stats}
+    _stats_cache[name] = stats_data  # Cache for future requests
+
+    return encode_json_response(stats_data)
+
+
+def handle_count(registry: dict[str, DataFrame], name: str) -> bytes:
+    """Return total row count without transferring data."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    df = registry[name]
+    start = time.perf_counter()
+    count = df.count()
+    count_ms = int((time.perf_counter() - start) * 1000)
+
+    return encode_json_response({"name": name, "count": count, "count_ms": count_ms})
+
+
+def handle_join_keys(registry: dict[str, DataFrame], name: str, columns: list[str]) -> bytes:
+    """Return key statistics for join analysis."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    df = registry[name]
+    for col in columns:
+        if col not in [f.name for f in df.schema.fields]:
+            return encode_error(f"Column '{col}' not found in '{name}'")
+
+    if len(columns) == 1:
+        key_expr = F.col(columns[0])
+    else:
+        key_expr = F.struct(*[F.col(c) for c in columns])
+
+    start = time.perf_counter()
+    result = df.agg(
+        F.countDistinct(key_expr).alias("cardinality"),
+        F.sum(F.when(key_expr.isNull(), 1).otherwise(0)).alias("null_count"),
+        F.count(F.lit(1)).alias("total_rows")
+    ).collect()[0]
+    compute_ms = int((time.perf_counter() - start) * 1000)
+
+    return encode_json_response({
+        "frame": name,
+        "columns": columns,
+        "cardinality": result["cardinality"],
+        "null_count": result["null_count"] or 0,
+        "total_rows": result["total_rows"],
+        "compute_ms": compute_ms
+    })
+
+
+def handle_join_temporal(
+    registry: dict[str, DataFrame], name: str, column: str, bucket: str
+) -> bytes:
+    """Return temporal distribution for join coverage analysis."""
+    if name not in registry:
+        return encode_error(f"DataFrame '{name}' not found")
+
+    df = registry[name]
+    if column not in [f.name for f in df.schema.fields]:
+        return encode_error(f"Column '{column}' not found in '{name}'")
+
+    field = next(f for f in df.schema.fields if f.name == column)
+    if not _is_temporal_type(str(field.dataType)):
+        return encode_error(f"Column '{column}' is not a temporal type")
+
+    if bucket not in ("day", "week", "month"):
+        return encode_error(f"Invalid bucket: {bucket}. Use day, week, or month")
+
+    start = time.perf_counter()
+    truncated = df.withColumn("__bucket", F.date_trunc(bucket, F.col(column)))
+    buckets_df = truncated.groupBy("__bucket").agg(F.count(F.lit(1)).alias("count"))
+    bucket_rows = buckets_df.orderBy("__bucket").collect()
+    compute_ms = int((time.perf_counter() - start) * 1000)
+
+    min_max = df.agg(
+        F.min(column).alias("min_val"),
+        F.max(column).alias("max_val")
+    ).collect()[0]
+
+    buckets = {}
+    for row in bucket_rows:
+        if row["__bucket"] is not None:
+            bucket_key = str(row["__bucket"])[:10]
+            buckets[bucket_key] = row["count"]
+
+    return encode_json_response({
+        "frame": name,
+        "column": column,
+        "bucket_size": bucket,
+        "min": str(min_max["min_val"]) if min_max["min_val"] else None,
+        "max": str(min_max["max_val"]) if min_max["max_val"] else None,
+        "buckets": buckets,
+        "compute_ms": compute_ms
+    })
+
+
+def handle_join_overlap(
+    registry: dict[str, DataFrame],
+    frame1: str, frame2: str,
+    cols1: list[str], cols2: list[str]
+) -> bytes:
+    """Return key overlap statistics between two frames."""
+    if frame1 not in registry:
+        return encode_error(f"DataFrame '{frame1}' not found")
+    if frame2 not in registry:
+        return encode_error(f"DataFrame '{frame2}' not found")
+    if len(cols1) != len(cols2):
+        return encode_error("Column count mismatch")
+
+    df1 = registry[frame1]
+    df2 = registry[frame2]
+
+    for col in cols1:
+        if col not in [f.name for f in df1.schema.fields]:
+            return encode_error(f"Column '{col}' not found in '{frame1}'")
+    for col in cols2:
+        if col not in [f.name for f in df2.schema.fields]:
+            return encode_error(f"Column '{col}' not found in '{frame2}'")
+
+    start = time.perf_counter()
+
+    if len(cols1) == 1:
+        key1 = F.col(cols1[0])
+        key2 = F.col(cols2[0])
+    else:
+        key1 = F.struct(*[F.col(c) for c in cols1])
+        key2 = F.struct(*[F.col(c) for c in cols2])
+
+    keys1 = df1.select(key1.alias("key")).distinct()
+    keys2 = df2.select(key2.alias("key")).distinct()
+
+    count1 = keys1.count()
+    count2 = keys2.count()
+    both = keys1.join(keys2, "key", "inner").count()
+    left_only = count1 - both
+    right_only = count2 - both
+
+    compute_ms = int((time.perf_counter() - start) * 1000)
+
+    overlap_pct = (both / max(count1, count2) * 100) if max(count1, count2) > 0 else 0.0
+
+    return encode_json_response({
+        "frame1": frame1,
+        "frame2": frame2,
+        "cols1": cols1,
+        "cols2": cols2,
+        "left_total": count1,
+        "right_total": count2,
+        "left_only": left_only,
+        "right_only": right_only,
+        "both": both,
+        "overlap_pct": round(overlap_pct, 2),
+        "compute_ms": compute_ms
+    })
+
+
+def dispatch_command(
+    registry: dict[str, DataFrame], command: str
+) -> bytes:
+    """Parse and dispatch a command to the appropriate handler."""
+    command = command.strip()
+
+    if command == "LIST":
+        return handle_list(registry)
+
+    if command.startswith("SCHEMA:"):
+        name = command[7:]
+        return handle_schema(registry, name)
+
+    if command.startswith("GET:"):
+        parts = command[4:].split(":")
+        if len(parts) != 2:
+            return encode_error("Invalid GET format. Use GET:name:limit")
+        name, limit_str = parts
+        try:
+            limit = int(limit_str)
+        except ValueError:
+            return encode_error(f"Invalid limit: {limit_str}")
+        return handle_get(registry, name, limit)
+
+    if command.startswith("STATS:"):
+        name = command[6:]
+        return handle_stats(registry, name)
+
+    if command.startswith("COUNT:"):
+        name = command[6:]
+        return handle_count(registry, name)
+
+    if command.startswith("JOIN_KEYS:"):
+        parts = command[10:].split(":")
+        if len(parts) != 2:
+            return encode_error("Invalid JOIN_KEYS format. Use JOIN_KEYS:name:col1,col2,...")
+        name, cols_str = parts
+        columns = [c.strip() for c in cols_str.split(",") if c.strip()]
+        if not columns:
+            return encode_error("At least one column required")
+        return handle_join_keys(registry, name, columns)
+
+    if command.startswith("JOIN_TEMPORAL:"):
+        parts = command[14:].split(":")
+        if len(parts) != 3:
+            return encode_error("Invalid JOIN_TEMPORAL format. Use JOIN_TEMPORAL:name:column:bucket")
+        name, column, bucket = parts
+        return handle_join_temporal(registry, name, column, bucket)
+
+    if command.startswith("JOIN_OVERLAP:"):
+        parts = command[13:].split(":")
+        if len(parts) != 4:
+            return encode_error("Invalid JOIN_OVERLAP format. Use JOIN_OVERLAP:f1:f2:cols1:cols2")
+        frame1, frame2, cols1_str, cols2_str = parts
+        cols1 = [c.strip() for c in cols1_str.split(",") if c.strip()]
+        cols2 = [c.strip() for c in cols2_str.split(",") if c.strip()]
+        if not cols1 or not cols2:
+            return encode_error("At least one column required per frame")
+        return handle_join_overlap(registry, frame1, frame2, cols1, cols2)
+
+    return encode_error(f"Unknown command: {command}")