Como definir o particionamento do DataFrame?

Comecei a usar o Spark SQL e os DataFrames no Spark 1.4.0. Estou querendo definir um particionador personalizado no DataFrames, no Scala, mas não vendo como fazer isso.

Uma das tabelas de dados com as quais estou trabalhando contém uma lista de transações, por conta, silimar no exemplo a seguir.

Account   Date       Type       Amount
1001    2014-04-01  Purchase    100.00
1001    2014-04-01  Purchase     50.00
1001    2014-04-05  Purchase     70.00
1001    2014-04-01  Payment    -150.00
1002    2014-04-01  Purchase     80.00
1002    2014-04-02  Purchase     22.00
1002    2014-04-04  Payment    -120.00
1002    2014-04-04  Purchase     60.00
1003    2014-04-02  Purchase    210.00
1003    2014-04-03  Purchase     15.00

Pelo menos inicialmente, a maioria dos cálculos ocorrerá entre as transações em uma conta. Então, eu gostaria de ter os dados particionados para que todas as transações de uma conta estejam na mesma partição Spark.

Mas não estou vendo uma maneira de definir isso. A classe DataFrame possui um método chamado 'repartition (Int)', onde você pode especificar o número de partições a serem criadas. Mas não estou vendo nenhum método disponível para definir um particionador personalizado para um DataFrame, como pode ser especificado para um RDD.

Os dados de origem são armazenados no Parquet. Eu vi que ao escrever um DataFrame no Parquet, você pode especificar uma coluna para particionar, então presumivelmente eu poderia dizer ao Parquet para particionar seus dados pela coluna 'Conta'. Mas pode haver milhões de contas e, se eu estiver entendendo o Parquet corretamente, ele criará um diretório distinto para cada conta, para que não pareça uma solução razoável.

Existe uma maneira de fazer com que o Spark particione esse DataFrame para que todos os dados de uma conta estejam na mesma partição?

Spark> = 2.3.0

O SPARK-22614 expõe o particionamento de intervalo.

val partitionedByRange = df.repartitionByRange(42, $"k")

partitionedByRange.explain
// == Parsed Logical Plan ==
// 'RepartitionByExpression ['k ASC NULLS FIRST], 42
// +- AnalysisBarrier Project [_1#2 AS k#5, _2#3 AS v#6]
// 
// == Analyzed Logical Plan ==
// k: string, v: int
// RepartitionByExpression [k#5 ASC NULLS FIRST], 42
// +- Project [_1#2 AS k#5, _2#3 AS v#6]
//    +- LocalRelation [_1#2, _2#3]
// 
// == Optimized Logical Plan ==
// RepartitionByExpression [k#5 ASC NULLS FIRST], 42
// +- LocalRelation [k#5, v#6]
// 
// == Physical Plan ==
// Exchange rangepartitioning(k#5 ASC NULLS FIRST, 42)
// +- LocalTableScan [k#5, v#6]

O SPARK-22389 expõe o particionamento de formato externo na API de fonte de dados v2 .

Spark> = 1.6.0

No Spark> = 1.6, é possível usar o particionamento por coluna para consulta e armazenamento em cache. Consulte: SPARK-11410 e SPARK-4849 usando o repartitionmétodo:

val df = Seq(
  ("A", 1), ("B", 2), ("A", 3), ("C", 1)
).toDF("k", "v")

val partitioned = df.repartition($"k")
partitioned.explain

// scala> df.repartition($"k").explain(true)
// == Parsed Logical Plan ==
// 'RepartitionByExpression ['k], None
// +- Project [_1#5 AS k#7,_2#6 AS v#8]
//    +- LogicalRDD [_1#5,_2#6], MapPartitionsRDD[3] at rddToDataFrameHolder at <console>:27
// 
// == Analyzed Logical Plan ==
// k: string, v: int
// RepartitionByExpression [k#7], None
// +- Project [_1#5 AS k#7,_2#6 AS v#8]
//    +- LogicalRDD [_1#5,_2#6], MapPartitionsRDD[3] at rddToDataFrameHolder at <console>:27
// 
// == Optimized Logical Plan ==
// RepartitionByExpression [k#7], None
// +- Project [_1#5 AS k#7,_2#6 AS v#8]
//    +- LogicalRDD [_1#5,_2#6], MapPartitionsRDD[3] at rddToDataFrameHolder at <console>:27
// 
// == Physical Plan ==
// TungstenExchange hashpartitioning(k#7,200), None
// +- Project [_1#5 AS k#7,_2#6 AS v#8]
//    +- Scan PhysicalRDD[_1#5,_2#6]

Ao contrário do RDDsSpark Dataset(incluindo Dataset[Row]aka DataFrame), não é possível usar o particionador personalizado por enquanto. Normalmente, você pode resolver isso criando uma coluna de particionamento artificial, mas ela não oferece a mesma flexibilidade.

