Memory issue

[kannankvsp]

spark.stop() is not releasing memory to os. i need to kill sail server each time to get the memory back.

Also it seems it taking more memory than apache spark.

[linhr]

Thanks for reaching out. Could you provide a minimum reproducible example?

[kannankvsp]

import time
from pathlib import Path

from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F


def timed(label: str, fn):
    start = time.perf_counter()
    result = fn()
    duration = time.perf_counter() - start
    print(f"[TIMER] {label}: {duration:.2f}s")
    return result


def build_spark(app_name: str, shuffle_partitions: int) -> SparkSession:
    spark = (
        SparkSession.builder.remote("sc://localhost:50051")
        .config("spark.sql.files.maxPartitionBytes", str(128 * 1024 * 1024))
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", str(128 * 1024 * 1024))
        .getOrCreate()
    )
    return spark


def normalize_columns(df: DataFrame) -> DataFrame:
    # Keep common ecommerce columns when present.
    preferred = [
        "event_time",
        "event_type",
        "product_id",
        "category_id",
        "category_code",
        "brand",
        "price",
        "user_id",
        "user_session",
    ]
    available = [c for c in preferred if c in df.columns]
    if available:
        df = df.select(*available)

    if "event_time" in df.columns:
        # Strip trailing UTC literal before strict timestamp parsing.
        event_time_clean = F.regexp_replace(F.col("event_time"), r"\s+UTC$", "")
        df = df.withColumn("event_ts", F.to_timestamp(event_time_clean, "yyyy-MM-dd HH:mm:ss"))
        df = df.withColumn("event_ts", F.coalesce(F.col("event_ts"), F.to_timestamp("event_time")))
    else:
        df = df.withColumn("event_ts", F.current_timestamp())

    if "price" in df.columns:
        df = df.withColumn("price", F.col("price").cast("double"))
    else:
        df = df.withColumn("price", F.lit(0.0))

    if "event_type" not in df.columns:
        df = df.withColumn("event_type", F.lit("unknown"))
    if "brand" not in df.columns:
        df = df.withColumn("brand", F.lit("unknown"))
    if "user_id" not in df.columns:
        df = df.withColumn("user_id", F.lit(None).cast("string"))

    df = (
        df.fillna({"event_type": "unknown", "brand": "unknown"})
        .withColumn("event_date", F.to_date("event_ts"))
        .withColumn("event_hour", F.hour("event_ts"))
        .withColumn(
            "price_band",
            F.when(F.col("price") < 50, "low")
            .when(F.col("price") < 200, "mid")
            .otherwise("high"),
        )
    )

    dedup_keys = [c for c in ["event_ts", "user_id", "product_id", "event_type"] if c in df.columns]
    if dedup_keys:
        df = df.dropDuplicates(dedup_keys)

    return df


def aggregate_metrics(df: DataFrame) -> DataFrame:
    return (
        df.groupBy("event_date", "event_type", "price_band")
        .agg(
            F.count("*").alias("events"),
            F.approx_count_distinct("user_id").alias("unique_users"),
            F.round(F.avg("price"), 2).alias("avg_price"),
            F.round(F.sum("price"), 2).alias("total_revenue"),
        )
        .orderBy("event_date", "event_type", "price_band")
    )


def main() -> None:
    input_path = Path("/home/kvsp/Desktop/sail/2019-Nov.csv")
    output_root = Path("/home/kvsp/Desktop/sail/output/pyspark")
    shuffle_partitions = 400

    spark = build_spark("PySpark-Perf-Test", shuffle_partitions)

    try:
        print(f"Input: {input_path}")
        print(f"Output: {output_root}")

        raw_df = timed(
            "Read CSV",
            lambda: spark.read.option("header", True)
            .option("inferSchema", False)
            .csv(str(input_path)).filter("brand='samsung'"),
        )

        raw_rows = timed("Raw row count", lambda: raw_df.count())
        print(f"Raw rows: {raw_rows}")

        transformed_df = timed("Transformations", lambda: normalize_columns(raw_df)).repartition(
            300, "event_date"
        )

        transformed_rows = timed("Transformed row count", lambda: transformed_df.count())
        print(f"Transformed rows: {transformed_rows}")

        metrics_df = timed("Build aggregates", lambda: aggregate_metrics(transformed_df))

        transformed_output = output_root / "transformed"
        metrics_output = output_root / "metrics"

        timed(
            "Write transformed Parquet",
            lambda: transformed_df.write.mode("overwrite")
            .partitionBy("event_date")
            .parquet(str(transformed_output)),
        )

        timed(
            "Write metrics Parquet",
            lambda: metrics_df.write.mode("overwrite")
            .partitionBy("event_date")
            .parquet(str(metrics_output)),
        )

        print("Performance test job completed successfully.")
    finally:
        spark.stop()


if __name__ == "__main__":
    from time import perf_counter
    a=perf_counter()
    main()
    print(f"Total execution time: {perf_counter() - a:.2f}s")

