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.

The original df is unchanged.

Parameters:

  • df (None)

  • column_name (str)

Return type:

None

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]

Removes brackets from bounding box strings and reorders coordinates for Sedona.

Assumes input is a comma separated list in the order:

min latitude, max latitude, min longitude, max longitude

Produces an array in the order: [min_lon, min_lat, max_lon, max_lat]

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]

Helper function to convert a Spark DataFrame into a Pandas DataFrame, format it using tabulate, and print the result with a title.

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 df to DEBUG_SUBDIRECTORY/{step_name}_full_persisted.

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 the configured debug directory in the project’s standard output format so it can be inspected after the run.

No return value — purely a side-effect (write + log line).

Parameters:

step_name (str)

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]

Performs atomic write operations using a staging directory to prevent partial/corrupted files.

Parameters:

  • df (None)

  • final_destination (str)

  • staging_directory (str)

  • file_format (str)

  • delimiter (str)

  • header (bool)

  • mode (str)

Return type:

None

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

Creates a unique staging directory for atomic operations.

Parameters:

operation_name (str)

Return type:

str

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.