Streaming tests done

- Added pandas and awkward streaming
- Tried to clean up and remove duplicate code

Author: gordonwatts

servicex/servicex.py

@@ -48,6 +48,14 @@ class StreamInfoPath:
     file: str
 
 
+@dataclass
+class StreamInfoData:
+    '''Contains information on accessing ServiceX data via converted data
+    '''
+    data: Any
+    file: str
+
+
 class ServiceXDataset(ServiceXABC):
     '''
     Used to access an instance of ServiceX at an end point on the internet. Support convieration
@@ -220,6 +228,16 @@ async def get_data_awkward_async(self, selection_query: str):
         return self._converter.combine_awkward(await self._data_return(
             selection_query, lambda f: self._converter.convert_to_awkward(f)))
 
+    async def get_data_awkward_stream(self, selection_query: str):
+        async for a in self._stream_return(selection_query,
+                                           lambda f: self._converter.convert_to_awkward(f)):
+            yield a
+
+    async def get_data_pandas_stream(self, selection_query: str):
+        async for a in self._stream_return(selection_query,
+                                           lambda f: self._converter.convert_to_pandas(f)):
+            yield a
+
     async def get_data_rootfiles_url_stream(self, selection_query: str) \
             -> AsyncIterator[StreamInfoUrl]:
         '''Returns, as an async iterator, each completed batch of work from ServiceX.
@@ -345,10 +363,10 @@ async def _data_return(self, selection_query: str,
             data                Data converted to the "proper" format, depending
                                 on the converter call.
         '''
-        as_data = ((f.file, asyncio.ensure_future(converter(f.path)))
-                   async for f in self._stream_local_files(selection_query, data_format))
-
-        all_data = {d[0]: await d[1] async for d in as_data}
+        all_data = {
+            f.file: f.data
+            async for f in self._stream_return(selection_query, converter, data_format)
+        }
 
         # Convert them to the proper format
 
@@ -358,6 +376,33 @@ async def _data_return(self, selection_query: str,
 
         return ordered_data
 
+    async def _stream_return(self, selection_query: str,
+                             converter: Callable[[Path], Awaitable[Any]],
+                             data_format: str = 'root-file') -> AsyncIterator[StreamInfoData]:
+        '''Given a query, return the data, in the order it arrives back
+        converted as appropriate.
+
+        For certian types of exceptions, the queries will be repeated. For example,
+        if `ServiceX` indicates that it was restarted in the middle of the query, then
+        the query will be re-submitted.
+
+        Arguments:
+
+            selection_query     `qastle` data that makes up the selection request.
+            converter           A `Callable` that will convert the data returned from
+                                `ServiceX` as a set of files.
+
+        Returns:
+
+            data                Data converted to the "proper" format, depending
+                                on the converter call.
+        '''
+        as_data = (StreamInfoData(await asyncio.ensure_future(converter(f.path)), f.file)
+                   async for f in self._stream_local_files(selection_query, data_format))
+
+        async for r in as_data:
+            yield r
+
     @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
     async def _stream_local_files(self, selection_query: str,
                                   data_format: str = 'root-file'):

tests/test_servicex.py

@@ -589,6 +589,52 @@ async def test_stream_parquet_files(mocker):
     assert mock_servicex_adaptor.query_json['result-format'] == 'parquet'
 
 
+@pytest.mark.asyncio
+async def test_stream_awkward_data(mocker):
+    'Get a parquet file pulling back minio info as it arrives'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+    data_adaptor.convert_to_awkward.return_value = {'JetPt': 10}
+
+    ds = fe.ServiceXDataset('localds://mc16_tev:13',
+                            servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+                            minio_adaptor=mock_minio_adaptor,  # type: ignore
+                            cache_adaptor=mock_cache,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+    lst = [f_info async for f_info in ds.get_data_awkward_stream('(valid qastle string)')]
+
+    assert len(lst) == 1
+    assert lst[0].file == 'one_minio_entry'
+    assert isinstance(lst[0].data, dict)
+
+
+@pytest.mark.asyncio
+async def test_stream_pandas_data(mocker):
+    'Get a parquet file pulling back minio info as it arrives'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+    data_adaptor.convert_to_pandas.return_value = {'JetPt': 10}
+
+    ds = fe.ServiceXDataset('localds://mc16_tev:13',
+                            servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+                            minio_adaptor=mock_minio_adaptor,  # type: ignore
+                            cache_adaptor=mock_cache,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+    lst = [f_info async for f_info in ds.get_data_pandas_stream('(valid qastle string)')]
+
+    assert len(lst) == 1
+    assert lst[0].file == 'one_minio_entry'
+    assert isinstance(lst[0].data, dict)
+
+
 @pytest.mark.asyncio
 async def test_status_exception(mocker):
     'Make sure status error - like transform not found - is reported all the way to the top'