kaggle dataset dowload link: link

[linhr]

Thanks for the code snippet.

Could you share the observation around memory usage when you run the script?

[shehabgamin]

spark.stop() shouldn’t release memory, as spark.stop() closes the Spark Connect client session. In Spark Connect's architecture, the client is just a gRPC session and stopping it wouldn’t directly free server-side state.

Further, the server not immediately releasing memory after stopping is expected behavior. The retained memory is held by allocator retention, where freed memory is kept in arenas rather than immediately returned to the OS, and is not an unrecoverable leak. This also allows subsequent runs in the same process to reuse already-allocated memory instead of re-requesting it from the OS, enabling better performance.

I ran a repeated-cycle test in one process with 50 cycles to demonstrate this.

Note: The script was generated with the help of AI:

import gc
import os
import time
import psutil
from pysail.spark import SparkConnectServer
from pyspark.sql import SparkSession


CYCLES = 50
ROWS = 50_000_000
GROUPS = 100_000
POLL_SECONDS = 1
POLL_INTERVAL = 1


def read_mem():
    gc.collect()
    p = psutil.Process(os.getpid())
    m = p.memory_full_info()
    rss = m.rss / 1024**2
    uss = getattr(m, "uss", None)
    if uss is not None:
        uss = uss / 1024**2
    return rss, uss


def print_mem(label):
    rss, uss = read_mem()
    print(
        f"{label:36s} rss={rss:9.1f} MiB"
        + (f" uss={uss:9.1f} MiB" if uss is not None else ""),
        flush=True,
        )
    return rss, uss


def poll_mem(label, seconds=POLL_SECONDS, interval=POLL_INTERVAL):
    print(f"\n{label}", flush=True)

    samples = []
    for i in range(seconds + 1):
        samples.append(print_mem(f"{label} +{i:02d}s"))
        if i < seconds:
            time.sleep(interval)

    rss_values = [rss for rss, _ in samples]
    uss_values = [uss for _, uss in samples if uss is not None]

    summary = (
        f"{label} summary "
        f"rss_start={rss_values[0]:.1f} MiB "
        f"rss_end={rss_values[-1]:.1f} MiB "
        f"rss_min={min(rss_values):.1f} MiB "
        f"rss_max={max(rss_values):.1f} MiB"
    )

    if uss_values:
        summary += (
            f" uss_start={uss_values[0]:.1f} MiB "
            f"uss_end={uss_values[-1]:.1f} MiB "
            f"uss_min={min(uss_values):.1f} MiB "
            f"uss_max={max(uss_values):.1f} MiB"
        )

    print(summary, flush=True)
    return samples


def run_workload(spark):
    return (
        spark.range(0, ROWS)
        .selectExpr(f"id % {GROUPS} as k")
        .groupBy("k")
        .count()
        .collect()
    )


cycle_results = []

print_mem("process start")

for cycle in range(1, CYCLES + 1):
    print(f"\n=== cycle {cycle}/{CYCLES} ===", flush=True)

    server = SparkConnectServer()
    server.start()
    _, port = server.listening_address
    print_mem(f"cycle {cycle} server started")

    spark = SparkSession.builder.remote(f"sc://localhost:{port}").getOrCreate()
    print_mem(f"cycle {cycle} spark created")

    t0 = time.perf_counter()
    run_workload(spark)
    elapsed = time.perf_counter() - t0

    after_workload_rss, after_workload_uss = print_mem(
        f"cycle {cycle} workload {elapsed:.3f}s"
    )

    spark.stop()
    del spark
    gc.collect()
    print_mem(f"cycle {cycle} after spark.stop")

    server.stop()
    del server
    gc.collect()

    cleanup_samples = poll_mem(f"cycle {cycle} after cleanup")
    cleanup_rss, cleanup_uss = cleanup_samples[-1]

    cycle_results.append(
        {
            "cycle": cycle,
            "elapsed": elapsed,
            "after_workload_rss": after_workload_rss,
            "after_workload_uss": after_workload_uss,
            "cleanup_rss": cleanup_rss,
            "cleanup_uss": cleanup_uss,
        }
    )


print("\n=== summary ===", flush=True)
print(
    "cycle elapsed_s workload_rss_mib workload_uss_mib cleanup_rss_mib cleanup_uss_mib",
    flush=True,
)

for row in cycle_results:
    workload_uss = row["after_workload_uss"]
    cleanup_uss = row["cleanup_uss"]

    print(
        f"{row['cycle']:5d} "
        f"{row['elapsed']:9.3f} "
        f"{row['after_workload_rss']:16.1f} "
        f"{workload_uss if workload_uss is not None else float('nan'):16.1f} "
        f"{row['cleanup_rss']:15.1f} "
        f"{cleanup_uss if cleanup_uss is not None else float('nan'):15.1f}",
        flush=True,
    )

rss_values = [row["cleanup_rss"] for row in cycle_results]
uss_values = [
    row["cleanup_uss"] for row in cycle_results if row["cleanup_uss"] is not None
]
elapsed_values = [row["elapsed"] for row in cycle_results]

