PySpark Dataframes-不转换为熊猫的枚举方式？

Question

我有一个非常大的pyspark.sql.dataframe.dataframe名为df。我需要一些枚举记录的方法--因此，能够访问具有特定索引的记录。（或选择具有索引范围的记录组）

在熊猫身上，我可以

indexes=[2,3,6,7] 
df[indexes]

indexes=np.arange(df.count())
df_indexed=df.withColumn('index', indexes)

null

>

为什么不起作用，如何使其起作用？如何向DataFrame添加行？

以后制作类似：

 indexes=[2,3,6,7] 
 df1.where("index in indexes").collect()

有没有更快更简单的处理方法？

Answer 1

它不起作用是因为：

1. With Column的第二个参数应该是Column，而不是集合。NP.array在这里不起作用
2. 当将“索引中的索引”作为SQL表达式传递给，其中indexs超出范围，且未将其解析为有效标识符
时

PySpark>=1.4.0

from pyspark.sql.functions import col, rowNumber
from pyspark.sql.window import Window

w = Window.orderBy()
indexed = df.withColumn("index", rowNumber().over(w))

# Using DSL
indexed.where(col("index").isin(set(indexes)))

# Using SQL expression
indexed.where("index in ({0})".format(",".join(str(x) for x in indexes)))

独立于PySpark版本，您可以尝试如下所示：

from pyspark.sql import Row
from pyspark.sql.types import StructType, StructField, LongType

row = Row("char")
row_with_index = Row("char", "index")

df = sc.parallelize(row(chr(x)) for x in range(97, 112)).toDF()
df.show(5)

## +----+
## |char|
## +----+
## |   a|
## |   b|
## |   c|
## |   d|
## |   e|
## +----+
## only showing top 5 rows

# This part is not tested but should work and save some work later
schema  = StructType(
    df.schema.fields[:] + [StructField("index", LongType(), False)])

indexed = (df.rdd # Extract rdd
    .zipWithIndex() # Add index
    .map(lambda ri: row_with_index(*list(ri[0]) + [ri[1]])) # Map to rows
    .toDF(schema)) # It will work without schema but will be more expensive

# inSet in Spark < 1.3
indexed.where(col("index").isin(indexes))