Spark Utilities

The Spark utilities module provides a comprehensive collection of PySpark functions for data manipulation, mathematical operations, and data processing.

Module Overview

siege_utilities.distributed.spark_utils.sanitise_dataframe_column_names(df)[source]

Cleans dataframe column names by converting them to lowercase and replacing slashes/spaces with underscores.

Parameters:

df (DataFrame) – Input Spark DataFrame.

Returns:

Sanitised DataFrame.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.tabulate_null_vs_not_null(df, column_name)[source]

Returns a dataframe showing the count of null and non-null values for a given column.

Parameters:

  • df (DataFrame) – Input Spark DataFrame.

  • column_name (str) – Name of the column to analyze.

Returns:

Resulting DataFrame with null vs non-null counts.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.get_row_count(df)[source]

Returns the count of rows in the dataframe.

Parameters:

df (DataFrame) – Input Spark DataFrame.

Returns:

Row count.

Return type:

int

siege_utilities.distributed.spark_utils.repartition_and_cache(df, partitions=100)[source]

Repartitions and caches a dataframe.

Parameters:

  • df (DataFrame) – Input Spark DataFrame.

  • partitions (int, optional) – Number of partitions. Default is 100.

Returns:

Repartitioned and cached DataFrame.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.register_temp_table(df, table_name)[source]

Registers a temporary view from a dataframe.

Parameters:

  • df (DataFrame) – Input Spark DataFrame.

  • table_name (str) – Name for the temporary view.

Return type:

None

siege_utilities.distributed.spark_utils.move_column_to_front_of_dataframe(df, column_name)[source]

Reorder df so column_name is the leftmost column.

The Spark schema order is what most CSV / parquet readers surface first, and downstream consumers (notebook displays, exports) read left-to-right. Moving the join key / identifier to the front makes the rest of the pipeline more readable without changing semantics.

Parameters:

  • df (pyspark.sql.DataFrame) – Input DataFrame. Not mutated.

  • column_name (str) – Name of the column to move to position 0. Must be present in df.columns (raises AnalysisException from the underlying select otherwise).

Returns:

New DataFrame with the same rows and columns as df but column_name in the first position. The original df is unchanged. The new schema is printed via .printSchema() as a side effect (for interactive inspection).

Return type:

pyspark.sql.DataFrame

siege_utilities.distributed.spark_utils.write_df_to_parquet(df, path, mode='overwrite')[source]

Writes a DataFrame to a Parquet file.

Parameters:

  • df (DataFrame) – Input Spark DataFrame.

  • path (str) – Output path.

  • mode (str) – Write mode. Defaults to “overwrite”.

Return type:

None

siege_utilities.distributed.spark_utils.read_parquet_to_df(spark, path)[source]

Reads a Parquet file into a Spark DataFrame.

Parameters:

  • spark (SparkSession) – Active Spark session.

  • path (str) – Path to the Parquet file.

Returns:

Loaded DataFrame.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.flatten_json_column_and_join_back_to_df(df, json_column, prefix='json_column_', logger=None, drop_original=True, explode_arrays=False, flatten_level='shallow', verbose=False, sample_size=5, show_samples=True)[source]

Flattens a JSON column in a Spark DataFrame, extracting fields and adding them as columns. Has fallback mechanisms for corrupt JSON data.

Parameters:

  • df (DataFrame) – The input Spark DataFrame.

  • json_column (str) – The name of the column containing JSON strings.

  • prefix (str, optional) – Prefix to add to the flattened column names. Defaults to “json_column_”.

  • logger (Optional[Any], optional) – Logger object for logging messages. Defaults to None.

  • drop_original (bool, optional) – Whether to drop the original JSON column after flattening. Defaults to True.

  • explode_arrays (bool, optional) – Whether to explode array columns. Defaults to False.

  • flatten_level (str, optional) – “shallow” or “deep” flattening. Defaults to “shallow”.

  • verbose (bool, optional) – Controls whether to log detailed messages. Defaults to False.

  • sample_size (int, optional) – Number of samples to check. Defaults to 5.

  • show_samples (bool, optional) – Whether to display sample data. Defaults to False.

Returns:

The DataFrame with the JSON column flattened.

Return type:

DataFrame

Raises:

  • ValueError – If all JSON samples are corrupt and schema cannot be inferred.

  • RuntimeError – If Spark analysis fails and the fallback path also fails.

