Combine awkward and panads separately

Not 100% sure this is going to work, but lets see!

Fixes #116

Author: gordonwatts

.vscode/settings.json

@@ -25,6 +25,7 @@
         "aexit",
         "aiohttp",
         "asyncio",
+        "awks",
         "backends",
         "cacheme",
         "codecov",

servicex/data_conversions.py

@@ -1,10 +1,14 @@
-from pathlib import Path
-from concurrent.futures import ThreadPoolExecutor
 import asyncio
-from servicex.utils import ServiceXException
-from typing import Dict, Optional, Union
-
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Dict, Iterable, Optional, Union
 from awkward.array.chunked import ChunkedArray
+from awkward.array.table import Table
+
+import pandas as pd
+import awkward as ak
+
+from servicex.utils import ServiceXException
 
 _conversion_pool = ThreadPoolExecutor(4)
 
@@ -53,6 +57,22 @@ async def convert_to_awkward(self, file: Path, file_type: Optional[str] = None):
             raise ServiceXException(f'Conversion from {file_type} into an awkward array is not '
                                     'yet supported')
 
+    def combine_pandas(self, dfs: Iterable[pd.DataFrame]) -> pd.DataFrame:
+        '''Combine many pandas dataframes into a single one, in order.
+
+        Args:
+            dfs (Iterable[pd.DataFrame]): The list of dataframes
+        '''
+        return pd.concat(dfs)
+
+    def combine_awkward(self, awks: Iterable[Union[Table, ChunkedArray]]) -> Table:
+        '''Combine many awkward arrays into a single one, in order.
+
+        Args:
+            awks (Iterable[ChunkedArray]): The input list of awkward arrays
+        '''
+        return ak.concatenate(awks)
+
     async def _convert_root_to_pandas(self, file: Path):
         '''
         Convert the contents of a ROOT file to pandas.

servicex/servicex.py

@@ -140,18 +140,14 @@ async def get_data_parquet_async(self, selection_query: str) -> List[Path]:
     @functools.wraps(ServiceXABC.get_data_pandas_df_async, updated=())
     @_wrap_in_memory_sx_cache
     async def get_data_pandas_df_async(self, selection_query: str):
-        import pandas as pd
-        return pd.concat(await self._data_return(
+        return self._converter.combine_pandas(await self._data_return(
             selection_query, lambda f: self._converter.convert_to_pandas(f)))
 
     @functools.wraps(ServiceXABC.get_data_awkward_async, updated=())
     @_wrap_in_memory_sx_cache
     async def get_data_awkward_async(self, selection_query: str):
-        import awkward
-        all_data = await self._data_return(
-            selection_query, lambda f: self._converter.convert_to_awkward(f))
-        col_names = all_data[0].keys()
-        return {c: awkward.concatenate([ar[c] for ar in all_data]) for c in col_names}
+        return self._converter.combine_awkward(await self._data_return(
+            selection_query, lambda f: self._converter.convert_to_awkward(f)))
 
     async def _file_return(self, selection_query: str, data_format: str):
         '''

tests/conftest.py

@@ -176,6 +176,7 @@ def good_pandas_file_data(mocker):
 
     converter = asyncmock.MagicMock(spec=DataConverterAdaptor)
     converter.convert_to_pandas.return_value = pd.DataFrame({'JetPt': [0, 1, 2, 3, 4, 5]})
+    converter.combine_pandas.return_value = converter.convert_to_pandas.return_value
 
     return converter
 
@@ -186,6 +187,7 @@ def good_awkward_file_data(mocker):
 
     converter = asyncmock.MagicMock(spec=DataConverterAdaptor)
     converter.convert_to_awkward.return_value = {'JetPt': awk.fromiter([0, 1, 2, 3, 4, 5])}
+    converter.combine_awkward.return_value = converter.convert_to_awkward.return_value
 
     return converter
 

tests/test_data_conversions.py

@@ -54,3 +54,63 @@ async def test_to_awkward_fail(good_root_file_path):
 async def test_to_panads_fail(good_root_file_path):
     with pytest.raises(ServiceXException):
         await DataConverterAdaptor('root').convert_to_pandas(good_root_file_path, 'notreally')
