У меня есть соединение JDBC с Apache Spark и PostgreSQL, и я хочу вставить некоторые данные в мою базу данных. Когда я использую режим append
, мне нужно указать id
для каждого DataFrame.Row
. Есть ли способ Spark создать первичные ключи?
Первичные ключи с Apache Spark
Ответ 1
Scala
Если вам нужны только уникальные номера, вы можете использовать zipWithUniqueId
и воссоздать DataFrame. Сначала импортируются данные и фиктивные данные:
import sqlContext.implicits._
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.{StructType, StructField, LongType}
val df = sc.parallelize(Seq(
("a", -1.0), ("b", -2.0), ("c", -3.0))).toDF("foo", "bar")
Извлечь схему для дальнейшего использования:
val schema = df.schema
Добавить поле id:
val rows = df.rdd.zipWithUniqueId.map{
case (r: Row, id: Long) => Row.fromSeq(id +: r.toSeq)}
Создать DataFrame:
val dfWithPK = sqlContext.createDataFrame(
rows, StructType(StructField("id", LongType, false) +: schema.fields))
То же самое в Python:
from pyspark.sql import Row
from pyspark.sql.types import StructField, StructType, LongType
row = Row("foo", "bar")
row_with_index = Row(*["id"] + df.columns)
df = sc.parallelize([row("a", -1.0), row("b", -2.0), row("c", -3.0)]).toDF()
def make_row(columns):
def _make_row(row, uid):
row_dict = row.asDict()
return row_with_index(*[uid] + [row_dict.get(c) for c in columns])
return _make_row
f = make_row(df.columns)
df_with_pk = (df.rdd
.zipWithUniqueId()
.map(lambda x: f(*x))
.toDF(StructType([StructField("id", LongType(), False)] + df.schema.fields)))
Если вы предпочитаете последовательный номер, вы можете заменить zipWithUniqueId
на zipWithIndex
, но он немного дороже.
Непосредственно с DataFrame
API:
(универсальный Scala, Python, Java, R с почти одинаковым синтаксисом)
Раньше я пропустил функцию monotonicallyIncreasingId
, которая должна работать нормально, если вам не нужны последовательные номера:
import org.apache.spark.sql.functions.monotonicallyIncreasingId
df.withColumn("id", monotonicallyIncreasingId).show()
// +---+----+-----------+
// |foo| bar| id|
// +---+----+-----------+
// | a|-1.0|17179869184|
// | b|-2.0|42949672960|
// | c|-3.0|60129542144|
// +---+----+-----------+
Хотя полезный monotonicallyIncreasingId
не является детерминированным. Не только идентификаторы могут отличаться от исполнения к исполнению, но без дополнительных трюков нельзя использовать для идентификации строк, когда последующие операции содержат фильтры.
Примечание:
Также возможно использовать функцию окна rowNumber
:
from pyspark.sql.window import Window
from pyspark.sql.functions import rowNumber
w = Window().orderBy()
df.withColumn("id", rowNumber().over(w)).show()
К сожалению:
Окно WARN: для работы с окном не определено раздел. Перемещение всех данных в один раздел может привести к серьезному ухудшению производительности.
Итак, если у вас нет естественного способа разбить ваши данные и обеспечить уникальность, в данный момент это не особенно полезно.
Ответ 2
from pyspark.sql.functions import monotonically_increasing_id
df.withColumn("id", monotonically_increasing_id()).show()
Обратите внимание, что второй аргумент df.withColumn является monotonically_increasing_id() не monotonically_increasing_id.
Ответ 3
Я нашел следующее решение относительно простым для случая, когда zipWithIndex() является желаемым поведением, т.е. для тех, которые задают последовательные целые числа.
В этом случае мы используем pyspark и полагаемся на понимание словаря, чтобы сопоставить исходный объект строки с новым словарем, который соответствует новой схеме, включая уникальный индекс.
# read the initial dataframe without index
dfNoIndex = sqlContext.read.parquet(dataframePath)
# Need to zip together with a unique integer
# First create a new schema with uuid field appended
newSchema = StructType([StructField("uuid", IntegerType(), False)]
+ dfNoIndex.schema.fields)
# zip with the index, map it to a dictionary which includes new field
df = dfNoIndex.rdd.zipWithIndex()\
.map(lambda (row, id): {k:v
for k, v
in row.asDict().items() + [("uuid", id)]})\
.toDF(newSchema)