siege_utilities.distributed.spark_utils.validate_geocode_data(df, lat_col_name, lon_col_name)[source]

Filters out rows with invalid geographic coordinates using string-based column names.

Raises:

ValueError – If lat_col_name or lon_col_name is not in the DataFrame.

Parameters:

  • lat_col_name (str)

  • lon_col_name (str)

siege_utilities.distributed.spark_utils.mark_valid_geocode_data(df, lat_col_name, lon_col_name, output_col_name='is_valid')[source]

Adds a boolean flag column to the DataFrame indicating whether the geographic coordinates are valid.

A set of coordinates is considered valid if: - The latitude and longitude columns are not null. - The latitude is between -90 and 90. - The longitude is between -180 and 180.

Unlike filtering functions, this function preserves all rows in the DataFrame by simply marking each row with a True (valid) or False (invalid) value in the new output column.

Parameters:

  • df (DataFrame) – The Spark DataFrame containing geocode data.

  • lat_col_name (str) – The name of the latitude column.

  • lon_col_name (str) – The name of the longitude column.

  • output_col_name (str, optional) – The name of the output column to store the validity flag. Defaults to “is_valid”.

Raises:

ValueError – If lat_col_name or lon_col_name is not in the DataFrame.

Returns:

A new DataFrame with an additional column indicating geocode validity.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.clean_and_reorder_bbox(df, bbox_col)[source]

Normalize a bracketed bounding-box string column into Sedona array form.

Parameters:

  • df (pyspark.sql.DataFrame) – Input DataFrame containing a bounding-box string column.

  • bbox_col (str) – Name of the string column holding the bbox. Each value is expected to be a comma-separated list optionally wrapped in [...], in the order: min_latitude, max_latitude, min_longitude, max_longitude.

Returns:

The input DataFrame extended with several derived columns ({bbox_col}_cleaned, _split, _min_lat, _max_lat, _min_lon, _max_lon) and a final {bbox_col}_reordered column containing an array in the Sedona-canonical [min_lon, min_lat, max_lon, max_lat] order. The original column is preserved.

Return type:

pyspark.sql.DataFrame

Notes

Sedona spatial predicates expect coordinates in (lon, lat) order, but many upstream data sources emit (lat, lon) first — this helper adapts the input to what Sedona will accept without manual reordering per pipeline.

siege_utilities.distributed.spark_utils.ensure_literal(value)[source]

Convert any value to a Spark literal (Column) unless it is already a Spark Column.

Parameters:

value – Any value to be converted.

Returns:

A pyspark.sql.Column containing the value (or its Spark literal), unless the value is already a Column.

Return type:

None

siege_utilities.distributed.spark_utils.reproject_geom_columns(df, geom_columns, source_srid, target_srid)[source]

Reprojects geometry columns using the three-argument version of ST_Transform: ST_Transform(geom, ‘source_srid’, ‘target_srid’)

Only reprojects if the current SRID is not equal to the target.

Parameters:

  • df (DataFrame) – Spark DataFrame containing the geometry columns.

  • geom_columns (list) – List of column names (strings) to reproject.

  • source_srid (str) – The source CRS (e.g. “EPSG:4326”).

  • target_srid (str) – The target CRS (e.g. “EPSG:27700”).

Raises:

ValueError – If source_srid or target_srid is not a valid EPSG identifier.

Returns:

The DataFrame with each specified geometry column conditionally reprojected.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.prepare_dataframe_for_export(df, logger_func=None)[source]
Prepares a DataFrame for export (e.g., to CSV) by:

  • Converting binary columns to Base64-encoded strings.

  • Casting simple scalar fields (non-string, non-complex) to strings.

  • Dropping intermediate columns (e.g., ‘parsed_json’) if present.

  • Converting complex (StructType/ArrayType) columns to JSON strings.

  • Handling null values appropriately.

Parameters:

  • df – Spark DataFrame to prepare

  • logger_func – Optional logging function (defaults to print)

Returns:

The transformed DataFrame with all columns as strings or JSON strings.

siege_utilities.distributed.spark_utils.prepare_summary_dataframe(data_tuples, column_names=None, logger_func=None)[source]

Helper function to create summary DataFrames with consistent string types. Prevents type merging errors by ensuring all values are strings.

Parameters:

  • data_tuples – List of tuples with data

  • column_names – Column names for the DataFrame

  • logger_func – Optional logging function