Spark <1.6.0:

Uma coisa que você pode fazer é pré-particionar os dados de entrada antes de criar um DataFrame

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

val schema = StructType(Seq(
  StructField("x", StringType, false),
  StructField("y", LongType, false),
  StructField("z", DoubleType, false)
))

val rdd = sc.parallelize(Seq(
  Row("foo", 1L, 0.5), Row("bar", 0L, 0.0), Row("??", -1L, 2.0),
  Row("foo", -1L, 0.0), Row("??", 3L, 0.6), Row("bar", -3L, 0.99)
))

val partitioner = new HashPartitioner(5) 

val partitioned = rdd.map(r => (r.getString(0), r))
  .partitionBy(partitioner)
  .values

val df = sqlContext.createDataFrame(partitioned, schema)

Como a DataFramecriação de um RDDrequer apenas uma fase simples do mapa, o layout da partição existente deve ser preservado *:

assert(df.rdd.partitions == partitioned.partitions)

Da mesma maneira que você pode reparticionar os existentes DataFrame:

sqlContext.createDataFrame(
  df.rdd.map(r => (r.getInt(1), r)).partitionBy(partitioner).values,
  df.schema
)

Então parece que não é impossível. A questão permanece se faz sentido. Argumentarei que na maioria das vezes isso não acontece:

1. O reparticionamento é um processo caro. Em um cenário típico, a maioria dos dados precisa ser serializada, embaralhada e desserializada. Por outro lado, o número de operações que podem se beneficiar de dados pré-particionados é relativamente pequeno e é ainda mais limitado se a API interna não for projetada para aproveitar essa propriedade.

   • ingressa em alguns cenários, mas exigiria um suporte interno,
   • funções de janela chamadas com o particionador correspondente. O mesmo que acima, limitado a uma única definição de janela. Já está particionado internamente, portanto, o pré-particionamento pode ser redundante,
   • agregações simples com GROUP BY- é possível reduzir a pegada de memória dos buffers temporários **, mas o custo geral é muito maior. Mais ou menos equivalente a groupByKey.mapValues(_.reduce)(comportamento atual) vs reduceByKey(pré-particionamento). É improvável que seja útil na prática.
   • compressão de dados com SqlContext.cacheTable. Como parece que está usando a codificação de duração da execução, a aplicação OrderedRDDFunctions.repartitionAndSortWithinPartitionspode melhorar a taxa de compactação.

2. O desempenho é altamente dependente da distribuição das chaves. Se estiver inclinado, resultará em uma utilização abaixo do ideal. Na pior das hipóteses, será impossível concluir o trabalho.

3. Um ponto importante do uso de uma API declarativa de alto nível é isolar-se dos detalhes de implementação de baixo nível. Como já mencionado por @dwysakowicz e @RomiKuntsman, uma otimização é uma tarefa do Catalyst Optimizer . É uma fera bastante sofisticada e eu realmente duvido que você possa facilmente melhorar isso sem mergulhar muito mais profundamente em suas partes internas.

Conceitos relacionados

Particionando com origens JDBC :

As fontes de dados JDBC suportam predicatesargumento . Pode ser usado da seguinte maneira:

sqlContext.read.jdbc(url, table, Array("foo = 1", "foo = 3"), props)

Ele cria uma única partição JDBC por predicado. Lembre-se de que se os conjuntos criados usando predicados individuais não forem disjuntos, você verá duplicatas na tabela resultante.

partitionBymétodo emDataFrameWriter :

O Spark DataFrameWriterfornece um partitionBymétodo que pode ser usado para "particionar" dados na gravação. Ele separa os dados na gravação usando o conjunto de colunas fornecido

val df = Seq(
  ("foo", 1.0), ("bar", 2.0), ("foo", 1.5), ("bar", 2.6)
).toDF("k", "v")

df.write.partitionBy("k").json("/tmp/foo.json")

Isso permite que o envio de predicado seja lido para consultas com base na chave:

val df1 = sqlContext.read.schema(df.schema).json("/tmp/foo.json")
df1.where($"k" === "bar")

mas não é equivalente a DataFrame.repartition. Em agregações específicas como:

val cnts = df1.groupBy($"k").sum()

ainda exigirá TungstenExchange:

cnts.explain

// == Physical Plan ==
// TungstenAggregate(key=[k#90], functions=[(sum(v#91),mode=Final,isDistinct=false)], output=[k#90,sum(v)#93])
// +- TungstenExchange hashpartitioning(k#90,200), None
//    +- TungstenAggregate(key=[k#90], functions=[(sum(v#91),mode=Partial,isDistinct=false)], output=[k#90,sum#99])
//       +- Scan JSONRelation[k#90,v#91] InputPaths: file:/tmp/foo.json

