S-JIS[2017-01-15/2020-10-08] 変更履歴

Dataset

Apache SparkのDatasetクラスについて。

概要

Spark2.0以降はDatasetを使ってプログラミングする。
（Spark2では、DataFrameは「Dataset[Row]」の別名である）

最初のDatasetはSparkSessionを使って生成する。
Scalaのコレクション（Seq）から作る方法と、ファイル等から読み込んで作る方法がある。

DatasetクラスにはfilterやmapといったScalaのコレクションと同様のメソッドが用意されている。

Datasetの生成

Seqから生成する例

import org.apache.spark.sql.SparkSession

case class Person(name: String, age: Long)

    val spark = SparkSession.builder().〜.getOrCreate()
    import spark.implicits._

    val seq = Seq(Person("hoge", 20), Person("foo", 30), Person("bar", 40))
    val ds = seq.toDS()

seq.toDF()でDataFrameに変換することも出来るが、内部的にはseq.toDS().toDF()と同じ。

テキストファイルから生成する例

import org.apache.spark.sql.SparkSession

  def readText(spark: SparkSession): Unit = {
    import spark.implicits._

    // DataFrameとして読み込む
    val df = spark.read.text("D:/tmp/person.csv")
//  df.printSchema()

    // Dataset[Person]に変換
    val ds = df.map(row => {
      val s = row.getString(0)
      val a = s.split(",")
      Person(a(0), a(1).toLong)
    })
    ds.show()
  }

read.textメソッドは可変長引数なので、複数ファイルを指定することも出来る。
また、ワイルドカード「*」を使うことも出来る。

read.textでファイルを読み込むと、スキーマは、valueという名前のStringフィールドがひとつあるだけの状態になっている。

root
 |-- value: string (nullable = true)

csvファイルから生成する例

import org.apache.spark.sql.SparkSession

  def readCsv(spark: SparkSession): Unit = {
    import spark.implicits._

    // DataFrameとして読み込む
    val df = spark.read.csv("D:/tmp/person.csv")
//  df.printSchema()

    // Dataset[Person]に変換
    val ds = df.map(row => {
      val name = row.getString(0)
      val age  = row.getString(1).toLong
      Person(name, age)
    })
    ds.show()
  }

csvにはいくつかのオプションがあり、例えばフィールドセパレーターをタブにすればtsvファイルが読める。[2017-01-17]
他にもエンコーディングやヘッダー有無の指定等、いくつかのオプションがある。
どのようなオプションがあるのかは、def csv(paths: String*)メソッドのScaladocを参照。

read.csvでファイルを読み込むと、スキーマは、_c0, _c1といったカラム名でデータ型は全てStringという状態になっている。

root
 |-- _c0: string (nullable = true)
 |-- _c1: string (nullable = true)

スキーマを指定してcsvファイルから読み込む例

DataFrameのスキーマ（カラム名とデータ型）がケースクラスと一致していれば、（自分でmapを書かなくても）そのケースクラスのDatasetに変換できる。

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types.LongType
import org.apache.spark.sql.types.StringType
import org.apache.spark.sql.types.StructType

  def readSchema(spark: SparkSession): Unit = {
    import spark.implicits._

    val schema = new StructType()
      .add("name", StringType)
      .add("age", LongType)
    val ds = spark.read.schema(schema).csv("D:/tmp/person.csv").as[Person]
    ds.show()
  }

import org.apache.spark.sql.types._

val schema = StructType(Seq(
  StructField("name", StringType),
  StructField("age", LongType)
))

import org.apache.spark.sql.types._

val schema = new StructType()
  .add("name", StringType)
  .add("age", LongType)

val schema = Seq.empty[Person].toDS().schema

val schema = spark.emptyDataset[Person].schema

import org.apache.spark.sql.Encoder

def getSchema[T](implicit enc: Encoder[T]) = enc.schema
val schema = getSchema[Person]

import org.apache.spark.sql.Encoder

val schema = implicitly[Encoder[Person]].schema

import org.apache.spark.sql.Encoders

val schema = Encoders.product[Person].schema

tsvファイルから生成する例

