Removing decode-raw and implementing it in file_peeking

Author: ArturU043

servicex_analysis_utils/decode-raw-test.py

@@ -28,6 +28,8 @@
 
 import servicex
 import uproot
+import numpy as np
+import awkward as ak
 
 def run_query(input_filenames=None):
     import uproot
@@ -86,9 +88,53 @@ def is_tree(obj):
     # Return str in an array
     return ak.Array([final_str])
 
+def build_deliver_spec(dataset):
+    """
+    Helper to build the servicex.deliver configuration.
+    Supports multiple inputs for multiple sample queries.
+
+    Parameters:
+    dataset (str, [str], or dict): Rucio DIDs to be checked by the servicex workers. 
+                                   If dict, custom names can be inputed
+
+    Returns:
+    spec_python (dict): The specification for the python function query containing Name, Query, Dataset, NFiles
+    """
+    #Servicex query using the PythonFunction backend
+    query_PythonFunction = servicex.query.PythonFunction().with_uproot_function(run_query)
+    
+    #Create a dict with sample name for ServiceX query & datasetID
+    dataset_dict={}
+    user_in=type(dataset)
+    
+    if user_in == str:
+        dataset_dict.update({"Sample":dataset})
+    elif user_in == list and type(dataset[0]) is str:
+        for i in range(len(dataset)):
+            name="Sample"+str(i+1) #write number for humans
+            dataset_dict.update({name:dataset[i]})
+    elif user_in == dict:
+        dataset_dict=dataset
+    else:
+        raise ValueError(f"Unsupported dataset input type: {user_in}.\nInput must be dict ('sample_name':'dataset_id'), str or list of str")
+    
+    sample_list = [
+        {
+            "NFiles": 1,
+            "Name": name,
+            "Dataset": servicex.dataset.Rucio(did),
+            "Query": query_PythonFunction,
+        }
+        for name, did in dataset_dict.items()
+    ]
+    spec_python = {"Sample": sample_list}
+
+    return spec_python   
+
 def print_structure_from_str(deliver_dict, filter_branch="", save_to_txt=False, do_print=False):
     """
-    Converts dataset file structures to a formatted string.
+    Re-formats the deliver-retrieve str structure for readability with a filter for branch selection. 
+    The string can be printed, written out or returned
 
     Parameters:
       deliver_dict (dict): ServiceX deliver output (keys: sample names, values: file paths or URLs).
@@ -146,50 +192,106 @@ def print_structure_from_str(deliver_dict, filter_branch="", save_to_txt=False,
     else:
         return result_str
 
-def build_deliver_spec(dataset):
+def parse_jagged_depth_and_dtype(dtype_str):
     """
-    Helper to build the servicex.deliver dict configuration.
-    Supports multiple inputs for multiple sample queries.
+    Helper to decode the dtype str for each branch.
+
+    Parses uproot-style interpretation strings such as:
+    - "AsJagged(AsJagged(AsDtype('>f4')))"
+
+    Returns the number of nested layers and the inner dtype. 
+    Used in str_to_array to reconstruct the ak.array.
 
     Parameters:
-    dataset (str, [str], or dict): Rucio DIDs to be checked by the servicex workers. 
-                                   If dict, custom names can be inputed
+        dtype_str (str): The dtype part of a branch info str; from the delivered file structure.
 
     Returns:
