For lazyarrays, keep the file handle open

Author: gordonwatts

servicex/data_conversions.py

@@ -145,7 +145,9 @@ async def _convert_root_to_awkward(self, file: Path):
         Note:
 
             - Work is done on a second thread.
-            - Pandas is only imported if this is called.
+            - Awkward is only imported if this is called.
+            - A LazyArray is returned, so it isn't completely loaded into memory. That also means this
+              will leak filehandles - as that has to be left open.
 
         '''
         from numpy import ndarray
@@ -155,11 +157,8 @@ def do_the_work(file: Path) -> Dict[Union[str, bytes], Union[ndarray, JaggedArra
             import uproot
 
             f_in = uproot.open(file)
-            try:
-                r = f_in[f_in.keys()[0]]
-                return r.lazyarrays()  # type: ignore
-            finally:
-                f_in._context.source.close()
+            r = f_in[f_in.keys()[0]]
+            return r.lazyarrays()  # type: ignore
 
         return await asyncio.wrap_future(_conversion_pool.submit(do_the_work, file))
 

tests/test_data_conversions.py

@@ -4,12 +4,21 @@
 import pandas as pd
 
 
+def check_awkward_accessible(col):
+    'Check to make sure we can look at every item in column'
+    col.flatten()
+
+
+def check_pandas_accessible(col):
+    assert len(col.array) > 0
+
+
 @pytest.mark.asyncio
 async def test_root_to_pandas(good_root_file_path):
     df = await DataConverterAdaptor('root').convert_to_pandas(good_root_file_path)
     assert isinstance(df, pd.DataFrame)
     assert len(df) == 283458
-
+    check_pandas_accessible(df['JetPt'])
 
 @pytest.mark.asyncio
 async def test_root_to_pandas_default(good_root_file_path):
@@ -30,18 +39,21 @@ async def test_parquet_to_pandas(good_uproot_file_path):
     df = await DataConverterAdaptor('parquet').convert_to_pandas(good_uproot_file_path)
     assert isinstance(df, pd.DataFrame)
     assert len(df) == 115714
+    check_pandas_accessible(df['JetPT'])
 
 
 @pytest.mark.asyncio
 async def test_parquet_to_awkward(good_uproot_file_path):
     df = await DataConverterAdaptor('parquet').convert_to_awkward(good_uproot_file_path)
     assert len(df['JetPT']) == 115714
+    check_awkward_accessible(df['JetPT'])
 
 
 @pytest.mark.asyncio
 async def test_root_to_awkward(good_root_file_path):
     df = await DataConverterAdaptor('root').convert_to_awkward(good_root_file_path)
     assert len(df['JetPt']) == 283458
+    check_awkward_accessible(df['JetPt'])
 
 
 @pytest.mark.asyncio
@@ -70,6 +82,7 @@ def load_df():
     combined = DataConverterAdaptor('root').combine_pandas([df1, df2])
 
     assert len(combined) == len(df1) + len(df2)
+    check_pandas_accessible(combined['JetPt'])
 
 
 def test_combine_pandas_from_parquet(good_uproot_file_path):
@@ -84,22 +97,24 @@ def load_df():
     combined = DataConverterAdaptor('root').combine_pandas([df1, df2])
 
     assert len(combined) == len(df1) + len(df2)
+    check_pandas_accessible(combined['JetPT'])
 
 
 def test_combine_awkward_from_root(good_root_file_path):
     'Load a dataframe from root files and make sure that they work when we ask them to combine'
     def load_df():
         import uproot
-        with uproot.open(good_root_file_path) as f_in:
-            df = f_in[f_in.keys()[0]].lazyarrays()  # type: ignore
-            return df
+        f_in = uproot.open(good_root_file_path)
+        df = f_in[f_in.keys()[0]].lazyarrays()  # type: ignore
+        return df
 
     df1 = load_df()
     df2 = load_df()
 
     combined = DataConverterAdaptor('root').combine_awkward([df1, df2])
 
     assert len(combined) == len(df1) + len(df2)
+    check_awkward_accessible(combined['JetPt'])
 
 
 def test_combine_awkward_from_parquet(good_uproot_file_path):
@@ -114,3 +129,4 @@ def load_df():
     combined = DataConverterAdaptor('root').combine_awkward([df1, df2])
 
     assert len(combined) == len(df1) + len(df2)
+    check_awkward_accessible(combined['JetPT'])