f9columnar 0.3.2__tar.gz → 0.3.3__tar.gz

{f9columnar-0.3.2 → f9columnar-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: f9columnar
-Version: 0.3.2
+Version: 0.3.3
 Summary: Columnar analysis utils.
 License: MIT
 Author: Jan Gavranovic

{f9columnar-0.3.2 → f9columnar-0.3.3}/f9columnar/arrays.py

@@ -4,7 +4,7 @@ import copy
 import logging
 from dataclasses import dataclass
 from dataclasses import field as dataclass_field
-from typing import Any
+from typing import Any, Self
 
 import awkward as ak
 import numpy as np
@@ -25,12 +25,12 @@ class BaseArrays:
     def shape(self) -> tuple[int, int]:
         return (len(self.array), len(self.fields))
 
-    def mask(self, mask: ak.Array | np.ndarray, inplace: bool = True) -> ak.Array | BaseArrays:
+    def mask(self, mask: ak.Array | np.ndarray, inplace: bool = True) -> Self:
         if inplace:
             self.array = self.array[mask]
             return self
         else:
-            return self.array[mask]
+            return self.__class__(self.array[mask])
 
     def _check_type(self, value: Any) -> bool:
         raise NotImplementedError
@@ -71,9 +71,9 @@ class BaseArrays:
         return self
 
     def __len__(self) -> int:
-        if type(self.array) is ak.Array:
+        if isinstance(self.array, ak.Array):
             return len(self.array)
-        elif type(self.array) is ak.Record:
+        elif isinstance(self.array, ak.Record):
             return 1
         else:
             raise ValueError("Invalid array type!")
@@ -169,9 +169,7 @@ class Arrays:
     def shape(self) -> tuple[int, int, int]:
         return (len(self), len(self.flat_fields), len(self.jagged_fields))
 
-    def remove_empty(
-        self, return_mask: bool = False, all_fields: bool = True
-    ) -> Arrays | ak.Array | tuple[ak.Array | None, ak.Array] | None:
+    def remove_empty(self, return_mask: bool = False, all_fields: bool = True) -> Arrays | ak.Array | None:
         if self.jagged_arrays is None:
             if return_mask:
                 return None
@@ -185,7 +183,7 @@ class Arrays:
         else:
             return self.mask_flat_inplace(non_empty_mask)
 
-    def mask_flat(self, mask: ak.Array | np.ndarray) -> tuple[ak.Array | None, ak.Array | None]:
+    def mask_flat(self, mask: ak.Array | np.ndarray) -> tuple[FlatArrays | None, JaggedArrays | None]:
         if self.flat_arrays is not None:
             flat_masked = self.flat_arrays.mask(mask, inplace=False)
         else:
@@ -207,7 +205,7 @@ class Arrays:
 
         return self
 
-    def mask_jagged(self, mask: ak.Array | np.ndarray) -> ak.Array | None:
+    def mask_jagged(self, mask: ak.Array | np.ndarray) -> JaggedArrays | None:
         if self.jagged_arrays is not None:
             jagged_masked = self.jagged_arrays.mask(mask, inplace=False)
         else:
@@ -661,3 +659,32 @@ class ArraysProcessor(Processor):
             arrays = self.arrays_handler.make_array_groups(arrays, self.groups)
 
         return {"arrays": arrays}
+
+
+def arrays_to_ak(arrays: Arrays) -> ak.Array:
+    fields: dict[str, ak.Array] = {}
+
+    if arrays.flat_arrays is not None:
+        for field in arrays.flat_arrays.fields:
+            fields[field] = arrays.flat_arrays.array[field]
+
+    if arrays.jagged_arrays is not None:
+        for field in arrays.jagged_arrays.fields:
+            fields[field] = arrays.jagged_arrays.array[field]
+
+    return ak.zip(fields, depth_limit=1)
+
+
+def group_arrays_to_ak(group_arrays: GroupArrays) -> ak.Array:
+    group_fields: dict[str, ak.Array] = {}
+
+    for _, arrays in group_arrays.group_arrays.items():
+        if arrays.flat_arrays is not None:
+            for field in arrays.flat_arrays.fields:
+                group_fields[field] = arrays.flat_arrays.array[field]
+
+        if arrays.jagged_arrays is not None:
+            for field in arrays.jagged_arrays.fields:
+                group_fields[field] = arrays.jagged_arrays.array[field]
+
+    return ak.zip(group_fields, depth_limit=1)

{f9columnar-0.3.2 → f9columnar-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "f9columnar"
-version = "0.3.2"
+version = "0.3.3"
 description = "Columnar analysis utils."
 authors = ["Jan Gavranovic <jan.gavranovic@cern.ch>"]
 readme = "README.md"