-    spec_python (dict): The specification for the python function query containing Name, Query, Dataset, NFiles
+        int, str: jagged_depth, base_numpy_dtype_str or None if not recognized.
     """
-    #Servicex query using the PythonFunction backend
-    query_PythonFunction = servicex.query.PythonFunction().with_uproot_function(run_query)
-    
-    #Create a dict with sample name for ServiceX query & datasetID
-    dataset_dict={}
-    user_in=type(dataset)
-    
-    if user_in == str:
-        dataset_dict.update({"Sample":dataset})
-    elif user_in == list and type(dataset[0]) is str:
-        for i in range(len(dataset)):
-            name="Sample"+str(i+1) #write number for humans
-            dataset_dict.update({name:dataset[i]})
-    elif user_in == dict:
-        dataset_dict=dataset
+    depth = 0
+    current = dtype_str.strip()
+
+    # Count how many nested AsJagged(...) wrappers exist
+    while current.startswith("AsJagged("):
+        depth += 1
+        current = current[len("AsJagged("):-1].strip()  # Strip outermost wrapper, up to -1 to remove )
+
+    # Extract the base dtype string from AsDtype('<np-format>')
+    if current.startswith("AsDtype('") and current.endswith("')"):
+        base_dtype = current[len("AsDtype('"):-2]
+        return depth, base_dtype
     else:
-        raise ValueError(f"Unsupported dataset input type: {user_in}.\nInput must be dict ('sample_name':'dataset_id'), str or list of str")
-    
-    sample_list = [
-        {
-            "NFiles": 1,
-            "Name": name,
-            "Dataset": servicex.dataset.Rucio(did),
-            "Query": query_PythonFunction,
-        }
-        for name, did in dataset_dict.items()
-    ]
-    spec_python = {"Sample": sample_list}
+        return depth, None
 
-    return spec_python   
+def str_to_array(encoded_str):
+    """
+    Helper to reconstruct ak.Arrays from an encoded file-structure string.
+    Retruned array mimicks TTrees, TBranches with correct field names and dtypes. 
+
+    Parameters:
+        encoded_str (str): The encoded string from run_query.
 
-def get_structure(dataset, **kwargs):
+    Returns:
+        reconstructed_data (ak.Array): Contains trees and branches with typed dumy values.
+    """
+
+    #Separate trees
+    tree_sections = encoded_str.strip().split("\n")
+    reconstructed_data = {}
+
+    for tree_section in tree_sections:
+        tree_section = tree_section.strip()
+        if not tree_section:
+            continue #skip empty lines
+
+        parts = tree_section.split(";", 1) #Tree and branches separated by ; 
+        tree_header = parts[0].strip()
+
+        # Extract tree name
+        treename = tree_header[len("Tree: "):]
+        branches = {}
+
+        if len(parts) > 1:
+            branches_str = parts[1].strip()
+            branch_infos = branches_str.split(",") #Branches separated by
+
+            for branch in branch_infos:
+                branch = branch.strip()
+
+                if " ; dtype: " in branch:  # line with branch info
+                    name_str, dtype_str = branch.split(" ; dtype: ", 1)
+                    # Extract name
+                    branch_name = name_str[len("TBranch: "):].strip()
+                    dtype_str = dtype_str.strip()
+
+                    # Get nesting depth and base dtype from interpretation string
+                    depth, base_dtype_str = parse_jagged_depth_and_dtype(dtype_str)
+                    if base_dtype_str is None:
+                        branches[branch_name] = None
+                        continue
+
+                    try:
+                        np_dtype = np.dtype(base_dtype_str)
+                    except TypeError:
+                        branches[branch_name] = None
+                        continue
+
+                    dummy = np_dtype.type(0) # Typed placeholder value
+
+                    # Simulate jagged structure by nesting the value in lists
+                    for _ in range(depth):
+                        dummy = [dummy] 
+
+                    # Wrap dummy in a length-1 ak.Array
+                    branches[branch_name] = ak.Array([dummy])
+
+        if branches:
+            # Each tree becomes a record array with 1 entry (dict of branch arrays)
+            reconstructed_data[treename] = ak.Array([branches])
+
+    return ak.Array(reconstructed_data)
+
+def get_structure(dataset, array_out=False, **kwargs):
     """
     Utility function. 
     Creates and sends the ServiceX request from user inputed datasets to retrieve file stucture.
@@ -204,4 +306,14 @@ def get_structure(dataset, **kwargs):
 
     output=servicex.deliver(spec_python)
 
-    return print_structure_from_str(output, **kwargs)
+    if array_out==True:
+        all_arrays={}
+        for sample, path in output.items():
+            with uproot.open(path[0]) as f:
+                structure_str = f["servicex"]["branch"].array()[0] 
+            sample_array=str_to_array(structure_str)
+            all_arrays[sample]=sample_array 
+        return all_arrays
+    
+    else:
+        return print_structure_from_str(output, **kwargs)