如何查询pyarrow表的struct字段

我有一个表，假设有两列a (list) B (list)和两行:

A: ["X", "Y"], ["Y", "Z"]
B: [1, 3], [5, 6]

我想实现像SELECT * FROM table WHERE A.Y = 5这样的东西，它会返回一个(第二)行。我如何使用pyarrow.compute实现这一点?

我正在研究pc.index_in和pc.equals，但他们不处理列表。我也试过C = StructArray.from_arrays([A, B])，但再次，我不知道如何写SELECT * FROM table WHERE C.Y = 5。

我也没有找到一个方法，来"normalize"/"unstack";通过使用Y，表将变成:

A: ["X"], ["Z"]
B: [1], [6]
Y: 3, 5

这可能是一个错误的设计，但基本上我正在寻找两个函数，如
Give me the vector of indexes of 'Y' in column A.-这将返回[1,0]在上面的例子中
然后Give me the values from column B at index {the vector from the previous result, so [1,0]}，这将导致[3,5]。

对于pyarrow中的列表数组，您可以使用list_flatten和list_flatten_indices来爆炸表:

import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
df = pd.concat(
[
pd.Series([["X", "Y"], ["Y", "Z"]], name="A"),
pd.Series([[1, 3], [5, 6]], name="B"),
],
axis=1,
)
table = pa.Table.from_pandas(df)

flat_table = pa.Table.from_arrays(
[
pc.list_flatten(table["A"]),
pc.list_flatten(table["B"]),
pc.list_parent_indices(table["A"]),
],
names=["A", "B", "index"],
)

相关内容

最新更新

热门标签：