21 · Gestire partizioni, cache e memoria

Misura come si distribuisce il lavoro, prova una tecnica contro lo sbilanciamento e osserva la cache. Usa una nuova sessione CPU su Kaggle con PySpark classico 4.0.x e Java 17+. Se serve, installa PySpark con %pip install "pyspark>=4.0,<4.1". Tutti i dati sono generati localmente.

from pathlib import Path
from tempfile import TemporaryDirectory

from pyspark import StorageLevel
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.master("local[2]").appName("21-partitions-cache-memory").getOrCreate()
spark.conf.set("spark.sql.shuffle.partitions", "4")
spark.conf.set("spark.sql.adaptive.enabled", "false")
sales = spark.range(1000, numPartitions=4).select(
    "id",
    F.when(F.col("id") < 900, 0).otherwise(F.col("id") % 7 + 1).alias("shop_id"),
    (F.col("id") % 10 + 1).alias("units"),
)

1. Contare le partizioni

In PySpark classico, df.rdd.getNumPartitions() indica quante partizioni ha la rappresentazione RDD del DataFrame. Il conteggio da solo non indica come si distribuiscono le righe.

2. Cache

Metti in cache un risultato quando più azioni lo riuseranno e il costo di memorizzazione è giustificato. Usa quel DataFrame per le azioni successive e libera la cache quando hai finito.

3. Sbilanciamento

Il lavoro è sbilanciato quando si distribuisce in modo disomogeneo. Dopo aver ripartito per negozio, confronta la partizione più grande con le altre. Il numero di righe è una stima del lavoro. I costi effettivi dei task possono variare.

4. Cambiare il numero di partizioni

repartition() richiede uno shuffle e può aumentare o ridurre il numero di partizioni. coalesce() può ridurlo senza introdurre uno shuffle, ma può lasciare il lavoro sbilanciato.

5. Chiavi che raccolgono molte righe

Il partizionamento per hash mantiene insieme le chiavi uguali. Il negozio 0 ha 900 righe, quindi almeno 900 finiscono nella stessa partizione anche aumentando il numero totale di partizioni.

print("Input partitions:", sales.rdd.getNumPartitions())
cached_sales = sales.cache()
try:
    assert cached_sales.count() == 1000
    assert cached_sales.filter("shop_id = 0").count() == 900
finally:
    cached_sales.unpersist()

keyed = sales.repartition(4, "shop_id")
assert keyed.rdd.getNumPartitions() == 4
counts = keyed.select(F.spark_partition_id().alias("partition")).groupBy("partition").count()
assert max(row["count"] for row in counts.collect()) >= 900
counts.orderBy("partition").show()
assert sales.coalesce(2).rdd.getNumPartitions() == 2

6. Dividere le partizioni sbilanciate con AQE

L’esecuzione adattiva può suddividere partizioni di shuffle sbilanciate per i join idonei e replicare il lavoro corrispondente quando necessario. Usa dimensioni osservate e soglie. Controlla il piano eseguito e le metriche dei task. Qui lo sbilanciamento si vede dai conteggi, anche senza una suddivisione da parte di AQE.

7. Aggiungere sale a una chiave frequente

Assegna le righe della chiave frequente a più gruppi intermedi. Qui un hash stabile dell’identificativo rende riproducibile l’assegnazione. Il sale suddivide una chiave logica.

8. Aggregazione parziale

Somma in ogni gruppo con sale, poi somma i risultati parziali per chiave originale. Spark può già calcolare somme parziali per le somme native. Qui il sale dimostra la correttezza del risultato senza misurare un guadagno di prestazioni.

9. Aggregazione senza un risultato parziale combinabile

Una funzione che richiede tutti i valori di un gruppo può doverli conservare. Per esempio, in generale le mediane dei gruppi con sale non bastano a ricavare la mediana globale. Usa un percentile integrato esatto o approssimato quando è adatto.

salt = F.when(F.col("shop_id") == 0, F.pmod(F.xxhash64("id"), F.lit(4))).otherwise(0)
partial = sales.withColumn("salt", salt).groupBy("shop_id", "salt").agg(F.sum("units").alias("units"))
combined = partial.groupBy("shop_id").agg(F.sum("units").alias("units"))
direct = sales.groupBy("shop_id").agg(F.sum("units").alias("units"))
assert {tuple(row) for row in combined.collect()} == {tuple(row) for row in direct.collect()}
assert partial.filter("shop_id = 0").count() == 4
partial.filter("shop_id = 0").orderBy("salt").show()

10. Riversamento su disco

Gli operatori possono scrivere dati intermedi su disco quando la memoria scarseggia. Controlla le metriche di spill in memoria e su disco per task. Un task lento con spill elevato può avere una partizione di input grande. L’esempio evita di forzare la pressione sulla memoria.

11. Dimensione delle partizioni in lettura

spark.sql.files.maxPartitionBytes influisce su come Spark raggruppa l’input suddivisibile in partizioni di scansione. Contano anche formato, dimensioni e costo di apertura dei file.

12. Partizioni dopo lo shuffle

