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
|
# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski.
#
# Distributed under the 3-clause BSD license, see accompanying file LICENSE
# or https://github.com/scikit-hep/vector for details.
from __future__ import annotations
import pytest
import vector
dak = pytest.importorskip("dask_awkward")
ak = pytest.importorskip("awkward")
def test_constructor():
x = dak.from_awkward(
ak.Array([{"x": 1, "y": 2}, {"x": 1.1, "y": 2.2}]), npartitions=1
)
vec = vector.Array(x)
assert isinstance(vec, dak.Array)
assert isinstance(vec.compute(), vector.backends.awkward.VectorAwkward2D)
assert ak.all(vec.x.compute() == ak.Array([1, 1.1]))
assert ak.all(vec.y.compute() == ak.Array([2, 2.2]))
def test_necessary_columns():
vec = vector.Array([[{"pt": 1, "phi": 2}], [], [{"pt": 3, "phi": 4}]])
dak_vec = dak.from_awkward(vec, npartitions=1)
cols = next(iter(dak.report_necessary_columns(dak_vec).values()))
# this may seem weird at first: why would one need "phi" and "rho", if one asked for "pt"?
# the reason is that vector will build internally a class with "phi" and "rho",
# see: https://github.com/scikit-hep/vector/blob/608da2d55a74eed25635fd408d1075b568773c99/src/vector/backends/awkward.py#L166-L167
# So, even if one asks for "pt", "phi" and "rho" are as well in order to build the vector class in the first place.
# (the same argument holds true for all other vector classes)
assert cols == frozenset({"phi", "rho"})
|
