def nestedindex_form(indices):
if not all(index["class"].startswith("ListOffset") for index in indices):
raise RuntimeError
form = {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": copy.deepcopy(indices[0]),
}
# steal offsets from first input
key = []
for index in indices:
key.append(index["content"]["form_key"])
key.append("!nestedindex")
key = concat(*key)
form["form_key"] = indices[0]["form_key"]
form["content"]["form_key"] = key
form["content"]["content"]["form_key"] = concat(key, "!content")
return form
def counts2offsets_form(counts_form):
form = {
"class": "NumpyArray",
"itemsize": 8,
"format": "i",
"primitive": "int64",
"parameters": counts_form.get("parameters", None),
"form_key": concat(counts_form["form_key"], "!counts2offsets"),
}
return form
def counts2nestedindex_form(local_counts, target_offsets):
if not local_counts["class"].startswith("ListOffset"):
raise RuntimeError
if not target_offsets["class"] == "NumpyArray":
raise RuntimeError
form = {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": copy.deepcopy(local_counts),
}
form["content"]["content"]["itemsize"] = 8
form["content"]["content"]["primitive"] = "int64"
form["content"]["content"]["parameters"] = {}
key = concat(local_counts["form_key"], target_offsets["form_key"],
"!counts2nestedindex")
form["form_key"] = local_counts["form_key"]
form["content"]["form_key"] = key
form["content"]["content"]["form_key"] = concat(key, "!content")
return form
def local2global_form(index, target_offsets):
if not index["class"].startswith("ListOffset"):
raise RuntimeError
if not target_offsets["class"] == "NumpyArray":
raise RuntimeError
form = copy.deepcopy(index)
form["content"]["form_key"] = concat(
index["form_key"], target_offsets["form_key"], "!local2global"
)
form["content"]["itemsize"] = 8
form["content"]["primitive"] = "int64"
return form
def children_form(offsets, globalparents):
if not globalparents["class"].startswith("ListOffset"):
raise RuntimeError
form = {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"itemsize": 8,
"format": "i",
"primitive": "int64",
},
},
}
form["form_key"] = offsets["form_key"]
key = concat(offsets["form_key"], globalparents["content"]["form_key"], "!children")
form["content"]["form_key"] = key
form["content"]["content"]["form_key"] = concat(key, "!content")
return form
def listarray_form(content, offsets):
if offsets["class"] != "NumpyArray":
raise ValueError
if offsets["primitive"] == "int32":
arrayclass = "ListOffsetArray32"
offsetstype = "i32"
elif offsets["primitive"] == "int64":
arrayclass = "ListOffsetArray64"
offsetstype = "i64"
else:
raise ValueError("Unrecognized offsets data type")
return {
"class": arrayclass,
"offsets": offsetstype,
"content": content,
"form_key": concat(offsets["form_key"], "!skip"),
}
def distinctParent_form(parents, pdg):
if not parents["class"].startswith("ListOffset"):
raise RuntimeError
if not pdg["class"].startswith("ListOffset"):
raise RuntimeError
form = {
"class": "ListOffsetArray64",
"offsets": "i64",
"content": {
"class": "NumpyArray",
"itemsize": 8,
"format": "i",
"primitive": "int64",
},
"form_key": parents["form_key"],
}
form["content"]["form_key"] = concat(
parents["content"]["form_key"],
pdg["content"]["form_key"],
"!distinctParent",
)
return form