tsvファイルを読み込む場合は、フィールドセパレーターをタブ文字にしてcsv読み込みメソッドを使用する。[2017-01-17]

import org.apache.spark.sql.Encoders
import org.apache.spark.sql.SparkSession

  def readTsv(spark: SparkSession): Unit = {
    import spark.implicits._

    val schema = Encoders.product[Person].schema
    val ds = spark.read.schema(schema).option("sep", "\t").csv("D:/tmp/person.tsv").as[Person]
//  ds.printSchema()
    ds.show()
  }

Datasetのメソッド

Dataset[T]の主なメソッド。

設定系

変換系

Datasetから別のDataset（DataFrame）を生成する処理。

アクション

ジョブを実行して値を返す処理。（いわゆる終端処理）

列操作系

カラムを増減・演算する処理。

行操作系

レコードを操作する処理。

結合処理

他のDatasetを使用する処理。

結合の注意

joinで結合キーとしてカラム名を指定するメソッドの場合は、双方のDatasetに同じ名前のカラムが存在する必要があるが、出力されるデータでは1カラムのみになる。
しかし結合条件を指定する方式（ds("name")===right("name")）の場合、双方に同名のカラムがあると、両方とも残る。 そして後続処理でそのカラム名を指定すると、一意に定まらなくてエラーになる。

この場合、Datasetに名前（テーブル名相当）を付けておけば、その名前を使ってカラムを一意に識別することが出来る。

    import spark.implicits._

    val ds1 = Seq(Person("hoge", 20), Person("foo", 30), Person("bar", 40)).toDS().as("person")
    val ds2 = Seq(Master("hoge", "abc"), Master("foo", "def"), Master("zzz", "zzz")).toDS().as("master")

    val df = ds1.join(ds2)
      .where($"person.name" === $"master.name")
      .select("person.name", "age", "sub")
    df.show()

↓実行結果

+----+---+---+
|name|age|sub|
+----+---+---+
|hoge| 20|abc|
| foo| 30|def|
+----+---+---+

インナージョインやレフトジョインでマスター側のデータに同一キーのレコードが複数有る場合、全て結合対象となる。
（Asakusa FrameworkのMasterJoinの場合、マスター側に複数レコードあると、どれか1レコードのみが結合対象となる）

ブロードキャスト

データ量の少ないDataset（DataFrame）なら、ブロードキャスト結合（hash join）をすることが出来る。[2017-01-16]

import org.apache.spark.sql.functions.broadcast

    val df = ds1.join(broadcast(ds2))

永続化（キャッシュ）処理

→RDDの永続化処理

キー毎に重複排除する例

同一キーで複数レコードある場合に1レコードのみ残すこと（重複排除・distinct）を考えてみる。[2017-01-16]

一番単純なのは、dropDuplicatesを使う方法。

    import spark.implicits._

    val ds = Seq(Person("hoge", 30), Person("hoge", 20), Person("hoge", 40), Person("foo", 30)).toDS()
    val ds2 = ds.dropDuplicates("name")
    ds2.show()

↓実行結果

+----+---+
|name|age|
+----+---+
| foo| 30|
|hoge| 30|
+----+---+

しかしこれだと、重複しているレコードの内のどのレコードが残るか分からない。
（先頭レコードが残るような感じがするけど、分散処理しているときにどれが先頭になるかは保証されるか？）

dropDuplicatesする前に全体をソートする方法は（少量データでは望んだ結果になったけど）、データがパーティションをまたがった場合に上手くいかないと思われる。

×  val ds2 = ds.sort($"name", $"age".asc).dropDuplicates("name")

    val ds2 = ds.repartition($"name").sortWithinPartitions($"name", $"age".asc).dropDuplicates("name")

キー毎にグルーピングして値一覧を処理すれば、そこで自分が望むデータだけを残すことが出来る。

    val ds2 = ds.groupByKey(_.name).mapGroups((k, i) => i.toSeq.sortBy(_.age).head)

iはIterator。ageでソートしたいがIteratorにはsortByが無いので、一旦Seqに変換している。
そして、ageでソートして先頭データ（つまりageが最小のPerson）を取得している。

