Distributed

Distributed computing: Spark utilities, HDFS operations and configuration.

Distributed functions package — lazy-loaded.

Contains Spark utilities, HDFS configuration, and HDFS operations. All submodules load on first attribute access via PEP 562 __getattr__.

siege_utilities.distributed.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.

siege_utilities.distributed.py_round(number, ndigits=None)

Round a number to a given precision in decimal digits.

The return value is an integer if ndigits is omitted or None. Otherwise the return value has the same type as the number. ndigits may be negative.

class siege_utilities.distributed.HDFSConfig[source]

Bases: object

Configuration class for HDFS and Spark operations.

Supports local, standalone cluster, and YARN deployments.

master

Spark master URL (‘local[*]’, ‘yarn’, ‘spark://host:7077’)

Type:

str

deploy_mode

Deployment mode for YARN (‘client’ or ‘cluster’)

Type:

str

yarn_queue

YARN queue name for resource allocation

Type:

str | None

driver_memory

Memory allocated to driver process

Type:

str

data_path: str | None = None
hdfs_base_directory: str = '/data/'
cache_directory: str | None = None
app_name: str = 'SparkDistributedProcessing'
spark_log_level: str = 'WARN'
master: str = 'local[*]'
deploy_mode: str = 'client'
yarn_queue: str | None = None
enable_sedona: bool = True
sedona_global_index_type: str = 'rtree'
sedona_join_broadcast_threshold: int = 10485760
num_executors: int | None = None
executor_cores: int | None = None
executor_memory: str = '2g'
driver_memory: str = '2g'
driver_cores: int = 1
network_timeout: str = '800s'
heartbeat_interval: str = '60s'
hdfs_timeout: int = 10
hdfs_copy_timeout: int = 300
force_sync: bool = False
log_info_func: Callable[[str], None] | None = None
log_error_func: Callable[[str], None] | None = None
hash_func: Callable[[str], str] | None = None
quick_signature_func: Callable[[str], str] | None = None
get_cache_path(filename)[source]

Get path for cache files.

Parameters:

filename (str)

Return type:

Path

get_optimal_partitions()[source]

Calculate optimal partitions for I/O heavy workloads

Return type:

int

property is_yarn: bool

Check if this config targets a YARN cluster.

property is_local: bool

Check if this config targets local mode.

log_info(message)[source]

Log info message using configured function

Parameters:

message (str)

log_error(message)[source]

Log error message using configured function

Parameters:

message (str)

__init__(data_path=None, hdfs_base_directory='/data/', cache_directory=None, app_name='SparkDistributedProcessing', spark_log_level='WARN', master='local[*]', deploy_mode='client', yarn_queue=None, enable_sedona=True, sedona_global_index_type='rtree', sedona_join_broadcast_threshold=10485760, num_executors=None, executor_cores=None, executor_memory='2g', driver_memory='2g', driver_cores=1, network_timeout='800s', heartbeat_interval='60s', hdfs_timeout=10, hdfs_copy_timeout=300, force_sync=False, log_info_func=None, log_error_func=None, hash_func=None, quick_signature_func=None)
Parameters:
  • data_path (str | None)

  • hdfs_base_directory (str)

  • cache_directory (str | None)

  • app_name (str)

  • spark_log_level (str)

  • master (str)

  • deploy_mode (str)

  • yarn_queue (str | None)

  • enable_sedona (bool)

  • sedona_global_index_type (str)

  • sedona_join_broadcast_threshold (int)

  • num_executors (int | None)

  • executor_cores (int | None)

  • executor_memory (str)

  • driver_memory (str)

  • driver_cores (int)

  • network_timeout (str)

  • heartbeat_interval (str)

  • hdfs_timeout (int)

  • hdfs_copy_timeout (int)

  • force_sync (bool)

  • log_info_func (Callable[[str], None] | None)

  • log_error_func (Callable[[str], None] | None)

  • hash_func (Callable[[str], str] | None)

  • quick_signature_func (Callable[[str], str] | None)

Return type:

None

siege_utilities.distributed.create_hdfs_config(**kwargs)[source]

Factory function to create HDFS configuration.

Parameters:

**kwargs – Forwarded to HDFSConfig(). See HDFSConfig for accepted fields (e.g. data_path, master, executor_memory, enable_sedona, etc.).

Return type:

HDFSConfig

siege_utilities.distributed.create_local_config(data_path, **kwargs)[source]

Create config optimized for local development.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of local-mode defaults (num_executors=2, executor_cores=1, executor_memory=’1g’, enable_sedona=False). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.create_cluster_config(data_path, **kwargs)[source]

Create config optimized for cluster deployment.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of cluster defaults (num_executors=8, executor_cores=4, executor_memory=’4g’, enable_sedona=True). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.create_geocoding_config(data_path, **kwargs)[source]

Create config optimized for geocoding workloads.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of geocoding defaults (app_name=’GeocodingPipeline’, num_executors=4, executor_memory=’2g’, enable_sedona=True, network_timeout=’1200s’). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.create_yarn_config(data_path, **kwargs)[source]

Create config optimized for YARN cluster deployment.

This configuration is designed for production YARN clusters with: - YARN as the resource manager - Client deploy mode (driver runs locally) - Higher resource allocation for executors - Sedona enabled for spatial operations

Parameters:
  • data_path (str) – Path to data (local or HDFS)

  • **kwargs – Override any default settings

Returns:

HDFSConfig configured for YARN deployment

Return type:

HDFSConfig

Example

>>> config = create_yarn_config('/data/census', yarn_queue='analytics')
>>> spark, path, _ = setup_distributed_environment(config)
siege_utilities.distributed.create_census_analysis_config(data_path, **kwargs)[source]

Create config optimized for Census longitudinal analysis.

Designed for processing large Census datasets with: - Sedona for spatial operations (boundary crosswalks) - Optimized partitioning for Census tract-level data - Memory settings for time-series joins

Parameters:
  • data_path (str) – Path to Census data

  • **kwargs – Override any default settings

Returns:

HDFSConfig optimized for Census analysis

Return type:

HDFSConfig

Example

>>> config = create_census_analysis_config('/data/census/acs')
>>> spark, path, _ = setup_distributed_environment(config)
class siege_utilities.distributed.AbstractHDFSOperations[source]

Bases: object

Configurable HDFS operations for distributed pipelines.

Wraps the small set of HDFS interactions the library needs — status check, Spark session construction, distributed dependency staging — behind a single class parameterized by project-specific settings (cache locations, Hadoop timeout, hash functions).

Parameters:

config (HDFSConfig) – Project-specific configuration object. Must expose a cache_directory path, get_cache_path method, hdfs_timeout, and optional hash_func / quick_signature_func overrides. See create_hdfs_operations() and setup_distributed_environment() for common construction patterns.

Notes

Instantiation eagerly creates config.cache_directory if absent. All HDFS interactions log through config.log_info / config.log_error rather than raising, matching the pre-check convention used elsewhere in the distributed module.

__init__(config)[source]

Initialize with HDFSConfig

check_hdfs_status()[source]

Check if HDFS is accessible

Return type:

bool

create_spark_session()[source]

Create Spark session using configuration.

Supports local, standalone cluster, and YARN deployments based on the master URL in the config.

sync_directory_to_hdfs(local_path=None, hdfs_subdir='inputs')[source]

Sync local directory/file to HDFS with proper verification.

Raises:
Parameters:
  • local_path (str | None)

  • hdfs_subdir (str)

Return type:

Tuple[str, Dict]

setup_distributed_environment(data_path=None, dependency_paths=None)[source]

Main setup function with proper verification.

Returns:

Tuple of (spark_session, data_url, None).

Raises:
Parameters:
  • data_path (str | None)

  • dependency_paths (List[str] | None)

siege_utilities.distributed.setup_distributed_environment(config, data_path=None, dependency_paths=None)[source]

One-shot helper to construct HDFS ops + stage data + build a Spark session.

Equivalent to AbstractHDFSOperations(config).setup_distributed_environment(...) — provided as a top-level convenience so callers who only need a single distributed-setup call don’t have to instantiate the class explicitly.

Parameters:
  • config (HDFSConfig) – Project-specific HDFS config (see AbstractHDFSOperations).

  • data_path (str, optional) – Local path to a data payload that should be staged onto HDFS before the Spark session starts. When None, no staging happens.

  • dependency_paths (list of str, optional) – Additional local paths (e.g. Python source directories) to zip and stage as PySpark dependencies. When None, only the library’s own default dependency bundle is staged.

Returns:

(spark_session, staged_data_url, dependency_zip_url) — matching the tuple returned by AbstractHDFSOperations.setup_distributed_environment(). Both URLs are None when their corresponding argument was omitted.

Return type:

tuple

siege_utilities.distributed.create_hdfs_operations(config)[source]

Factory for AbstractHDFSOperations instances.

Wraps the constructor so callers can request an HDFS operations object without importing the class name directly — useful for keeping consumer code decoupled from the concrete implementation.

Parameters:

config (HDFSConfig) – Project-specific HDFS config (see AbstractHDFSOperations).

Returns:

A ready-to-use HDFS operations instance.

Return type:

AbstractHDFSOperations

Submodules

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.

Abstract HDFS Operations - Fully Configurable and Reusable Zero hard-coded project dependencies

class siege_utilities.distributed.hdfs_operations.AbstractHDFSOperations[source]

Bases: object

Configurable HDFS operations for distributed pipelines.

Wraps the small set of HDFS interactions the library needs — status check, Spark session construction, distributed dependency staging — behind a single class parameterized by project-specific settings (cache locations, Hadoop timeout, hash functions).

Parameters:

config (HDFSConfig) – Project-specific configuration object. Must expose a cache_directory path, get_cache_path method, hdfs_timeout, and optional hash_func / quick_signature_func overrides. See create_hdfs_operations() and setup_distributed_environment() for common construction patterns.

Notes

Instantiation eagerly creates config.cache_directory if absent. All HDFS interactions log through config.log_info / config.log_error rather than raising, matching the pre-check convention used elsewhere in the distributed module.

__init__(config)[source]

Initialize with HDFSConfig

check_hdfs_status()[source]

Check if HDFS is accessible

Return type:

bool

create_spark_session()[source]

Create Spark session using configuration.

Supports local, standalone cluster, and YARN deployments based on the master URL in the config.

sync_directory_to_hdfs(local_path=None, hdfs_subdir='inputs')[source]

Sync local directory/file to HDFS with proper verification.

Raises:
Parameters:
  • local_path (str | None)

  • hdfs_subdir (str)

Return type:

Tuple[str, Dict]

setup_distributed_environment(data_path=None, dependency_paths=None)[source]

Main setup function with proper verification.

Returns:

Tuple of (spark_session, data_url, None).

Raises:
Parameters:
  • data_path (str | None)

  • dependency_paths (List[str] | None)

siege_utilities.distributed.hdfs_operations.create_hdfs_operations(config)[source]

Factory for AbstractHDFSOperations instances.

Wraps the constructor so callers can request an HDFS operations object without importing the class name directly — useful for keeping consumer code decoupled from the concrete implementation.

Parameters:

config (HDFSConfig) – Project-specific HDFS config (see AbstractHDFSOperations).

Returns:

A ready-to-use HDFS operations instance.

Return type:

AbstractHDFSOperations

siege_utilities.distributed.hdfs_operations.setup_distributed_environment(config, data_path=None, dependency_paths=None)[source]

One-shot helper to construct HDFS ops + stage data + build a Spark session.

Equivalent to AbstractHDFSOperations(config).setup_distributed_environment(...) — provided as a top-level convenience so callers who only need a single distributed-setup call don’t have to instantiate the class explicitly.

Parameters:
  • config (HDFSConfig) – Project-specific HDFS config (see AbstractHDFSOperations).

  • data_path (str, optional) – Local path to a data payload that should be staged onto HDFS before the Spark session starts. When None, no staging happens.

  • dependency_paths (list of str, optional) – Additional local paths (e.g. Python source directories) to zip and stage as PySpark dependencies. When None, only the library’s own default dependency bundle is staged.

Returns:

(spark_session, staged_data_url, dependency_zip_url) — matching the tuple returned by AbstractHDFSOperations.setup_distributed_environment(). Both URLs are None when their corresponding argument was omitted.

Return type:

tuple

Abstract HDFS Configuration Configurable settings for HDFS operations - no hard-coded project dependencies

class siege_utilities.distributed.hdfs_config.HDFSConfig[source]

Bases: object

Configuration class for HDFS and Spark operations.

Supports local, standalone cluster, and YARN deployments.

master

Spark master URL (‘local[*]’, ‘yarn’, ‘spark://host:7077’)

Type:

str

deploy_mode

Deployment mode for YARN (‘client’ or ‘cluster’)

Type:

str

yarn_queue

YARN queue name for resource allocation

Type:

str | None

driver_memory

Memory allocated to driver process

Type:

str

data_path: str | None = None
hdfs_base_directory: str = '/data/'
cache_directory: str | None = None
app_name: str = 'SparkDistributedProcessing'
spark_log_level: str = 'WARN'
master: str = 'local[*]'
deploy_mode: str = 'client'
yarn_queue: str | None = None
enable_sedona: bool = True
sedona_global_index_type: str = 'rtree'
sedona_join_broadcast_threshold: int = 10485760
num_executors: int | None = None
executor_cores: int | None = None
executor_memory: str = '2g'
driver_memory: str = '2g'
driver_cores: int = 1
network_timeout: str = '800s'
heartbeat_interval: str = '60s'
hdfs_timeout: int = 10
hdfs_copy_timeout: int = 300
force_sync: bool = False
log_info_func: Callable[[str], None] | None = None
log_error_func: Callable[[str], None] | None = None
hash_func: Callable[[str], str] | None = None
quick_signature_func: Callable[[str], str] | None = None
get_cache_path(filename)[source]

Get path for cache files.

Parameters:

filename (str)

Return type:

Path

get_optimal_partitions()[source]

Calculate optimal partitions for I/O heavy workloads

Return type:

int

property is_yarn: bool

Check if this config targets a YARN cluster.

property is_local: bool

Check if this config targets local mode.

log_info(message)[source]

Log info message using configured function

Parameters:

message (str)

log_error(message)[source]

Log error message using configured function

Parameters:

message (str)

__init__(data_path=None, hdfs_base_directory='/data/', cache_directory=None, app_name='SparkDistributedProcessing', spark_log_level='WARN', master='local[*]', deploy_mode='client', yarn_queue=None, enable_sedona=True, sedona_global_index_type='rtree', sedona_join_broadcast_threshold=10485760, num_executors=None, executor_cores=None, executor_memory='2g', driver_memory='2g', driver_cores=1, network_timeout='800s', heartbeat_interval='60s', hdfs_timeout=10, hdfs_copy_timeout=300, force_sync=False, log_info_func=None, log_error_func=None, hash_func=None, quick_signature_func=None)
Parameters:
  • data_path (str | None)

  • hdfs_base_directory (str)

  • cache_directory (str | None)

  • app_name (str)

  • spark_log_level (str)

  • master (str)

  • deploy_mode (str)

  • yarn_queue (str | None)

  • enable_sedona (bool)

  • sedona_global_index_type (str)

  • sedona_join_broadcast_threshold (int)

  • num_executors (int | None)

  • executor_cores (int | None)

  • executor_memory (str)

  • driver_memory (str)

  • driver_cores (int)

  • network_timeout (str)

  • heartbeat_interval (str)

  • hdfs_timeout (int)

  • hdfs_copy_timeout (int)

  • force_sync (bool)

  • log_info_func (Callable[[str], None] | None)

  • log_error_func (Callable[[str], None] | None)

  • hash_func (Callable[[str], str] | None)

  • quick_signature_func (Callable[[str], str] | None)

Return type:

None

siege_utilities.distributed.hdfs_config.create_census_analysis_config(data_path, **kwargs)[source]

Create config optimized for Census longitudinal analysis.

Designed for processing large Census datasets with: - Sedona for spatial operations (boundary crosswalks) - Optimized partitioning for Census tract-level data - Memory settings for time-series joins

Parameters:
  • data_path (str) – Path to Census data

  • **kwargs – Override any default settings

Returns:

HDFSConfig optimized for Census analysis

Return type:

HDFSConfig

Example

>>> config = create_census_analysis_config('/data/census/acs')
>>> spark, path, _ = setup_distributed_environment(config)
siege_utilities.distributed.hdfs_config.create_cluster_config(data_path, **kwargs)[source]

Create config optimized for cluster deployment.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of cluster defaults (num_executors=8, executor_cores=4, executor_memory=’4g’, enable_sedona=True). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.hdfs_config.create_geocoding_config(data_path, **kwargs)[source]

Create config optimized for geocoding workloads.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of geocoding defaults (app_name=’GeocodingPipeline’, num_executors=4, executor_memory=’2g’, enable_sedona=True, network_timeout=’1200s’). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.hdfs_config.create_hdfs_config(**kwargs)[source]

Factory function to create HDFS configuration.

Parameters:

**kwargs – Forwarded to HDFSConfig(). See HDFSConfig for accepted fields (e.g. data_path, master, executor_memory, enable_sedona, etc.).

Return type:

HDFSConfig

siege_utilities.distributed.hdfs_config.create_local_config(data_path, **kwargs)[source]

Create config optimized for local development.

Parameters:
  • data_path (str) – Path to data directory.

  • **kwargs – Override any HDFSConfig default. Merged on top of local-mode defaults (num_executors=2, executor_cores=1, executor_memory=’1g’, enable_sedona=False). See HDFSConfig for accepted fields.

Return type:

HDFSConfig

siege_utilities.distributed.hdfs_config.create_yarn_config(data_path, **kwargs)[source]

Create config optimized for YARN cluster deployment.

This configuration is designed for production YARN clusters with: - YARN as the resource manager - Client deploy mode (driver runs locally) - Higher resource allocation for executors - Sedona enabled for spatial operations

Parameters:
  • data_path (str) – Path to data (local or HDFS)

  • **kwargs – Override any default settings

Returns:

HDFSConfig configured for YARN deployment

Return type:

HDFSConfig

Example

>>> config = create_yarn_config('/data/census', yarn_queue='analytics')
>>> spark, path, _ = setup_distributed_environment(config)