numpy 什么是最快的方法来获得非重复行在pyarrow表?

t8e9dugd  于 2023-10-19  发布在  其他
关注(0)|答案(2)|浏览(93)

我试图获取关于我的pyarrow表中两列中的值的不同组合的信息。
我目前正在做的是:

import pandas as pd
import pyarrow as pa
my_table = pa.Table.from_pandas(
  pd.DataFrame(
    {
      'col1':['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'], 
      'col2':[1,1,2,2,1,1,2,3],
      'col3':[1,2,3,4,5,6,7,8]
    }
  )
)
a = [i.to_numpy().astype('str') for i in my_table.select(['col1', 'col2']).columns]
unique = np.unique(np.array(a), axis = 1)

它返回以下项的预期结果:

unique
>array([['a', 'a', 'b', 'b', 'b'],
       ['1', '2', '1', '2', '3']], dtype='<U21')

但是对于大table来说这是相当慢的,我希望有更快的方法?
或者,我真正想知道的是,当我试图写一个分区的数据集时,如何提前知道它将写入哪些目录(即,哪些分区在我的表中有一些数据)
编辑:
它可以更快地转换为pandas而不是多个numpy数组,然后使用drop_duplicates()

my_table.select(['col1', 'col2']).to_pandas().drop_duplicates()
ql3eal8s

ql3eal8s1#

对结构体直接编码的支持由https://issues.apache.org/jira/browse/ARROW-3978跟踪
同时,这里有一个解决方案,它在计算上类似于pandas的unique-ing功能,但通过使用pyarrow自己的计算内核避免了转换到pandas的成本。

import pyarrow as pa
import pyarrow.compute as pc

def _dictionary_and_indices(column):
    assert isinstance(column, pa.ChunkedArray)

    if not isinstance(column.type, pa.DictionaryType):
        column = pc.dictionary_encode(column, null_encoding_behavior='encode')

    dictionary = column.chunk(0).dictionary
    indices = pa.chunked_array([c.indices for c in column.chunks])

    if indices.null_count != 0:
        # We need nulls to be in the dictionary so that indices can be
        # meaningfully multiplied, so we must round trip through decoded
        column = pc.take(dictionary, indices)
        return _dictionary_and_indices(column)

    return dictionary, indices

def unique(table):
    "produce a table containing only the unique rows from the input"
    if table.num_columns == 0:
        return None

    table = table.unify_dictionaries()

    dictionaries = []
    fused_indices = None

    for c in table.columns:
        dictionary, indices = _dictionary_and_indices(c)

        if fused_indices is None:
            fused_indices = indices
        else:
            # pack column's indices into fused_indices
            fused_indices = pc.add(
                pc.multiply(fused_indices, len(dictionary)),
                indices)

        dictionaries.append(dictionary)

    uniques = []

    # pc.unique can now be invoked on the single array of fused_indices
    fused_indices = pc.unique(fused_indices)

    for dictionary in reversed(dictionaries):
        # unpack the column's indices from fused_indices
        quotient = pc.divide(fused_indices, len(dictionary))
        remainder = pc.subtract(fused_indices,
                                pc.multiply(quotient, len(dictionary)))

        # decode this column's uniques
        uniques.insert(0, pc.take(dictionary, remainder))
        fused_indices = quotient

    return pa.Table.from_arrays(uniques, names=table.column_names)

if __name__ == '__main__':
    my_table = pa.table({
        'col1': ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'],
        'col2': [1,   1,   2,   2,   1,   1,   2,   3],
        'col3': [1,   2,   3,   4,   5,   6,   7,   8],
    })

    assert unique(my_table.select(['col1', 'col2'])).equals(pa.table({
        'col1': ['a', 'a', 'b', 'b', 'b'],
        'col2': [1,   2,   1,   2,   3],
    }))
jtoj6r0c

jtoj6r0c2#

也可以使用聚合函数

table.group_by([<column_names>]).aggregate([])

相关问题