+
+
+def test_combine_pandas_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]].pandas.df()  # type: ignore
+            return df
+
+    df1 = load_df()
+    df2 = load_df()
+
+    combined = DataConverterAdaptor('root').combine_pandas([df1, df2])
+
+    assert len(combined) == len(df1) + len(df2)
+
+
+def test_combine_pandas_from_parquet(good_uproot_file_path):
+    'Load a dataframe from a parquet file and make sure they work when we ask them to combine'
+    def load_df():
+        import pandas as pd
+        return pd.read_parquet(good_uproot_file_path)
+
+    df1 = load_df()
+    df2 = load_df()
+
+    combined = DataConverterAdaptor('root').combine_pandas([df1, df2])
+
+    assert len(combined) == len(df1) + len(df2)
+
+
+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
+
+    df1 = load_df()
+    df2 = load_df()
+
+    combined = DataConverterAdaptor('root').combine_awkward([df1, df2])
+
+    assert len(combined) == len(df1) + len(df2)
+
+
+def test_combine_awkward_from_parquet(good_uproot_file_path):
+    'Load a dataframe from a parquet file and make sure they work when we ask them to combine'
+    def load_df():
+        import awkward as ak
+        return ak.fromparquet(good_uproot_file_path)
+
+    df1 = load_df()
+    df2 = load_df()
+
+    combined = DataConverterAdaptor('root').combine_awkward([df1, df2])
+
+    assert len(combined) == len(df1) + len(df2)

tests/test_servicex.py

@@ -262,6 +262,10 @@ async def test_good_run_single_ds_1file_pandas(mocker, good_pandas_file_data):
     assert isinstance(r, pd.DataFrame)
     assert len(r) == 6
 
+    good_pandas_file_data.combine_pandas.assert_called_once()
+    good_pandas_file_data.convert_to_pandas.assert_called_once()
+    assert len(good_pandas_file_data.combine_pandas.call_args[0][0]) == 1
+
 
 @pytest.mark.asyncio
 async def test_good_run_single_ds_1file_awkward(mocker, good_awkward_file_data):
@@ -283,6 +287,10 @@ async def test_good_run_single_ds_1file_awkward(mocker, good_awkward_file_data):
     assert 'JetPt' in r
     assert len(r['JetPt']) == 6
 
+    good_awkward_file_data.combine_awkward.assert_called_once()
+    good_awkward_file_data.convert_to_awkward.assert_called_once()
+    assert len(good_awkward_file_data.combine_awkward.call_args[0][0]) == 1
+
 
 @pytest.mark.asyncio
 async def test_good_run_single_ds_2file_pandas(mocker, good_pandas_file_data):
@@ -298,9 +306,9 @@ async def test_good_run_single_ds_2file_pandas(mocker, good_pandas_file_data):
                             cache_adaptor=mock_cache,
                             data_convert_adaptor=good_pandas_file_data,
                             local_log=mock_logger)
-    r = await ds.get_data_pandas_df_async('(valid qastle string)')
-    assert isinstance(r, pd.DataFrame)
-    assert len(r) == 6 * 2
+    await ds.get_data_pandas_df_async('(valid qastle string)')
+    good_pandas_file_data.combine_pandas.assert_called_once()
+    assert len(good_pandas_file_data.combine_pandas.call_args[0][0]) == 2
 
 
 @pytest.mark.asyncio
@@ -317,11 +325,8 @@ async def test_good_run_single_ds_2file_awkward(mocker, good_awkward_file_data):
                             cache_adaptor=mock_cache,
                             data_convert_adaptor=good_awkward_file_data,
                             local_log=mock_logger)
-    r = await ds.get_data_awkward_async('(valid qastle string)')
-    assert isinstance(r, dict)
-    assert len(r) == 1
-    assert 'JetPt' in r
-    assert len(r['JetPt']) == 6 * 2
+    await ds.get_data_awkward_async('(valid qastle string)')
+    assert len(good_awkward_file_data.combine_awkward.call_args[0][0]) == 2
 
 
 @pytest.mark.asyncio