spark.sql.shuffle.partitions imposta il numero predefinito per molti shuffle SQL. AQE può cambiare come viene raggruppato il lavoro successivo. Qui è disabilitata per rendere riproducibile il conteggio.

13. coalesce(1) prima di un calcolo pesante

Ridurre a una partizione può far eseguire il lavoro successivo a un solo task, fino a un altro scambio. Svolgi prima il lavoro in parallelo se alla fine serve davvero un unico piccolo file.

with TemporaryDirectory(prefix="nb21-") as directory:
    path = str(Path(directory) / "sales")
    sales.write.parquet(path)
    file_scan = spark.read.parquet(path)
    print("Max scan partition bytes:", spark.conf.get("spark.sql.files.maxPartitionBytes"))
    print("Scan partitions:", file_scan.rdd.getNumPartitions())

print("Configured shuffle partitions:", spark.conf.get("spark.sql.shuffle.partitions"))
assert direct.rdd.getNumPartitions() == 4
sales.coalesce(1).select(F.sqrt("units")).explain()

14. Livello di memorizzazione

persist() permette di scegliere un livello di memorizzazione. Qui MEMORY_AND_DISK_DESER conserva i blocchi in memoria quando possibile e consente l’uso del disco.

15. Liberare la cache

unpersist() rimuove la richiesta di persistenza e libera i blocchi. Usa blocking=True per attendere la rimozione.

16. Controllare se un DataFrame è in cache

is_cached mostra l’indicatore di persistenza dell’oggetto DataFrame. Il numero di partizioni già calcolate va verificato separatamente.

17. Cache pigra

La persistenza contrassegna un piano per la cache. Un’azione la popola man mano che valuta le partizioni. Se accessibile, la vista Storage dell’interfaccia Spark mostra i blocchi materializzati.

18. Cache per un DataFrame usato una volta

La prima azione deve calcolare i dati e popolare la cache. Un risultato usato una sola volta spesso sostiene il costo di memorizzazione senza beneficiare di una seconda lettura più veloce.

reused = sales.filter("units >= 5").persist(StorageLevel.MEMORY_AND_DISK_DESER)
try:
    assert reused.is_cached
    print("Requested storage level:", reused.storageLevel)
    print("Spark UI:", spark.sparkContext.uiWebUrl)
    assert reused.count() == 600
    assert reused.agg(F.max("units")).first()[0] == 10
finally:
    reused.unpersist(blocking=True)
assert not reused.is_cached

19. Ridurre lo shuffle

Filtra o escludi colonne superflue prima di uno scambio se la semantica lo permette. Aggregare alla grana richiesta può ridurre le righe spostate.

20. Riuso del partizionamento

Spark può riusare una distribuzione compatibile con i requisiti dell’operatore a valle. Qui l’aggregazione può usare il partizionamento per hash esplicito sul negozio. Cerca gli scambi nel piano per verificare se due operazioni sulla stessa chiave condividono la distribuzione.

prepared = sales.filter("units >= 5").select("shop_id", "units").repartition(4, "shop_id")
prepared.groupBy("shop_id").agg(F.sum("units")).explain()

21. Memoria riservata

Nella configurazione standard Spark calcola la memoria unificata sottraendo una riserva di 300 MiB dall’heap JVM. L’overhead del container è una risorsa separata.

22. Memoria unificata

Esecuzione e memorizzazione condividono una regione dimensionata da spark.memory.fraction. L’esecuzione include shuffle, join e ordinamenti. La memorizzazione include i blocchi in cache.

23. Memoria protetta per la cache

spark.memory.storageFraction imposta la quota di memoria unificata sotto la quale l’esecuzione non può espellere blocchi. La quota rappresenta una soglia di protezione, non una preallocazione esclusiva.

24. Memoria del driver

Il driver richiede spazio per pianificazione, metadati e risultati restituiti. L’heap JVM e la memoria del processo Python sono separati. Configura l’heap del driver all’avvio.

25. Memoria esaurita nel driver

collect() o toPandas() senza limiti possono sovraccaricare il driver. Controlla i log e usa anteprime limitate o scritture distribuite.

26. Memoria esaurita negli executor

Una grande partizione di shuffle, un broadcast, un risultato in cache o un worker Python possono mettere sotto pressione le risorse. Individua il processo o container fallito, controlla sbilanciamento e metriche dei task e adegua la disposizione dei dati o le risorse. Gli esempi evitano di provocare un esaurimento di memoria.

conf = spark.sparkContext.getConf()
print("Memory fraction:", conf.get("spark.memory.fraction", "0.6"))
print("Storage fraction:", conf.get("spark.memory.storageFraction", "0.5"))
print("Bounded preview:", sales.orderBy("id").take(3))

Verifica la comprensione

1. Perché più partizioni possono lasciare sovraccarico un negozio?
2. Perché la somma delle somme parziali con sale recupera il totale?
3. Che cosa indica is_cached prima della prima azione?

Risposte

1. Il partizionamento per hash manda comunque chiavi uguali nella stessa partizione.
2. Ogni riga contribuisce a una sola somma parziale e la somma finale le combina tutte.
3. Indica la richiesta di persistenza. La cache deve ancora essere materializzata da un’azione.