Combine awkward and panads separately (#117)

gordonwatts · web-flow · commit ce602b1ee1ff · 2020-09-24T13:04:03.000-07:00
Pull out how we do combining pandas and awkward - turns out to need more testing than we had been doing as there are some odd edge cases.
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -25,6 +25,7 @@
         "aexit",
         "aiohttp",
         "asyncio",
+        "awks",
         "backends",
         "cacheme",
         "codecov",
diff --git a/README.md b/README.md
@@ -59,7 +59,7 @@ The following lines will return a `pandas.DataFrame` containing all the jet pT's
     from servicex import ServiceX
     query = "(call ResultTTree (call Select (call SelectMany (call EventDataset (list 'localds:bogus')) (lambda (list e) (call (attr e 'Jets') 'AntiKt4EMTopoJets'))) (lambda (list j) (/ (call (attr j 'pt')) 1000.0))) (list 'JetPt') 'analysis' 'junk.root')"
     dataset = "mc15_13TeV:mc15_13TeV.361106.PowhegPythia8EvtGen_AZNLOCTEQ6L1_Zee.merge.DAOD_STDM3.e3601_s2576_s2132_r6630_r6264_p2363_tid05630052_00"
-    ds = ServiceXDataset(dataset, 'xaod')
+    ds = ServiceXDataset(dataset)
     r = ds.get_data_pandas_df(query)
     print(r)
 ```
@@ -162,8 +162,9 @@ Everything is based around the `ServiceXDataset` object. Below is the documentat
           dataset                     Name of a dataset from which queries will be selected.
           backend_type                The type of backend. Used only if we need to find an
                                       end-point. If we do not have a `servicex_adaptor` then this
-                                      cannot be null. Possible types are `uproot`, `xaod`,
-                                      and anything that finds a match in the `.servicex` file.
+                                      will default to xaod, unless you have any endpoint listed
+                                      in your servicex file. It will default to best match there,
+                                      in that case.
           image                       Name of transformer image to use to transform the data
           max_workers                 Maximum number of transformers to run simultaneously on
                                       ServiceX.
diff --git a/servicex/config_default.yaml b/servicex/config_default.yaml
@@ -1,11 +1,14 @@
 # Default settings for servicex. This will point you to a developer end-point, that
 # you've setup on your own machine (usually using k8's port-forward command):
 
-api_endpoint:
-  endpoint: http://localhost:5000
-  # email: xxx
-  # password: yyy
+api_endpoints:
+  - endpoint: http://localhost:5000
+    # email: xxx
+    # password: yyy
 
+# These are default settings, depreciated, and should not be used.
+# They will be removed in the next version.
+api_endpoint:
   minio_endpoint: localhost:9000
   # The username and password for accessing files generated by servicex.
   # NOTE:
diff --git a/servicex/data_conversions.py b/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.
@@ -125,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
@@ -135,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))
 
diff --git a/servicex/servicex.py b/servicex/servicex.py
@@ -2,7 +2,6 @@
 import asyncio
 import functools
 import logging
-from servicex.servicex_config import ServiceXConfigAdaptor
 import time
 from datetime import timedelta
 from pathlib import Path
@@ -13,12 +12,14 @@
 import backoff
 from backoff import on_exception
 
+from servicex.servicex_config import ServiceXConfigAdaptor
+
 from .cache import Cache
 from .data_conversions import DataConverterAdaptor
 from .minio_adaptor import (MinioAdaptor, MinioAdaptorFactory,
                             find_new_bucket_files)
-from .servicex_adaptor import (ServiceXAdaptor, servicex_adaptor_factory,
-                               transform_status_stream, trap_servicex_failures)
+from .servicex_adaptor import (ServiceXAdaptor, transform_status_stream,
+                               trap_servicex_failures)
 from .servicex_utils import _wrap_in_memory_sx_cache
 from .servicexabc import ServiceXABC
 from .utils import (ServiceXException, ServiceXFailedFileTransform,
@@ -55,8 +56,9 @@ def __init__(self,
             dataset                     Name of a dataset from which queries will be selected.
             backend_type                The type of backend. Used only if we need to find an
                                         end-point. If we do not have a `servicex_adaptor` then this
-                                        cannot be null. Possible types are `uproot`, `xaod`,
-                                        and anything that finds a match in the `.servicex` file.
+                                        will default to xaod, unless you have any endpoint listed
+                                        in your servicex file. It will default to best match there,
+                                        in that case.
             image                       Name of transformer image to use to transform the data
             max_workers                 Maximum number of transformers to run simultaneously on
                                         ServiceX.
@@ -94,10 +96,6 @@ def __init__(self,
                              status_callback_factory,
                              )
 
-        # Make sure the arguments are reasonable
-        if backend_type is None and servicex_adaptor is None:
-            raise ServiceXException('Specify backend_type or servicex_adaptor')
-
         # Get the local settings
         config = config_adaptor if config_adaptor is not None \
             else ServiceXConfigAdaptor()
@@ -109,8 +107,8 @@ def __init__(self,
 
         if not servicex_adaptor:
             # Given servicex adaptor is none, this should be ok. Fixes type checkers
-            assert backend_type is not None
-            servicex_adaptor = servicex_adaptor_factory(config.settings, backend_type)
+            end_point, email, password = config.get_servicex_adaptor_config(backend_type)
+            servicex_adaptor = ServiceXAdaptor(end_point, email, password)
         self._servicex_adaptor = servicex_adaptor
 
         if not minio_adaptor:
@@ -142,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):
         '''
diff --git a/servicex/servicex_adaptor.py b/servicex/servicex_adaptor.py
@@ -4,7 +4,6 @@
 import logging
 
 import aiohttp
-from confuse import ConfigView
 from google.auth import jwt
 
 from .utils import (
@@ -249,35 +248,3 @@ async def trap_servicex_failures(stream: AsyncIterator[TransformTuple]) \
                                               f'processed: {processed}).')
 
         yield p
-
-
-def servicex_adaptor_factory(c: ConfigView, backend_type: str) -> ServiceXAdaptor:
-    '''Given a configuration and the backend, find an appropriate configuration
-    for us to grab and create a `servicex_adaptor`.
-
-    Args:
-        c (ConfigView): The config information loaded form files.
-        backend_type (str): The backend type we need to match
-
-    Returns:
-        [ServiceXAdaptor]: A servicex adaptor.
-    '''
-    # Find a list of all endpoints.
-    # It is an error if this is not specified somewhere.
-    endpoints = c['api_endpoints']
-    seen_types = []
-    for ep in endpoints:
-        if ep['type'].as_str_expanded() == backend_type:
-            endpoint = ep['endpoint'].as_str_expanded()
-            email = ep['email'].as_str_expanded() if 'email' in ep else None
-            password = ep['password'].as_str_expanded() if 'password' in ep \
-                else None
-
-            # We can default these to "None"
-            return ServiceXAdaptor(endpoint, email, password)
-        else:
-            seen_types.append(ep['type'].as_str_expanded())
-
-    # If we are here, we found no matching type.
-    raise ServiceXException(f'Unable to find type {backend_type} '
-                            f'in configuration. Saw: {", ".join(seen_types)}')
diff --git a/servicex/servicex_config.py b/servicex/servicex_config.py
@@ -1,4 +1,5 @@
-from typing import Optional, cast
+import logging
+from typing import Optional, Tuple, cast
 
 from confuse.core import ConfigView
 
@@ -70,11 +71,76 @@ def get_backend_info(self, backend_type: str, key: str) -> Optional[str]:
         def find_in_list(c, key) -> Optional[str]:
             if c.exists():
                 for ep in c:
-                    if ep['type'].as_str() == backend_type:
-                        if key in ep:
-                            return ep[key].as_str_expanded()
+                    if ep['type'].exists():
+                        if ep['type'].as_str() == backend_type:
+                            if key in ep:
+                                return ep[key].as_str_expanded()
             return None
 
         a = find_in_list(self._settings['api_endpoints'], key)
         a = find_in_list(self._settings['backend_types'], key) if a is None else a
         return a
+
+    def get_servicex_adaptor_config(self, backend_type: Optional[str] = None) -> \
+            Tuple[str, Optional[str], Optional[str]]:
+        '''Return the servicex (endpoint, username, email) from a given backend configuration.
+
+        Args:
+            backend_type (str): The backend name (like `xaod`) which we hopefully can find in the
+            `.servicex` file.
+
+        Returns:
+            Tuple[str, str, str]: The tuple of info to create a `ServiceXAdaptor`: end point,
+            email, and password.
+        '''
+        # Find a list of all endpoints.
+        # It is an error if this is not specified somewhere.
+        endpoints = self._settings['api_endpoints']
+
+        def extract_info(ep) -> Tuple[str, Optional[str], Optional[str]]:
+            endpoint = ep['endpoint'].as_str_expanded()
+            email = ep['email'].as_str_expanded() if 'email' in ep else None
+            password = ep['password'].as_str_expanded() if 'password' in ep \
+                else None
+
+            # We can default these to "None"
+            return (endpoint, email, password)  # type: ignore
+
+        # If we have a good name, look for exact match
+        if backend_type is not None:
+            for ep in endpoints:
+                if ep['type'].exists() and ep['type'].as_str_expanded() == backend_type:
+                    return extract_info(ep)
+
+        # See if one is unlabeled.
+        log = logging.getLogger(__name__)
+        for ep in endpoints:
+            if not ep['type'].exists():
+                if backend_type is None:
+                    log.warning('No backend type requested, '
+                                f'using {ep["endpoint"].as_str_expanded()} - please be explicit '
+                                'in the ServiceXDataset constructor')
+                else:
+                    log.warning(f"No '{backend_type}' backend type found, "
+                                f'using {ep["endpoint"].as_str_expanded()} - please add to '
+                                'the .servicex file')
+                return extract_info(ep)
+
+        if backend_type is not None:
+            # They have a labeled backend, and all the end-points are labeled. So that means
+            # there really is not match. So throw!
+            seen_types = [str(ep['type'].as_str_expanded()) for ep in endpoints
+                          if ep['type'].exists()]
+            raise ServiceXException(f'Unable to find type {backend_type} '
+                                    f'in configuration. Saw: {", ".join(seen_types)}')
+
+        # Nope - now we are going to have to just use the first one there.
+        for ep in endpoints:
+            log.warning('No backend type requested, '
+                        f'using {ep["endpoint"].as_str_expanded()} - please be explicit '
+                        'in the ServiceXDataset constructor')
+            return extract_info(ep)
+
+        # If we are here - then... it just isn't going to work!
+        raise ServiceXException('Not even a default set of configurations are here! Bad install '
+                                ' of the servicex package!')
diff --git a/tests/conftest.py b/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
 
diff --git a/tests/test_data_conversions.py b/tests/test_data_conversions.py
diff --git a/tests/test_servicex.py b/tests/test_servicex.py
diff --git a/tests/test_servicex_adaptor.py b/tests/test_servicex_adaptor.py
diff --git a/tests/test_servicex_config.py b/tests/test_servicex_config.py
