python之spark posexplode 函数运行速度很慢

我有一个存储为 orc 的 spark 数据框，其中包含大约 10000 行和以下架构:

>>> df.printSchema() 
root 
 |-- contig: string (nullable = true) 
 |-- start: integer (nullable = true) 
 |-- ref: string (nullable = true) 
 |-- alt: string (nullable = true) 
 |-- gt: array (nullable = true) 
 |    |-- element: integer (containsNull = true) 

其中 arrayField 是一个包含 200000 个整数的列表。我想将其转换为具有平面结构的数据框:

>>> from pyspark.sql.functions import posexplode 
>>> flat = df.select('contig', 'start', 'ref', 'alt', posexplode(df.gt)) 
>>> flat.explain() 
== Physical Plan == 
*Project [contig#0, start#1, ref#2, alt#3, pos#11, col#12] 
+- Generate posexplode(gt#4), true, false, [pos#11, col#12] 
   +- *FileScan orc [contig#0,start#1,ref#2,alt#3,gt#4] Batched: false, Format: ORC, Location: InMemoryFileIndex[file:/path/to/data], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<contig:string,start:int,ref:string,alt:string,gt:array<int>> 
>>> flat.write.orc('/path/to/output/file') 

在具有 24 个 CPU 内核和超过 100GB 内存的机器上，将扁平化数据帧写入文件需要五个多小时。这只是 posexplode 函数的一个特性还是有其他问题？

请您参考如下方法:

似乎 spark 对这里的行做了一些疯狂的事情。使用 RDD，我能够获得更好的性能(每个 CPU 内核每秒 1/3 行，而数据帧的每个 CPU 内核每秒 1/40 行)。不过，这仍然不是很快。

df = sql_context.read.orc('/path/to/source/file') 
rdd = df.rdd 
 
def expand(row): 
    contig, start, ref, alt, gt = row 
    def getrow(index, genotype): 
        return contig, start, ref, alt, index, genotype 
    return [getrow(index, genotype) for index, genotype in enumerate(gt)] 
 
rdd_flat = rdd.flatMap(expand) 
schema = ('contig', 'start', 'ref', 'alt', 'index', 'genotype') 
sqlc.createDataFrame(rdd_flat, schema=schema).write.orc('/path/to/output/file') 

有趣的是，如果我将 expand-function 重新定义为

def expand(row): 
    def getrow(index, genotype): 
        return Row( 
            contig=row.contig, 
            start=row.start, 
            ref=row.ref, 
            alt=row.alt, 
            index=index, 
            genotype=genotype 
        ) 
    return [getrow(index, genotype) for index, genotype in enumerate(row.gt)] 

它的运行速度慢了大约 13 倍(单个函数调用大约需要 1.4 秒)。

很明显，行对象效率极低。

不过，还有更多问题需要解决。单核应该能够每秒运行9次expand-function，但实际性能是每3秒1行。

编辑:找到一个“解决方案”:使用 prestodb 查询而不是 spark。每个 cpu 内核每秒运行 1 行多一点——比数据帧快 20 多倍，比 RDD 快 4 倍:

create table flat ( 
  contig varchar, 
  start int, 
  ref varchar, 
  alt varchar, 
  index bigint, 
  genotype tinyint 
)  
WITH (format = 'ORC'); 
 
insert into flat 
select contig, start, ref, alt, index, genotype, partition_name 
from nested cross join unnest(gt) with ordinality as g (genotype, index) 
where partition_name='10-70329347';