1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import os
import pytest
import awkward as ak
pa = pytest.importorskip("pyarrow")
pq = pytest.importorskip("pyarrow.parquet")
def test_parquet_subcolumn_select(tmp_path):
ak_tbl = ak.Array(
{
"a": [
{"lbl": "item 1", "idx": 11, "ids": [1, 2, 3]},
{"lbl": "item 2", "idx": 12, "ids": [51, 52]},
{"lbl": "item 3", "idx": 13, "ids": [61, 62, 63, 64]},
],
"b": [
[[111, 112], [121, 122]],
[[211, 212], [221, 222]],
[[311, 312], [321, 322]],
],
}
)
parquet_file = os.path.join(tmp_path, "test_3514.parquet")
ak.to_parquet(ak_tbl, parquet_file)
selection = ak.from_parquet(parquet_file, columns=["a.ids", "b"])
assert selection["a"].to_list() == [
{"ids": [1, 2, 3]},
{"ids": [51, 52]},
{"ids": [61, 62, 63, 64]},
]
assert selection["b"].to_list() == ak_tbl["b"].to_list()
|
