enable cloud storage ops with fsspec

pyproject.toml

@@ -10,6 +10,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "falsa>=0.0.3",
+    "fsspec>=2024.1.0",
     "pandas>=2.2.3",
     "polars>=1.24.0",
     "pydantic>=2.10.6",
@@ -20,6 +21,12 @@ dependencies = [
     "dash-ag-grid>=31.3.0",
 ]
 
+[project.optional-dependencies]
+s3 = ["s3fs>=2024.1.0"]
+azure = ["adlfs>=2024.1.0"]
+gcs = ["gcsfs>=2024.1.0"]
+cloud = ["s3fs>=2024.1.0", "adlfs>=2024.1.0", "gcsfs>=2024.1.0"]
+
 
 [tool.hatch.build.targets.wheel]
 packages = ["sparkparse"]
@@ -45,3 +52,9 @@ ignore = ["E501"]
 
 [tool.ruff.lint.isort]
 known-first-party = ["sparkparse"]
+
+[tool.pyrefly]
+project-includes = [
+    "**/*.py*",
+    "**/*.ipynb",
+]

sparkparse/analyze.py

@@ -6,6 +6,7 @@
 import polars as pl
 
 from sparkparse.models import NodeType, ParsedLogDataFrames
+from sparkparse.storage import write_text
 
 _JOIN_NODE_TYPES = frozenset(
     [
@@ -26,13 +27,20 @@ def _detail_dict(details_str: str | None) -> dict[str, Any] | None:
         return None
 
 
-def to_plan_summary(dfs: ParsedLogDataFrames, log_name: str) -> dict[str, Any]:
+def to_plan_summary(
+    dfs: ParsedLogDataFrames,
+    log_name: str,
+    out_path: str | None = None,
+) -> dict[str, Any]:
     """
     Return a token-efficient dict of execution plan data for LLM analysis.
 
     Presents raw facts (nodes, durations, bytes, join types, paths) without
     pre-assigned severity. The value over df.explain() is runtime metrics
     (per-node durations correlated from accumulator updates).
+
+    When ``out_path`` is provided, the JSON-serialized summary is also written
+    there (local path or cloud URI via storage.write_text).
     """
     dag = dfs.dag
     combined = dfs.combined
@@ -110,12 +118,17 @@ def to_plan_summary(dfs: ParsedLogDataFrames, log_name: str) -> dict[str, Any]:
         pl.sum("jvm_gc_time_seconds").alias("jvm_gc_time_seconds"),
     ).row(0, named=True)
 
-    return {
+    summary = {
         "log_name": log_name,
         "queries": queries,
         "totals": agg,
     }
 
+    if out_path is not None:
+        write_text(out_path, json.dumps(summary, indent=2, default=str))
+
+    return summary
+
 
 def find_cartesian_joins(dfs: ParsedLogDataFrames) -> pl.DataFrame:
     """
@@ -186,16 +199,12 @@ def find_largest_scans(dfs: ParsedLogDataFrames, n: int = 10) -> pl.DataFrame:
     )
 
     return (
-        scan_nodes.select(
-            "query_id", "node_id", "node_name", "details", "node_duration_minutes"
-        )
+        scan_nodes.select("query_id", "node_id", "node_name", "details", "node_duration_minutes")
         .join(node_bytes, on=["query_id", "node_name"], how="left")
         .with_columns(
             pl.col("details")
             .map_elements(
-                lambda s: json.loads(s)["detail"]["location"]["location"]
-                if s is not None
-                else [],
+                lambda s: json.loads(s)["detail"]["location"]["location"] if s is not None else [],
                 return_dtype=pl.List(pl.String),
             )
             .alias("paths")

sparkparse/app.py

@@ -10,6 +10,12 @@
 from sparkparse.dashboard import init_dashboard, run_app
 from sparkparse.models import OutputFormat, ParsedLogDataFrames
 from sparkparse.parse import get_parsed_metrics
+from sparkparse.storage import (
+    get_path_name,
+    get_path_stem,
+    is_cloud_path,
+    write_text,
+)
 
 __version__ = "0.1.0"
 
@@ -98,8 +104,8 @@ def analyze(
         typer.Option(help="Analyze a single log file instead of the whole directory."),
     ] = None,
     out_file: Annotated[
-        Path | None,
-        typer.Option(help="Write analysis JSON to this file instead of stdout."),
+        str | None,
+        typer.Option(help="Write analysis output to this file (local path or cloud URI)."),
     ] = None,
     format: Annotated[
         AnalysisFormat,
@@ -116,7 +122,7 @@ def analyze(
         verbose=False,
     )
 
-    log_name = Path(log_file).stem if log_file else Path(log_dir).name
+    log_name = get_path_stem(log_file) if log_file else get_path_name(log_dir)
     summary = to_plan_summary(dfs, log_name)
 
     if format == AnalysisFormat.json:
@@ -128,9 +134,7 @@ def analyze(
         lines.append(f"Bytes read: {totals.get('bytes_read', 0):,}")
         lines.append(f"Bytes written: {totals.get('bytes_written', 0):,}")
         lines.append(f"Shuffle bytes read: {totals.get('shuffle_bytes_read', 0):,}")
-        lines.append(
-            f"Shuffle bytes written: {totals.get('shuffle_bytes_written', 0):,}"
-        )
+        lines.append(f"Shuffle bytes written: {totals.get('shuffle_bytes_written', 0):,}")
         lines.append(f"Memory spilled: {totals.get('memory_bytes_spilled', 0):,}")
         lines.append(f"Disk spilled: {totals.get('disk_bytes_spilled', 0):,}")
         for q in summary.get("queries", []):
@@ -149,8 +153,9 @@ def analyze(
         output = "\n".join(lines)
 
     if out_file is not None:
-        out_file.parent.mkdir(parents=True, exist_ok=True)
-        out_file.write_text(output)
+        if not is_cloud_path(out_file):
+            Path(out_file).parent.mkdir(parents=True, exist_ok=True)
+        write_text(out_file, output)
         typer.echo(f"Analysis written to {out_file}", err=True)
     else:
         sys.stdout.write(output + "\n")

sparkparse/capture.py

@@ -1,21 +1,29 @@
 import functools
 import logging
 import os
-import shutil
 import subprocess
 import sys
 import tempfile
 import time
 import webbrowser
 from collections.abc import Callable
-from pathlib import Path
 from typing import Any, Literal, TypeVar, overload
 
 from pyspark.sql import SparkSession
 
 from sparkparse.analyze import to_plan_summary
 from sparkparse.app import get
 from sparkparse.models import ParsedLogDataFrames
+from sparkparse.storage import (
+    copy_file,
+    ensure_dir,
+    get_path_name,
+    get_path_stem,
+    join_path,
+    list_files,
+    path_exists,
+    remove_dir,
+)
 
 _log = logging.getLogger(__name__)
 
@@ -46,9 +54,7 @@ def __init__(
         self._parsed_logs = None
         self._analysis: dict[str, Any] | None = None
 
-    def __call__(
-        self, func: Callable[..., R]
-    ) -> Callable[..., tuple[R, "SparkparseCapture"]]:
+    def __call__(self, func: Callable[..., R]) -> Callable[..., tuple[R, "SparkparseCapture"]]:
         @functools.wraps(func)
         def get_wrapper(*args: Any, **kwargs: Any) -> tuple[R, "SparkparseCapture"]:
             with self:
@@ -70,7 +76,7 @@ def __enter__(self):
         if self.temp_dir is None:
             self._log_dir = tempfile.mkdtemp(prefix="sparkparse_")
         else:
-            os.makedirs(self.temp_dir, exist_ok=True)
+            ensure_dir(self.temp_dir)
             self._log_dir = self.temp_dir
 
         if self.spark or SparkSession.getActiveSession():
@@ -130,14 +136,14 @@ def __exit__(self, exc_type, *args):
         if self._log_dir is None:
             raise ValueError("log directory is not set")
 
-        log_dir_contents = [i for i in Path(self._log_dir).glob("*")]
-        if not any(log_dir_contents):
+        log_dir_contents = list_files(self._log_dir)
+        if not log_dir_contents:
             raise ValueError("no logs found in log directory")
 
         if self._orig_log_dir is not None:
             for f in log_dir_contents:
-                out_path = Path(self._orig_log_dir) / f.stem
-                shutil.copy2(f.as_posix(), out_path)
+                out_path = join_path(self._orig_log_dir, get_path_stem(f))
+                copy_file(f, out_path)
 
         if self.action == "viz":
             self._run_dashboard_in_background()
@@ -152,17 +158,13 @@ def __exit__(self, exc_type, *args):
                 raise ValueError("log directory is not set")
             result = get(log_dir=self._log_dir)
             self._parsed_logs = result
-            log_name = Path(self._log_dir).name
+            log_name = get_path_name(self._log_dir)
             self._analysis = to_plan_summary(result, log_name)
         else:
             raise ValueError(f"Invalid action: {self.action}")
 
-        if (
-            self._should_cleanup
-            and self._log_dir is not None
-            and os.path.exists(self._log_dir)
-        ):
-            shutil.rmtree(self._log_dir)
+        if self._should_cleanup and self._log_dir is not None and path_exists(self._log_dir):
+            remove_dir(self._log_dir)
 
 
 def capture_context(
@@ -217,9 +219,7 @@ def decorator(
         else:
             _spark = spark
 
-        cap = SparkparseCapture(
-            action, spark=_spark, temp_dir=temp_dir, headless=headless
-        )
+        cap = SparkparseCapture(action, spark=_spark, temp_dir=temp_dir, headless=headless)
         return cap(func)
 
     if func is None:

sparkparse/clean.py

@@ -1,4 +1,5 @@
 import logging
+from pathlib import Path
 
 import polars as pl
 
@@ -12,6 +13,7 @@
     Stage,
     Task,
 )
+from sparkparse.storage import is_cloud_path, join_path
 
 
 def clean_jobs(jobs: list[Job]) -> pl.DataFrame:
@@ -814,10 +816,13 @@ def write_parsed_log(
     parsed_name: str,
     suffix: str,
 ) -> None:
-    out_dir_path = resolve_dir(out_dir, 1)
-    out_dir_path.mkdir(parents=True, exist_ok=True)
-
-    out_path = out_dir_path / f"{parsed_name}_{suffix}"
+    if is_cloud_path(out_dir):
+        out_path = join_path(out_dir, f"{parsed_name}_{suffix}")
+    else:
+        out_dir_path = resolve_dir(out_dir, 1)
+        assert isinstance(out_dir_path, Path)
+        out_dir_path.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir_path / f"{parsed_name}_{suffix}"
 
     logging.info(f"Writing parsed log: {out_path}")
     logging.debug(f"Output format: {out_format}")

sparkparse/common.py

@@ -7,6 +7,7 @@
 from pyspark.sql import SparkSession
 
 from sparkparse.models import OutputFormat
+from sparkparse.storage import is_cloud_path, open_file
 
 
 def get_current_time() -> datetime.datetime:
@@ -21,28 +22,44 @@ def timeit_wrapper(*args, **kwargs):
         result = func(*args, **kwargs)
         end_time = time.perf_counter()
         total_time = end_time - start_time
-        print(
-            f"{get_current_time()} -- Function {func.__name__} Took {total_time * 1000:.2f} ms"
-        )
+        print(f"{get_current_time()} -- Function {func.__name__} Took {total_time * 1000:.2f} ms")
         return result
 
     return timeit_wrapper
 
 
 def write_dataframe(
-    df: pl.DataFrame, out_path: Path, out_format: OutputFormat, overwrite: bool = True
+    df: pl.DataFrame, out_path: str | Path, out_format: OutputFormat, overwrite: bool = True
 ) -> None:
-    if overwrite:
-        out_path.unlink(missing_ok=True)
+    out_path_str = str(out_path)
+    cloud = is_cloud_path(out_path_str)
+
+    if overwrite and not cloud:
+        Path(out_path_str).unlink(missing_ok=True)
 
     if out_format == OutputFormat.csv:
-        df.write_csv(out_path.with_suffix(".csv").as_posix(), include_header=True)
+        target = out_path_str + ".csv"
+        if cloud:
+            with open_file(target, "wb") as f:
+                df.write_csv(f)
+        else:
+            df.write_csv(target, include_header=True)
     elif out_format == OutputFormat.parquet:
-        df.write_parquet(out_path.with_suffix(".parquet").as_posix())
+        target = out_path_str + ".parquet"
+        if cloud:
+            with open_file(target, "wb") as f:
+                df.write_parquet(f)
+        else:
+            df.write_parquet(target)
     elif out_format == OutputFormat.delta:
-        df.write_delta(out_path.as_posix())
+        df.write_delta(out_path_str)
     elif out_format == OutputFormat.json:
-        df.write_json(out_path.with_suffix(".json").as_posix())
+        target = out_path_str + ".json"
+        if cloud:
+            with open_file(target, "wb") as f:
+                df.write_json(f)
+        else:
+            df.write_json(target)
 
 
 def get_spark(log_dir: Path) -> SparkSession:
@@ -74,10 +91,13 @@ def create_header(header_length: int, title: str, center: bool, spacer: str):
     return output
 
 
-def resolve_dir(incoming_dir: str | Path, default_nesting=2) -> Path:
+def resolve_dir(incoming_dir: str | Path, default_nesting=2) -> str | Path:
     # resolves path of incoming dir_str
     # if provided path does not exist, will attempt to resolve relative to sparkparse root
 
+    if is_cloud_path(str(incoming_dir)):
+        return str(incoming_dir)
+
     if isinstance(incoming_dir, Path):
         initial_path = incoming_dir
     else:
@@ -89,9 +109,7 @@ def resolve_dir(incoming_dir: str | Path, default_nesting=2) -> Path:
     path = Path(__file__).parents[default_nesting] / incoming_dir
     if not path.exists():
         if not path.parent.exists():
-            raise ValueError(
-                f"directory {path} does not exist and parent is also missing"
-            )
+            raise ValueError(f"directory {path} does not exist and parent is also missing")
         path.mkdir(exist_ok=True, parents=True)
 
     return path