↓実行結果

+----+---+
|name|age|
+----+---+
| foo| 30|
|hoge| 20|
+----+---+

ただ、自前でIteratorをソートしているので、件数が多いときは嫌なことになりそうな気がする。

自前でソートするとデータ件数が多い場合に不安だし、そもそもソートはシャッフル処理に任せたい気がするので、キーでパーティションを分けてからソートしておけば大丈夫か？（groupByKeyでパーティションが変わる可能性があるような気もする…）

    val ds2 = ds.repartition($"name").sortWithinPartitions($"name", $"age".asc).groupByKey(_.name).mapGroups((k, i) => i.next)

並び順は気にせずに最小のものを探すという方法も考えられる。

    val ds2 = ds.groupByKey(_.name).mapGroups { (k, i) =>
      var min: Person = null
      i.foreach { p =>
        if ((min eq null) || p.age < min.age) {
          min = p
        }
      }
      min
    }

しかし、なんだかとてもクドい感じ(爆)

IteratorにsortByは無いがminBy（やmaxBy）はあるので、これを使うのが一番シンプルかも。[2017-01-23]

    val ds2 = ds.groupByKey(_.name).mapGroups((k, i) => i.minBy(_.age))

minByなら（toSeq.sortByと違ってソートせずに）データ全体を1回走査するだけなので、ソートを伴う操作より早いかも？

自分で最小のものを探すなら、reduceを使う方法もある。

    val ds2 = ds.groupByKey(_.name).reduceGroups((p, r) => if (p.age <= r.age) p else r).map(_._2)

（reduceGroupsはキーと値（この例ではPerson）のタプルを返すので、mapを使ってタプルの2番目（Person）だけ取得している）

reduce内でif式を使いたくないなら、以下のような方法も考えられる。

    // 新しいPersonを作って返す
    val ds2 = ds.groupByKey(_.name).reduceGroups((p, r) => Person(p.name, math.min(p.age, r.age))).map(_._2)

    // 一部だけを変更したPersonを作って返す
    val ds2 = ds.groupByKey(_.name).reduceGroups((p, r) => p.copy(age = math.min(p.age, r.age))).map(_._2)

ケースクラスが可変（case class Person(name: String, var age: Long)）であれば、該当フィールドだけ書き換えることも出来る。[2017-01-17]

    val ds2 = ds.groupByKey(_.name).reduceGroups { (p, r) => p.age = math.min(p.age, r.age); p }.map(_._2)

Window関数を使うという方法もあるようだ。（参考： stackoverflowのSPARK DataFrame: select the first row of each group）

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.row_number

    val w = Window.partitionBy($"name").orderBy($"age".asc)
    val ds2 = ds.withColumn("rn", row_number.over(w)).where($"rn" === 1).drop("rn").as[Person]

行番号を表すrow_numberカラムをwithColumnで追加し、それが1になるもの（つまり先頭行）だけをwhereで抽出している。
（row_numberカラムは最終的には不要なので、dropで削除している）
（withColumnを使うとDataFrameになってしまうので、（drop後に）as[Person]を呼び出してDataset[Person]にしている）

参考までに、Asakusa Framework（Java）のGroupSortだとグルーピングとソートが同時に指定できるので、以下のように書ける。

    @GroupSort
    public void distinctByName(@Key(group = {"name"}, order = {"age ASC"}) List<Person> personList, Result<Person> out) {
        out.add(personList.get(0));
    }

Asakusa Framewokで書いたアプリケーションをSparkに変換して動くAsakusa on Sparkがこれを実現できているんだから、何か方法はあると思うんだけど＾＾；

ファイル出力

Dataset#writeでファイルに出力することが出来る。[2017-01-31]

CSVファイルを出力する例

import org.apache.spark.sql.SaveMode

  def writeCsv(ds: Dataset[Person], path: String): Unit = {
    ds.write.mode(SaveMode.Overwrite).csv(path)
  }

pathは出力先ディレクトリー。この下にパーティションの個数分の複数のファイル（拡張子はcsv）が出力される。
modeをOverwrite（上書き）にしておかないと、ディレクトリーが既に存在している場合にエラーになる。

