Início / coding / Perguntas / 78908629
Accepted
ignoring_gravity
Asked: 2024-08-24 17:39:07 +0800 CST

Como compactar dois arrays PyArrow?

  • 772

Em Polars, posso usar zip_widthpara tirar valores de s1ou s2de acordo com uma máscara:

In [1]: import polars as pl

In [2]: import pyarrow as pa

In [3]: import pyarrow as pc

In [4]: s1 = pl.Series([1,2,3])

In [5]: mask = pl.Series([True, False, False])

In [6]: s2 = pl.Series([4, 5, 6])

In [7]: s1.zip_with(mask, s2)
Out[7]:
shape: (3,)
Series: '' [i64]
[
        1
        5
        6
]

Como posso fazer isso com PyArrow? Eu tentei pyarrow.compute.replace_with_mask, mas funciona de maneira diferente:

In [10]: import pyarrow.compute as pc

In [11]: import pyarrow as pa

In [12]: a1 = pa.array([1,2,3])

In [13]: mask = pa.array([True, False, False])

In [14]: a2 = pa.array([4,5,6])

In [15]: pc.replace_with_mask(a1, pc.invert(mask), a2)
Out[15]:
<pyarrow.lib.Int64Array object at 0x7f69d411afe0>
[
  1,
  4,
  5
]

Como replicar zip_withno PyArrow?

python

1 respostas

  1. Best Answer

    Você pode usar if_elsea função de computação do PyArrow:

    import pyarrow as pa
import pyarrow.compute as pc

# Your input arrays
a1 = pa.array([1, 2, 3])
mask = pa.array([True, False, False])
a2 = pa.array([4, 5, 6])

# Apply the if_else function
result = pc.if_else(mask, a1, a2)

print(result)

    [1, 5, 6]

    A if_elsefunção compacta efetivamente os arrays com base na máscara, replicando o comportamento de Polars' zip_with.

    • 2

relate perguntas