Raises:

RuntimeError – If no active Spark session is found.

Returns:

Spark DataFrame with all string columns

siege_utilities.distributed.spark_utils.export_pyspark_df_to_excel(df, file_name='output.xlsx', sheet_name='Sheet1')[source]

Converts a PySpark DataFrame to a Pandas DataFrame and exports it to an Excel file.

Parameters:

  • spark_df (pyspark.sql.DataFrame) – The PySpark DataFrame to export.

  • file_name (str) – The name of the output Excel file.

  • sheet_name (str) – The sheet name in the Excel file.

siege_utilities.distributed.spark_utils.pivot_summary_table_for_bools(df, columns, spark)[source]

Generate a pivot table summary for given boolean flag columns in a DataFrame. The pivot table includes three metrics:

  • “Count”: Sum of rows where the flag is True.

  • “Percentage (%)”: Percentage relative to total records.

  • “Total”: The total number of records (repeated for each column).

All numeric values are converted to float to ensure a consistent type.

Parameters:

  • df (DataFrame) – The source Spark DataFrame.

  • columns (list) – List of column names (assumed to be boolean flags) to summarize.

  • spark (SparkSession) – The active Spark session.

Returns:

A Spark DataFrame representing the pivot table.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.pivot_summary_with_metrics(df, group_col, pivot_col, spark)[source]

Generate a pivot summary for a categorical column against one or more grouping columns, including rows for “Count”, “Percentage (%)”, and “Total” for each group.

Parameters:

  • df (DataFrame) – The source Spark DataFrame.

  • group_col (str or list) – The column name (or list of column names) used for grouping. For example, “geocode_granularity” or [“state”, “region”].

  • pivot_col (str) – The categorical column to pivot on (e.g., “final_geocode_choice”).

  • spark (SparkSession) – The active Spark session.

Returns:

A Spark DataFrame in which each original group appears as three rows:

one for the counts, one for the percentages, and one for the total count. The non-grouping columns represent each distinct pivot column value.

Return type:

DataFrame

siege_utilities.distributed.spark_utils.export_prepared_df_as_csv_to_path_using_delimiter(df, write_path, delimiter=',')[source]

Exports DataFrame with necessary transformations to ensure Spark compatibility.

Parameters:

  • df (None) – Target DataFrame

  • write_path (Path) – Pathlib object for export destination

  • delimiter (str) – CSV delimiter (default is comma)

Return type:

None

Applies prepare_dataframe_for_export() to prevent Spark export issues.

siege_utilities.distributed.spark_utils.print_debug_table(spark_df, title)[source]

Log a Spark DataFrame as a formatted table for interactive debugging.

Converts spark_df to pandas, formats it with tabulate in psql grid style, and logs it via log_info() under the given title.

Parameters:

  • spark_df (pyspark.sql.DataFrame) – DataFrame to display. Should be small enough to fit in driver memory (calls .toPandas() internally).

  • title (str) – Header line printed before the table.

Returns:

Side effect only — writes to the configured logger.

Return type:

None

Raises:

ImportError – If the optional tabulate dependency is not installed. The error message names the pip install command that unblocks the call site.

siege_utilities.distributed.spark_utils.compute_walkability(distance)[source]

Bucket a distance-in-meters into a walkability grade label.

Used to classify how far one place is from another in pedestrian terms — the thresholds come from urban-planning literature: Trivial (<100m), Tolerable, Moderate, Borderline, Outside (>500m). Returns a {"grade": ..., "label": ...} dict, or None when distance is None (so the caller’s .withColumn(... apply udf) produces a null instead of crashing).

The actual thresholds live in walkability_config above for audit / per-deployment tweaking.

Return type:

Dict[str, str] | None

siege_utilities.distributed.spark_utils.validate_geometry(df, geom_col, step_name)[source]

Validates a single geometry column.

Parameters: - df (DataFrame): Spark DataFrame containing geometry data. - geom_col (str): Name of the geometry column to check. - step_name (str): Label for the debug output.

siege_utilities.distributed.spark_utils.backup_full_dataframe(df, step_name)[source]

Persist an intermediate DataFrame snapshot for post-run inspection.

