Run the AEF-BNG pipeline distributed across a Spark cluster.
Creates a DataFrame of chunk specs with flat columns, serialises the
AEF index, and processes partitions in parallel using mapInArrow.
Writes results to a Unity Catalog Delta table.
Parameters:
Source code in aef_bng/spark.py
| def process_with_spark(config: AEFBNGConfig) -> None:
"""Run the AEF-BNG pipeline distributed across a Spark cluster.
Creates a DataFrame of chunk specs with flat columns, serialises the
AEF index, and processes partitions in parallel using ``mapInArrow``.
Writes results to a Unity Catalog Delta table.
Args:
config: Pipeline configuration.
"""
import time
from pyspark.sql import SparkSession
t_total = time.perf_counter()
spark = SparkSession.builder.getOrCreate()
# Phase 1: Load tile index
t0 = time.perf_counter()
index = AEFBNGIndex()
index.load_for_bounds(config.bounds, config.years)
index_bytes = pickle.dumps(index)
logger.info("Index loaded in %.1fs", time.perf_counter() - t0)
# Phase 2: Build chunk grid and partition
t0 = time.perf_counter()
chunks_df = _build_chunks_dataframe(spark, config)
# Count actual rows (boundary filtering may have dropped chunks) - tiny DataFrame.
num_tasks = chunks_df.count()
# Each chunk is ~64MB in memory (64 bands x 1000x1000 int8) + reprojection overhead.
# Target 2-4 chunks per partition to balance parallelism vs memory on DS4_v2 (28GB, 8 cores).
# With autoscaling (1-10 workers x 8 cores = 8-80 slots), use ~num_tasks/3 partitions
# so each partition processes ~3 chunks sequentially within a single event loop.
num_partitions = max(1, min(num_tasks, num_tasks // 3))
chunks_df = chunks_df.repartition(num_partitions)
logger.info(
"Grid built: %d tasks across %d partitions in %.1fs",
num_tasks,
num_partitions,
time.perf_counter() - t0,
)
# Phase 3: Distributed processing (mapInArrow)
t0 = time.perf_counter()
spark_schema = _spark_output_schema()
process_fn = _make_process_partition(index_bytes, config.resampling)
result_df = chunks_df.mapInArrow(process_fn, schema=spark_schema)
from pyspark.databricks.sql import functions as dbf # type: ignore[import-not-found]
from pyspark.sql import functions as F # type: ignore # noqa: PGH003
result_df = result_df.withColumn(
"geometry", dbf.st_geomfromwkb(F.col("geometry_wkb"), F.lit(27700))
).drop("geometry_wkb")
# Phase 4: Write to Delta (triggers execution)
table_name = config.table_name
if table_name:
(
result_df.write.format("delta")
.mode("append")
.option("mergeSchema", "true")
.saveAsTable(table_name)
)
write_elapsed = time.perf_counter() - t0
logger.info("Processing + write complete in %.1fs -> %s", write_elapsed, table_name)
else:
(
result_df.write.format("delta")
.mode("append")
.option("mergeSchema", "true")
.save(config.output_path)
)
write_elapsed = time.perf_counter() - t0
logger.info("Processing + write complete in %.1fs -> %s", write_elapsed, config.output_path)
logger.info("Spark pipeline total: %.1fs", time.perf_counter() - t_total)
|