31
31
import numpy as np
32
32
import awkward as ak
33
33
import json
34
+ import logging
34
35
35
36
36
37
def run_query (input_filenames ):
@@ -123,6 +124,30 @@ def build_deliver_spec(datasets):
123
124
return spec_python
124
125
125
126
127
+ def open_delivered_file (sample , path ):
128
+ """
129
+ Opens the first file delivered by ServiceX for a given sample and returns the
130
+ structure encoded in the "servicex/branch" branch.
131
+ If no files are found, logs a warning and returns None.
132
+ Parameters:
133
+ sample (str): The sample name for which to open the file.
134
+ path (list): List of file paths delivered by ServiceX for the sample.
135
+ """
136
+
137
+ if not path :
138
+ logging .warning (
139
+ f"Warning: No files found for sample '{ sample } ' in delivered results. Skipping."
140
+ )
141
+ return None
142
+
143
+ try :
144
+ with uproot .open (path [0 ]) as f :
145
+ return f ["servicex" ]["branch" ].array ()[0 ]
146
+ except Exception as e :
147
+ logging .error (f"Error opening file for sample '{ sample } ': { e } " )
148
+ return None
149
+
150
+
126
151
def print_structure_from_str (
127
152
deliver_dict , filter_branch = "" , save_to_txt = False , do_print = False
128
153
):
@@ -147,16 +172,18 @@ def print_structure_from_str(
147
172
)
148
173
149
174
for sample_name , path in deliver_dict .items ():
175
+ structure_str = open_delivered_file (sample_name , path )
176
+ if structure_str is None :
177
+ continue
178
+ # Parse the JSON string into a dictionary
179
+ structure_dict = json .loads (structure_str )
180
+
150
181
output_lines .append (
151
182
f"\n ---------------------------\n "
152
183
f"\U0001f4c1 Sample: { sample_name } \n "
153
184
f"---------------------------"
154
185
)
155
186
156
- with uproot .open (path [0 ]) as f :
157
- json_str = f ["servicex" ]["branch" ].array ()[0 ]
158
- structure_dict = json .loads (json_str )
159
-
160
187
for tree_name , branches in structure_dict .items ():
161
188
output_lines .append (f"\n \U0001f333 Tree: { tree_name } " )
162
189
output_lines .append (" ├── Branches:" )
@@ -272,8 +299,9 @@ def get_structure(datasets, array_out=False, **kwargs):
272
299
if array_out == True :
273
300
all_arrays = {}
274
301
for sample , path in output .items ():
275
- with uproot .open (path [0 ]) as f :
276
- structure_str = f ["servicex" ]["branch" ].array ()[0 ]
302
+ structure_str = open_delivered_file (sample , path )
303
+ if structure_str is None :
304
+ continue
277
305
sample_array = str_to_array (structure_str )
278
306
all_arrays [sample ] = sample_array
279
307
return all_arrays
0 commit comments