A debug-mode helper: long Spark pipelines occasionally need an out-of-band snapshot of an intermediate frame (the canonical .cache() lives in memory and dies with the job). This writes the snapshot to DEBUG_SUBDIRECTORY/{step_name}_full_persisted in the project’s standard results output format (see RESULTS_OUTPUT_FORMAT / RESULTS_OUTPUT_DELIMITER) so it can be inspected after the run.

Parameters:

  • df (pyspark.sql.DataFrame) – DataFrame to snapshot. Wide/deep frames are written in full — do not use on frames larger than debug-storage budget.

  • step_name (str) – Human-readable label used to build the output directory name. Ensure uniqueness across steps in the same run.

Returns:

Side effect only — writes to DEBUG_SUBDIRECTORY and emits a confirming INFO log line.

Return type:

None

siege_utilities.distributed.spark_utils.atomic_write_with_staging(df, final_destination, staging_directory, file_format='csv', delimiter=',', header=True, mode='overwrite')[source]

Write a DataFrame atomically via a staging directory.

Writes df to staging_directory first, then moves the result into final_destination. If any step raises, the staging directory is cleaned up and final_destination remains untouched — no partial/corrupted files ever appear at the target.

Parameters:

  • df (pyspark.sql.DataFrame) – DataFrame to write.

  • final_destination (str) – Target path where the completed output should live. Must be empty or absent (unless mode='overwrite').

  • staging_directory (str) – Working directory where the write happens before the atomic move. Recommended: build via create_unique_staging_directory() to avoid concurrent collisions. Removed on success and on failure.

  • file_format (str) – Output format. Currently supports 'csv' and 'parquet'; defaults to 'csv'.

  • delimiter (str) – Field delimiter for CSV output. Ignored for parquet. Defaults to ','.

  • header (bool) – Include header row in CSV output. Ignored for parquet. Defaults to True.

  • mode (str) – Spark write mode passed to df.write.mode(...). Defaults to 'overwrite'.

Returns:

Side effect only — writes to final_destination.

Return type:

None

Raises:

RuntimeError – Wrapped from any exception during the staging write or the atomic move. Original cause is preserved via raise ... from     e.

siege_utilities.distributed.spark_utils.create_unique_staging_directory(base_path, operation_name='operation')[source]

Create a unique staging directory for atomic write operations.

Parameters:

  • base_path (str or pathlib.Path) – Parent directory under which the staging directory is created. Must be writable.

  • operation_name (str) – Human-readable label included in the directory name so concurrent operations remain distinguishable in logs. Defaults to 'operation'.

Returns:

Absolute path to the newly created staging directory. The name has the shape {base_path}/staging_{operation_name}_{uuid4hex} and is guaranteed unique.

Return type:

str

Notes

Used by atomic_write_with_staging() to isolate in-progress writes from the final destination so partial or failed writes never leave a corrupt file at the target path.

Functions by Category

Mathematical Functions

Array Functions

Aggregation Functions

Cryptographic Functions

Usage Examples

Basic mathematical operations:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
import siege_utilities

spark = SparkSession.builder.appName("MathExample").getOrCreate()

# Create sample data
data = [("A", 1.5), ("B", -2.3), ("C", 0.0)]
df = spark.createDataFrame(data, ["id", "value"])

# Apply mathematical functions
df = df.withColumn("abs_value", siege_utilities.abs(col("value")))
df = df.withColumn("acos_value", siege_utilities.acos(col("value")))

df.show()

Array operations:

# Array manipulation
df = df.withColumn("array_col", siege_utilities.array(col("id"), col("value")))
df = df.withColumn("distinct_array", siege_utilities.array_distinct(col("array_col")))
df = df.withColumn("array_contains", siege_utilities.array_contains(col("array_col"), "A"))

# Array aggregation
df = df.groupBy("id").agg(
    siege_utilities.array_agg(col("value")).alias("all_values")
)

Date operations:

from pyspark.sql.functions import current_date

# Add months to current date
df = df.withColumn("future_date",
                   siege_utilities.add_months(current_date(), 3))

Unit Tests

The Spark utilities module has comprehensive test coverage:

✅ test_spark_utils.py - All Spark utility tests pass

Test Coverage:
- Mathematical functions (abs, acos, acosh)
- Array operations (creation, manipulation, aggregation)
- Date functions (add_months)
- Aggregation functions (aggregate, any_value)
- Cryptographic functions (AES encryption/decryption)
- Edge cases and error handling

Test Results: All Spark utility tests pass successfully with comprehensive coverage.