print(
    "\ncleanup RSS "
    f"first={rss_values[0]:.1f} MiB "
    f"last={rss_values[-1]:.1f} MiB "
    f"min={min(rss_values):.1f} MiB "
    f"max={max(rss_values):.1f} MiB "
    f"delta={rss_values[-1] - rss_values[0]:.1f} MiB",
    flush=True,
)

if uss_values:
    print(
        "cleanup USS "
        f"first={uss_values[0]:.1f} MiB "
        f"last={uss_values[-1]:.1f} MiB "
        f"min={min(uss_values):.1f} MiB "
        f"max={max(uss_values):.1f} MiB "
        f"delta={uss_values[-1] - uss_values[0]:.1f} MiB",
        flush=True,
    )

print(
    "elapsed "
    f"first={elapsed_values[0]:.3f}s "
    f"last={elapsed_values[-1]:.3f}s "
    f"min={min(elapsed_values):.3f}s "
    f"max={max(elapsed_values):.3f}s",
    flush=True,
)

=== summary ===
cycle elapsed_s workload_rss_mib workload_uss_mib cleanup_rss_mib cleanup_uss_mib
    1     0.410            254.6            191.7           254.7           191.3
    2     0.391            265.9            202.5           265.5           202.1
    3     0.393            274.5            211.1           274.1           210.7
    4     0.392            277.7            214.3           277.3           213.9
    5     0.387            288.6            225.2           288.2           224.8
    6     0.388            289.7            226.2           289.3           225.9
    7     0.397            291.1            227.7           290.8           227.3
    8     0.392            297.1            233.6           296.7           233.2
    9     0.393            299.2            235.8           298.8           235.4
   10     0.392            305.2            241.8           304.8           241.4
   11     0.396            308.6            245.2           308.2           244.8
   12     0.394            309.7            246.3           309.3           245.9
   13     0.392            312.4            249.0           312.0           248.6
   14     0.400            318.7            255.3           318.3           254.9
   15     0.397            324.0            260.6           323.6           260.2
   16     0.397            327.5            264.1           327.1           263.7
   17     0.396            333.2            269.8           331.8           268.3
   18     0.394            337.4            274.0           336.0           272.6
   19     0.403            340.5            277.1           340.1           276.7
   20     0.403            345.5            282.1           345.1           281.7
   21     0.402            348.7            285.2           348.2           284.8
   22     0.390            350.2            286.7           349.8           286.3
   23     0.391            353.8            290.4           353.4           290.0
   24     0.390            353.6            290.2           353.2           289.8
   25     0.392            357.9            294.5           357.6           294.2
   26     0.393            364.7            301.3           364.3           300.9
   27     0.388            369.0            305.6           368.6           305.2
   28     0.388            373.3            309.8           372.9           309.5
   29     0.389            375.4            312.0           375.0           311.6
   30     0.386            377.6            314.2           377.2           313.8
   31     0.392            377.8            314.4           377.4           314.0
   32     0.393            381.3            317.9           381.0           317.6
   33     0.386            381.8            318.4           381.4           318.0
   34     0.389            353.0            279.6           352.6           279.2
   35     0.384            354.9            281.6           354.5           281.2
   36     0.392            355.1            281.8           354.7           281.4
   37     0.385            356.8            283.4           356.4           283.0
   38     0.390            358.0            284.7           357.6           284.3
   39     0.388            362.5            289.1           362.1           288.8
   40     0.394            362.6            289.2           362.2           288.9
   41     0.389            335.7            230.0           335.3           229.6
   42     0.394            335.3            224.8           334.9           224.4
   43     0.389            335.5            227.2           335.2           226.8
   44     0.389            338.2            228.8           337.7           228.3
   45     0.386            336.0            225.5           335.6           225.1
   46     0.385            336.2            225.6           335.8           225.2
   47     0.388            338.6            228.4           338.2           228.0
   48     0.385            302.2            189.1           301.9           188.7
   49     0.386            302.5            189.3           302.1           188.9
   50     0.390            303.5            190.4           303.1           190.0

cleanup RSS first=254.7 MiB last=303.1 MiB min=254.7 MiB max=381.4 MiB delta=48.4 MiB
cleanup USS first=191.3 MiB last=190.0 MiB min=188.7 MiB max=318.0 MiB delta=-1.3 MiB
elapsed first=0.410s last=0.390s min=0.384s max=0.410s

USS climbed, then dropped back below the first-cycle cleanup value. That means the process did release/reclaim a large amount of private memory during the run, just not immediately after every stop().

RSS doesn't fully return to baseline,but since USS does recover, that residual is shared memory, which is consistent with normal process behavior and not indicative of a leak.

One caveat: this run doesn’t demonstrate the “much faster” effect because the process had already executed prior workloads before the 50-cycle test began, so cycle 1 isn’t a true cold start.​​​​​​​​​​​​​​​​