From cc7504ea35517c599f0937bc54d28d6df8a9e4f9 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 30 Jul 2024 10:03:23 -0500 Subject: [PATCH] fix: ak.to_parquet for a record-like non-Record (#3175) * fix: ak.to_parquet for a record-like non-Record * style: pre-commit fixes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Ianna Osborne --- src/awkward/operations/ak_to_parquet.py | 2 +- tests/test_3174_to_parquet_record.py | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 tests/test_3174_to_parquet_record.py diff --git a/src/awkward/operations/ak_to_parquet.py b/src/awkward/operations/ak_to_parquet.py index 35bf3995fe..971b578208 100644 --- a/src/awkward/operations/ak_to_parquet.py +++ b/src/awkward/operations/ak_to_parquet.py @@ -314,7 +314,7 @@ def get_layout_and_table(x): else: column_prefix = () - if isinstance(data, ak.Record): + if isinstance(layout, ak.record.Record): form = layout.array.form else: form = layout.form diff --git a/tests/test_3174_to_parquet_record.py b/tests/test_3174_to_parquet_record.py new file mode 100644 index 0000000000..5ef025c548 --- /dev/null +++ b/tests/test_3174_to_parquet_record.py @@ -0,0 +1,25 @@ +# 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 + +pytest.importorskip("pyarrow.parquet") + + +def test(tmp_path): + filename = os.path.join(tmp_path, "whatever.parquet") + + px = ak.Array([[1.0, 2.0], [3.0], [], [4.0, 5.0, 6.0]]) + py = ak.Array([[0.5, 1.5], [2.5], [], [3.5, 4.5, 5.5]]) + pz = ak.Array([[0.2, 1.2], [2.2], [], [3.2, 4.2, 5.2]]) + e = ak.Array([[0.1, 1.1], [2.1], [], [3.1, 4.1, 5.1]]) + data = {"px": px, "py": py, "pz": pz, "e": e} # not quite a record + + ak.to_parquet(data, filename) + + assert ak.from_parquet(filename).tolist() == ak.to_list(ak.to_layout(data))