ただ、このcsvメソッドでデータを出力すると、各フィールドがtrimされてしまう。（つまり、空白1文字を出力したいと思っても、空文字列になってしまう）
write（DataFrameWriter）にはoptionを設定することは出来る（def csv(path: String)メソッドのScaladocを参照）のだけれども、trimしないようにするオプションは無さげ。

常にダブルクォーテーションで囲むようにすればtrimされないかと思って試してみたけど、trimされてからダブルクォーテーションで囲まれてたorz[2017-02-02]

    ds.write.option("quoteAll", true).mode(SaveMode.Overwrite).csv(path)

ちなみにsortByというメソッドで出力順を制御できそうな感じがしたが、csvファイルは対象外だったorz[2017-02-01]

テキストファイルを出力する例

csvメソッドで出力すると各フィールドがtrimされてしまうが、自前で単一のStringに変換してテキストファイルとして出力することは出来る。[2017-02-01]

import org.apache.spark.sql.SaveMode

  def writeCsv(ds: Dataset[Person], path: String): Unit = {
    ds.map(_.productIterator.mkString(",")).write.mode(SaveMode.Overwrite).text(path)
  }

ケースクラスはProductトレイトをミックスインしているので、productIteratorで各フィールドの値が取れる。
これをmkStringで“カンマ区切りの1つのString”に変換する。

pathは出力先ディレクトリー。この下にパーティションの個数分の複数のファイル（拡張子はtxt）が出力される。

write前にrepartitionでパーティション数を1にしてやれば、出力するファイルを1つにすることは出来る。

  def writeCsv(ds: Dataset[Person], path: String): Unit = {
    ds.map(_.productIterator.mkString(",")).repartition(1).write.mode(SaveMode.Overwrite).text(path)
  }

ちなみに、write（DataFrameWriter）にはpartitionByというパーティションを変更できそうなメソッドがあるのだが、出力ファイルを1つにすることは出来なかった。

Writerを使って出力する失敗例

csvメソッドやtextメソッドではファイル名が自由に出来ないので、自前でWriterを使って出力しようとしてみた。

import java.nio.charset.Charset
import java.nio.file.Files
import java.nio.file.Paths

  def writeByWriter(ds: Dataset[Person], path: String): Unit = {
    Files.createDirectories(Paths.get(path))

    val n = path.lastIndexOf('/');
    val name = path.substring(n + 1) + ".csv"

    val writer = Files.newBufferedWriter(Paths.get(path, name), Charset.forName("UTF-8"))
    try {
      ds.map(_.productIterator.mkString(",")).repartition(1).foreach { s =>
        writer.write(s)
        writer.write('\n')
      }
    } finally {
      writer.close()
    }

foreachでBufferedWirterにwriteする。

しかしいざ実行すると、BufferedWirterがシリアライザブルでない為、例外が発生する。
（foreachの中でwriterを使おうとしているので、シリアライズしてexecutorに渡そうとしているのだと思う）

Writerを使って出力する例

import java.nio.charset.Charset
import java.nio.file.Files
import java.nio.file.Paths

  def writeByIterator(ds: Dataset[Person], path: String): Unit = {
    Files.createDirectories(Paths.get(path))

    val n = path.lastIndexOf('/');
    val name = path.substring(n + 1) + ".csv"

    val writer = Files.newBufferedWriter(Paths.get(path, name), Charset.forName("UTF-8"))
    try {
      import scala.collection.JavaConverters._

      ds.map(_.productIterator.map {
        case null      => ""
        case s: String => if (s.contains(",") || s.contains("\r") || s.contains("\n")) '"' + s.replaceAll("\"", "\"\"") + '"' else s
        case x         => x.toString
      }.mkString(","))
        .toLocalIterator().asScala.foreach { s =>
          writer.write(s)
          writer.write("\r\n")
        }
    } finally {
      writer.close()
    }
  }

Datasetのforeachの中にはwriterを置くことが出来ないので、toLocalIteratorメソッドでDatasetからローカルのIteratorに変換し、そのforeachでBufferedWirterに出力する。