bucketBymétodo emDataFrameWriter (Spark> = 2.0):

bucketBypossui aplicativos semelhantes, partitionBymas está disponível apenas para tabelas ( saveAsTable). As informações de bucket podem ser usadas para otimizar junções:

// Temporarily disable broadcast joins
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", -1)

df.write.bucketBy(42, "k").saveAsTable("df1")
val df2 = Seq(("A", -1.0), ("B", 2.0)).toDF("k", "v2")
df2.write.bucketBy(42, "k").saveAsTable("df2")

// == Physical Plan ==
// *Project [k#41, v#42, v2#47]
// +- *SortMergeJoin [k#41], [k#46], Inner
//    :- *Sort [k#41 ASC NULLS FIRST], false, 0
//    :  +- *Project [k#41, v#42]
//    :     +- *Filter isnotnull(k#41)
//    :        +- *FileScan parquet default.df1[k#41,v#42] Batched: true, Format: Parquet, Location: InMemoryFileIndex[file:/spark-warehouse/df1], PartitionFilters: [], PushedFilters: [IsNotNull(k)], ReadSchema: struct<k:string,v:int>
//    +- *Sort [k#46 ASC NULLS FIRST], false, 0
//       +- *Project [k#46, v2#47]
//          +- *Filter isnotnull(k#46)
//             +- *FileScan parquet default.df2[k#46,v2#47] Batched: true, Format: Parquet, Location: InMemoryFileIndex[file:/spark-warehouse/df2], PartitionFilters: [], PushedFilters: [IsNotNull(k)], ReadSchema: struct<k:string,v2:double>

* Por layout da partição, quero dizer apenas uma distribuição de dados. partitionedO RDD não possui mais um particionador. ** Assumindo que não haja projeção antecipada. Se a agregação cobre apenas um pequeno subconjunto de colunas, provavelmente não há ganho algum.

No Spark <1.6 Se você criar um HiveContext, não o antigo, SqlContextpode usar o HiveQL DISTRIBUTE BY colX... (garante que cada um dos N redutores obtenha intervalos sem sobreposição de x) & CLUSTER BY colX...(atalho para Distribuir por e Classificar por), por exemplo;

df.registerTempTable("partitionMe")
hiveCtx.sql("select * from partitionMe DISTRIBUTE BY accountId SORT BY accountId, date")

Não tenho certeza de como isso se encaixa na API do Spark DF. Essas palavras-chave não são suportadas no SqlContext normal (observe que você não precisa ter um meta store da seção para usar o HiveContext)

EDIT: Spark 1.6+ agora tem isso na API nativa do DataFrame

Então, para começar com algum tipo de resposta:) - Você não pode

Eu não sou um especialista, mas, tanto quanto eu entendo os DataFrames, eles não são iguais ao rdd e o DataFrame não possui o Partitioner.

Geralmente, a idéia do DataFrame é fornecer outro nível de abstração que lide com esses problemas. As consultas no DataFrame são convertidas em plano lógico que é convertido em operações em RDDs. O particionamento que você sugeriu provavelmente será aplicado automaticamente ou pelo menos deveria ser.

Se você não confiar no SparkSQL, ele fornecerá algum tipo de trabalho ideal, sempre é possível transformar o DataFrame em RDD [Row], conforme sugerido nos comentários.

Use o DataFrame retornado por:

yourDF.orderBy(account)

Não há uma maneira explícita de usar partitionByem um DataFrame, apenas em um PairRDD, mas quando você classifica um DataFrame, ele será usado no LogicalPlan e ajudará quando você precisar fazer cálculos em cada Conta.

Acabei de me deparar com o mesmo problema exato, com um quadro de dados que quero particionar por conta. Eu suponho que quando você diz "deseja que os dados sejam particionados para que todas as transações de uma conta estejam na mesma partição Spark", você o deseja para escala e desempenho, mas seu código não depende disso (como usar mapPartitions()etc), certo?

Consegui fazer isso usando RDD. Mas não sei se essa é uma solução aceitável para você. Depois de ter o DF disponível como um RDD, você pode aplicar repartitionAndSortWithinPartitionspara executar o reparticionamento personalizado dos dados.

Aqui está uma amostra que eu usei:

class DatePartitioner(partitions: Int) extends Partitioner {

  override def getPartition(key: Any): Int = {
    val start_time: Long = key.asInstanceOf[Long]
    Objects.hash(Array(start_time)) % partitions
  }

  override def numPartitions: Int = partitions
}

myRDD
  .repartitionAndSortWithinPartitions(new DatePartitioner(24))
  .map { v => v._2 }
  .toDF()
  .write.mode(SaveMode.Overwrite)