Skip to content

Commit

Permalink
add tests for for each array type
Browse files Browse the repository at this point in the history
  • Loading branch information
@pfackeldey
pfackeldey committed Nov 22, 2024
1 parent 5246841 commit 2ee9feb
Showing 1 changed file with 257 additions and 0 deletions.
257 changes: 257 additions & 0 deletions tests/test_3311_form_with_key_path.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,257 @@
from __future__ import annotations

import numpy as np

import awkward as ak


def test_record_tuple():
form = ak.forms.from_dict(
{
"class": "RecordArray",
"fields": None,
"contents": [
{"class": "NumpyArray", "primitive": "int64", "form_key": "('0',)"},
{"class": "NumpyArray", "primitive": "int64", "form_key": "('1',)"},
],
"form_key": "()",
}
)
array = ak.Array([(1, 2)])
assert array.layout.form_with_key_path() == form


def test_record_dict():
form = ak.forms.from_dict(
{
"class": "RecordArray",
"fields": ["foo", "bar"],
"contents": [
{"class": "NumpyArray", "primitive": "int64", "form_key": "('foo',)"},
{"class": "NumpyArray", "primitive": "int64", "form_key": "('bar',)"},
],
"form_key": "()",
}
)
array = ak.Array({"foo": [1], "bar": [2]})
assert array.layout.form_with_key_path() == form


def test_numpy():
form = ak.forms.from_dict(
{"class": "NumpyArray", "primitive": "int64", "form_key": "()"}
)
array = ak.Array([1, 2, 3])
assert array.layout.form_with_key_path() == form


def test_listoffset():
form = ak.forms.from_dict(
{
"class": "ListOffsetArray",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"primitive": "int64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
array = ak.Array([[1, 2], [3]])
assert array.layout.form_with_key_path() == form


def test_empty():
form = ak.forms.from_dict({"class": "EmptyArray", "form_key": "()"})
array = ak.Array([])
assert array.layout.form_with_key_path() == form


def test_bitmasked():
form = ak.forms.from_dict(
{
"class": "BitMaskedArray",
"mask": "u8",
"valid_when": True,
"lsb_order": False,
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
content = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4]).layout
mask = ak.index.IndexU8(
np.packbits(np.array([False, True, True, False, False], dtype=np.int8))
)
bitmaskedarray = ak.contents.BitMaskedArray(mask, content, True, 5, False)
assert bitmaskedarray.form_with_key_path() == form


def test_bytemasked():
form = ak.forms.from_dict(
{
"class": "ByteMaskedArray",
"mask": "i8",
"valid_when": True,
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
content = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4]).layout
mask = ak.index.Index8(np.array([False, True, True, False, False], dtype=np.int8))
bytemaskedarray = ak.contents.ByteMaskedArray(mask, content, True)
assert bytemaskedarray.form_with_key_path() == form


def test_indexedarray():
form = ak.forms.from_dict(
{
"class": "IndexedArray",
"index": "i64",
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(None,)",
},
"form_key": "()",
}
)

content = ak.Array([0.0, 1.1, 2.2, 3.3, 4.4]).layout
index = ak.index.Index64(np.array([3, 1, 1, 4, 2], dtype=np.int64))
indexedarray = ak.contents.IndexedArray(index, content)
assert indexedarray.form_with_key_path() == form


def test_indexedoptionarray():
form = ak.forms.from_dict(
{
"class": "IndexedOptionArray",
"index": "i64",
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
content = ak.Array(
np.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
).layout
index = ak.index.Index64(np.array([2, -1, 4, 0, 8], dtype=np.int64))
layout = ak.Array(ak.contents.IndexedOptionArray(index, content)).layout
assert layout.form_with_key_path() == form


def test_listarray():
form = ak.forms.from_dict(
{
"class": "ListArray",
"starts": "i64",
"stops": "i64",
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
content = ak.contents.NumpyArray(
np.array([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])
)
starts = ak.index.Index64(np.array([0, 3, 3, 5, 6]))
stops = ak.index.Index64(np.array([3, 3, 5, 6, 9]))
layout = ak.contents.ListArray(starts, stops, content)
assert layout.form_with_key_path() == form


def test_regulararray():
form = ak.forms.from_dict(
{
"class": "RegularArray",
"size": 3,
"content": {
"class": "NumpyArray",
"primitive": "int64",
"form_key": "(None,)",
},
"form_key": "()",
}
)
content = ak.Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]).layout
layout = ak.contents.RegularArray(content, 3, zeros_length=0)
assert layout.form_with_key_path() == form


def test_unionarray():
form = ak.forms.from_dict(
{
"class": "UnionArray",
"tags": "i8",
"index": "i64",
"contents": [
{
"class": "ListOffsetArray",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"primitive": "float64",
"form_key": "(0, None)",
},
"form_key": "(0,)",
},
{
"class": "ListOffsetArray",
"offsets": "i64",
"content": {
"class": "ListOffsetArray",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"primitive": "uint8",
"parameters": {"__array__": "char"},
"form_key": "(1, None, None)",
},
"parameters": {"__array__": "string"},
"form_key": "(1, None)",
},
"form_key": "(1,)",
},
],
"form_key": "()",
}
)
content1 = ak.operations.from_iter([[], [1.1], [2.2, 2.2]], highlevel=False)
content2 = ak.operations.from_iter([["two", "two"], ["one"], []], highlevel=False)
tags = ak.index.Index8(np.array([0, 1, 0, 1, 0, 1], dtype=np.int8))
index = ak.index.Index64(np.array([0, 0, 1, 1, 2, 2], dtype=np.int64))
layout = ak.contents.UnionArray(tags, index, [content1, content2])
assert layout.form_with_key_path() == form


def test_unmaskedarray():
form = ak.forms.from_dict(
{
"class": "UnmaskedArray",
"content": {
"class": "NumpyArray",
"primitive": "int64",
"form_key": "(None,)",
},
"form_key": "()",
}
)

content = ak.Array([1, 2, 3, 4, 5]).layout
layout = ak.contents.UnmaskedArray(content)
assert layout.form_with_key_path() == form

0 comments on commit 2ee9feb

Please sign in to comment.