diff --git a/src/uproot/_dask.py b/src/uproot/_dask.py index aec1b6e02..14923de35 100644 --- a/src/uproot/_dask.py +++ b/src/uproot/_dask.py @@ -751,6 +751,7 @@ def __init__( interp_options, form_mapping, rendered_form, + original_form=None, ) -> None: self.ttrees = ttrees self.common_keys = common_keys @@ -758,12 +759,14 @@ def __init__( self.interp_options = interp_options self.form_mapping = form_mapping self.rendered_form = rendered_form + self.original_form = original_form def __call__(self, i_start_stop): i, start, stop = i_start_stop if self.form_mapping is not None: awkward = uproot.extras.awkward() + dask_awkward = uproot.extras.dask_awkward() if set(self.common_keys) != set(self.rendered_form.columns()): actual_form = self.rendered_form.select_columns(self.common_keys) @@ -774,22 +777,43 @@ def __call__(self, i_start_stop): self.ttrees[i], start, stop, self.interp_options ) - return awkward.from_buffers( + layout = awkward.from_buffers( actual_form, stop - start, mapping, buffer_key=buffer_key, + highlevel=False, + ) + + return awkward.Array( + dask_awkward.lib.unproject_layout.unproject_layout( + self.rendered_form, + layout, + ), behavior=self.form_mapping.behavior, ) - return self.ttrees[i].arrays( + array = self.ttrees[i].arrays( self.common_keys, entry_start=start, entry_stop=stop, ak_add_doc=self.interp_options["ak_add_doc"], ) - def project_columns(self, common_keys): + if self.original_form is not None: + awkward = uproot.extras.awkward() + dask_awkward = uproot.extras.dask_awkward() + + return awkward.Array( + dask_awkward.lib.unproject_layout.unproject_layout( + self.original_form, + array.layout, + ) + ) + + return array + + def project_columns(self, common_keys=None, original_form=None): common_base_keys = self.common_base_keys if self.form_mapping is not None: awkward = uproot.extras.awkward() @@ -821,6 +845,7 @@ def project_columns(self, common_keys): self.interp_options, self.form_mapping, self.rendered_form, + original_form, ) @@ -835,6 +860,7 @@ def __init__( interp_options, form_mapping, rendered_form, + original_form=None, ) -> None: self.custom_classes = custom_classes self.allow_missing = allow_missing @@ -844,6 +870,7 @@ def __init__( self.interp_options = interp_options self.form_mapping = form_mapping self.rendered_form = rendered_form + self.original_form = original_form def __call__(self, file_path_object_path_istep_nsteps_ischunk): ( @@ -882,6 +909,7 @@ def __call__(self, file_path_object_path_istep_nsteps_ischunk): if self.form_mapping is not None: awkward = uproot.extras.awkward() + dask_awkward = uproot.extras.dask_awkward() if set(self.common_keys) != set(self.rendered_form.columns()): actual_form = self.rendered_form.select_columns(self.common_keys) @@ -892,22 +920,43 @@ def __call__(self, file_path_object_path_istep_nsteps_ischunk): ttree, start, stop, self.interp_options ) - return awkward.from_buffers( + layout = awkward.from_buffers( actual_form, stop - start, mapping, buffer_key=buffer_key, + highlevel=False, + ) + + return awkward.Array( + dask_awkward.lib.unproject_layout.unproject_layout( + self.rendered_form, + layout, + ), behavior=self.form_mapping.behavior, ) - return ttree.arrays( + array = ttree.arrays( self.common_keys, entry_start=start, entry_stop=stop, ak_add_doc=self.interp_options["ak_add_doc"], ) - def project_columns(self, common_keys): + if self.original_form is not None: + awkward = uproot.extras.awkward() + dask_awkward = uproot.extras.dask_awkward() + + return awkward.Array( + dask_awkward.lib.unproject_layout.unproject_layout( + self.original_form, + array.layout, + ) + ) + + return array + + def project_columns(self, columns=None, original_form=None): common_base_keys = self.common_base_keys if self.form_mapping is not None: awkward = uproot.extras.awkward() @@ -918,8 +967,8 @@ def project_columns(self, common_keys): ) = awkward._nplikes.typetracer.typetracer_with_report(self.rendered_form) tt = awkward.Array(new_meta_labelled) - if common_keys is not None: - for key in common_keys: + if columns is not None: + for key in columns: tt[tuple(key.split("."))].layout._touch_data(recursive=True) common_base_keys = [ @@ -930,18 +979,19 @@ def project_columns(self, common_keys): if x in self.common_base_keys ] - elif common_keys is not None: - common_keys = [x for x in common_keys if x in self.common_keys] + elif columns is not None: + columns = [x for x in columns if x in self.common_keys] return _UprootOpenAndRead( self.custom_classes, self.allow_missing, self.real_options, - common_keys, - common_keys if self.form_mapping is None else common_base_keys, + columns, + columns if self.form_mapping is None else common_base_keys, self.interp_options, self.form_mapping, self.rendered_form, + original_form=original_form, )