Apache Spark / Typová kontrola s knihovnou Wick

Ať už preferujete psaní dotazů ve Spark SQL, nebo jste zvyklí spíše na volání funkcí přes Dataframe API, dříve nebo později narazíte na klasický problém. Pokud se přepíšete v názvu sloupce nebo se omylem pokusíte sečíst číslo s textem, kód se bez potíží zkompiluje. O chybě se tak dozvíte až za běhu programu. V praxi to často znamená zdlouhavé čekání na zabalení aplikace a alokaci zdrojů na clusteru jen proto, abyste zjistili, že job po pár vteřinách selhal kvůli banálnímu překlepu.

Tuto situaci elegantně řeší knihovna Wick, za kterou stojí inženýři z Netflixu. Jde o typově bezpečné API pro Spark, které přesouvá odhalování chyb z runtimu do fáze kompilace, čímž šetří hodiny zbytečné práce a frustrace.

S typovou bezpečností jde ruku v ruce i lepší podpora vývojového prostředí – IDE díky tomu samo napovídá existující sloupce a jejich datové typy. Pokud znáte starší knihovnu Frameless (určenou pro Scalu 2), Wick funguje jako její moderní nástupce. Využívá totiž novou funkcionalitu Named Tuples dostupnou od Scaly 3.7. Na rozdíl od standardního Dataset API navíc Wick funguje jako „zero-cost“ abstrakce. Znamená to, že veškerá typová kontrola probíhá výhradně při kompilaci a nepřináší s sebou žádnou výkonnostní penalizaci za běhu. Je ovšem důležité zmínit jedno aktuální omezení: knihovna zatím nepodporuje Spark 4.

Základem pro práci s knihovnou Wick je vytvoření obyčejné case classy, která reprezentuje strukturu dat. Následně se klasický netypový DataFrame ze Sparku obalí do struktury DataSeq, čímž se do hry vnese kýžená typová bezpečnost.

case class Livestock(
  entity: String,
  code: String,
  year: Int,
  asses: Double,
  buffalo: Double,
  cattle: Double,
  goats: Double,
  horses: Double,
  mules: Double,
  pigs: Double,
  sheep: Double,
  turkeys: Double,
  chickens: Double
)

val df = sparkSession.read
  .option("header", value = true)
  .schema(summon[Encoder[Livestock]].schema)
  .csv("data/livestock.csv")

val livestock = DataSeq[Livestock](df)

Jakmile jsou data obalená v DataSeq, můžeme s nimi pracovat za pomoci funkcí, které se velmi podobají standardnímu Dataframe API. Při přístupu k hodnotám (např. row.entity) už IDE samo napovídá, jaké vlastnosti jsou k dispozici. Operátor sčítání (+) je definován pouze pro číselné typy, což zaručuje, že výpočet celkového počtu dobytka v milionech neselže na nekompatibilních datech.

livestock
  .filter(_.year === 2014)
  .select(row =>
    (
      entity = row.entity,
      total_millions = (row.cattle + row.sheep + row.chickens) / 1_000_000
    )
  )
  .orderBy(row => desc(row.total_millions))
  .show()

Pro srovnání se podívejme, jak by vypadal tentýž dotaz zapsaný v čistém netypovém Spark SQL. Případný překlep ve slově „cattle“ by zde prošel bez povšimnutí a chyba by vyletěla až za běhu.

sparkSession
  .sql(
    """SELECT entity
      |     , (cattle + sheep + chickens) / 1000000 AS total_millions
      |  FROM livestock
      |  WHERE year = 2014
      |  ORDER BY total_millions DESC
      |""".stripMargin
  )
  .show()

Typová bezpečnost se pochopitelně vztahuje i na komplexnější operace, jakými jsou agregace. Metoda agg vyžaduje, aby její vstupy tvořily skalární výrazy. Pokud byste v následujícím kódu vynechali agregační funkci sum(...) a pokusili se přímo vydělit hodnotu row.cattle, kompilátor okamžitě ohlásí chybu. Nedovolí vám totiž agregovat neagregovanou hodnotu.

livestock
  .groupBy(row => (entity = row.entity))
  .agg(row => (cattle_total_millions = sum(row.cattle) / 1_000_000))
  .orderBy(row => desc(row.cattle_total_millions))
  .show()

A opět SQL alternativa téhož případu. Výsledek je sice identický, ale jakákoliv garance správnosti názvů a typů před spuštěním zde chybí.

sparkSession
  .sql(
    """SELECT entity
      |     , SUM(cattle) / 1000000 AS cattle_total_millions
      |  FROM livestock
      |  GROUP BY entity
      |  ORDER BY cattle_total_millions DESC
      |""".stripMargin
  )
  .show()

Ačkoliv typová kontrola ve Sparku existovala i dříve, často znamenala ústupky v pohodlí nebo výkonu. Wick naproti tomu naplno těží z Named Tuples ve Scale 3 a přináší moderní, zero-cost abstrakci. Výsledkem je nejen čistší kód, ale hlavně mnohem rychlejší a spolehlivější vývojový cyklus.

Celý projekt s příklady je k dispozici na githubu.