Getting parquet files working

- Normalize test names
- Add better tests and code for files.

Author: gordonwatts

.vscode/settings.json

@@ -36,6 +36,7 @@
         "getenv",
         "gitlab",
         "inmem",
+        "isabstractmethod",
         "jupyter",
         "jupyterlab",
         "leftfoot",

servicex/servicex.py

@@ -40,6 +40,14 @@ class StreamInfoUrl:
     bucket: str
 
 
+@dataclass
+class StreamInfoPath:
+    '''Contains information on accessing ServiceX data via a local Path
+    '''
+    path: Path
+    file: str
+
+
 class ServiceXDataset(ServiceXABC):
     '''
     Used to access an instance of ServiceX at an end point on the internet. Support convieration
@@ -159,11 +167,47 @@ def ignore_cache(self):
     async def get_data_rootfiles_async(self, selection_query: str) -> List[Path]:
         return await self._file_return(selection_query, 'root-file')
 
+    async def get_data_rootfiles_stream(self, selection_query: str) \
+            -> AsyncIterator[StreamInfoPath]:
+        '''Returns, as an async iterator, each completed batch of work from Servicex.
+        The `StreamInfoPath` contains a path where downstream consumers can directly
+        access the data.
+
+        Args:
+            selection_query (str): The `qastle` query for the data to retreive.
+
+        Yields:
+            AsyncIterator[StreamInfoPath]: As ServiceX completes the data, and it is downloaded
+                                           to the local machine, the async iterator returns
+                                           a `StreamInfoPath` which can be used to access the
+                                           file locally.
+        '''
+        async for f_info in self._stream_local_files(selection_query, 'root-files'):
+            yield f_info
+
     @functools.wraps(ServiceXABC.get_data_parquet_async, updated=())
     @_wrap_in_memory_sx_cache
     async def get_data_parquet_async(self, selection_query: str) -> List[Path]:
         return await self._file_return(selection_query, 'parquet')
 
+    async def get_data_parquet_stream(self, selection_query: str) \
+            -> AsyncIterator[StreamInfoPath]:
+        '''Returns, as an async iterator, each completed batch of work from Servicex.
+        The `StreamInfoPath` contains a path where downstream consumers can directly
+        access the data.
+
+        Args:
+            selection_query (str): The `qastle` query for the data to retreive.
+
+        Yields:
+            AsyncIterator[StreamInfoPath]: As ServiceX completes the data, and it is downloaded
+                                           to the local machine, the async iterator returns
+                                           a `StreamInfoPath` which can be used to access the
+                                           file locally.
+        '''
+        async for f_info in self._stream_local_files(selection_query, 'parquet'):
+            yield f_info
+
     @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):
@@ -185,18 +229,18 @@ async def get_data_rootfiles_url_stream(self, selection_query: str) \
         Args:
             selection_query (str): The ServiceX Selection
         '''
-        async for f_info in self._get_minio_buckets(selection_query, 'root-files'):
+        async for f_info in self._stream_url_buckets(selection_query, 'root-files'):
             yield f_info
 
-    async def get_data_parquet_minio_stream(self, selection_query: str) \
+    async def get_data_parquet_url_stream(self, selection_query: str) \
             -> AsyncIterator[StreamInfoUrl]:
         '''Returns, as an async iterator, each of the files from the minio bucket,
         as the files are added there.
 
         Args:
             selection_query (str): The ServiceX Selection
         '''
-        async for f_info in self._get_minio_buckets(selection_query, 'parquet'):
+        async for f_info in self._stream_url_buckets(selection_query, 'parquet'):
             yield f_info
 
     async def _file_return(self, selection_query: str, data_format: str):
@@ -224,7 +268,7 @@ async def convert_to_file(f: Path) -> Path:
         return await self._data_return(selection_query, convert_to_file, data_format)
 
     @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
-    async def _get_minio_buckets(self, selection_query: str, data_format: str) \
+    async def _stream_url_buckets(self, selection_query: str, data_format: str) \
             -> AsyncIterator[StreamInfoUrl]:
         '''Get a list of files back for a request
 
@@ -280,12 +324,10 @@ async def _get_minio_buckets(self, selection_query: str, data_format: str) \
                 await self._servicex_adaptor.dump_query_errors(client, request_id)
                 raise ServiceXException(f'Failed to transform all files in {request_id}') from e
 
-    @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
     async def _data_return(self, selection_query: str,
                            converter: Callable[[Path], Awaitable[Any]],
-                           data_format: str = 'root-file'):
-        '''
-        Given a query, return the data, in a unique order, that hold
+                           data_format: str = 'root-file') -> List[Any]:
+        '''Given a query, return the data, in a unique order, that hold
         the data for the query.
 
         For certian types of exceptions, the queries will be repeated. For example,
@@ -303,25 +345,51 @@ async def _data_return(self, selection_query: str,
             data                Data converted to the "proper" format, depending
                                 on the converter call.
         '''
-        # Get a notifier to update anyone who wants to listen.
-        notifier = self._create_notifier()
+        as_data = ((f.file, asyncio.ensure_future(converter(f.path)))
+                   async for f in self._stream_local_files(selection_query, data_format))
 
-        # Get all the files
-        as_files = \
-            (f async for f in
-             self._get_files(selection_query, data_format, notifier))
+        all_data = {d[0]: await d[1] async for d in as_data}
 
         # Convert them to the proper format
-        as_data = ((f[0], asyncio.ensure_future(converter(await f[1])))
-                   async for f in as_files)
 
         # Finally, we need them in the proper order so we append them
         # all together
-        all_data = {f[0]: await f[1] async for f in as_data}
         ordered_data = [all_data[k] for k in sorted(all_data.keys())]
 
         return ordered_data
 
+    @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'):
+        '''
+        Given a query, return the data as a list of paths pointing to local files
+        that contain the results of the query. This is an async generator, and files
+        are returned as they arrive.
+
+        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.
+
+        Returns:
+
+            data                Data converted to the "proper" format, depending
+                                on the converter call.
+        '''
+        # Get a notifier to update anyone who wants to listen.
+        notifier = self._create_notifier()
+
+        # Get all the files
+        as_files = \
+            (f async for f in
+             self._get_files(selection_query, data_format, notifier))
+
+        async for name, a_path in as_files:
+            yield StreamInfoPath(Path(await a_path), name)
+
     async def _get_files(self, selection_query: str, data_type: str,
                          notifier: _status_update_wrapper) \
             -> AsyncIterator[Tuple[str, Awaitable[Path]]]:

tests/test_servicex.py

@@ -1,7 +1,7 @@
 import asyncio
 from contextlib import contextmanager
 from pathlib import Path
-from servicex.servicex import StreamInfoUrl
+from servicex.servicex import StreamInfoPath, StreamInfoUrl
 from servicex.cache import Cache
 
 from confuse.core import Configuration
@@ -141,7 +141,7 @@ async def test_minio_back(mocker):
         "/foo/bar.root")
 
     assert len(r) == 1
-    assert r[0] == '/foo/bar.root'
+    assert r[0] == Path('/foo/bar.root')
 
 
 @pytest.mark.asyncio
@@ -185,7 +185,7 @@ def test_good_run_root_files_no_async(mocker):
 
     r = ds.get_data_rootfiles('(valid qastle string)')
     assert len(r) == 2
-    assert r[0] == '/foo/bar.root'
+    assert r[0] == Path('/foo/bar.root')
 
 
 @pytest.mark.asyncio
@@ -392,7 +392,7 @@ async def test_good_run_single_ds_2file_awkward(mocker, good_awkward_file_data):
 
 
 @pytest.mark.asyncio
-async def test_good_run_root_files_from_minio(mocker):
+async def test_stream_root_files_from_minio(mocker):
     'Get a root 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")
@@ -420,7 +420,34 @@ async def test_good_run_root_files_from_minio(mocker):
 
 
 @pytest.mark.asyncio
-async def test_bad_request_id_run_root_files_from_minio(mocker):
+async def test_stream_root_files(mocker):
+    'Get a root 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)
+
+    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 async for f in ds.get_data_rootfiles_stream('(valid qastle string)')]
+
+    assert len(lst) == 1
+    r = lst[0]
+    assert isinstance(r, StreamInfoPath)
+    assert r.file == 'one_minio_entry'
+    assert 'foo' in r.path.parts
+    assert 'bar.root' in r.path.parts
+
+    assert mock_servicex_adaptor.query_json['result-format'] == 'root-file'
+
+
+@pytest.mark.asyncio
+async def test_stream_bad_request_id_run_root_files_from_minio(mocker):
     'Using the minio interface - the request_id is not known'
     mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
     transform_status = mocker.MagicMock(side_effect=ServiceXUnknownRequestID('boom'))
@@ -446,7 +473,7 @@ async def test_bad_request_id_run_root_files_from_minio(mocker):
 
 
 @pytest.mark.asyncio
-async def test_bad_transform_run_root_files_from_minio(mocker):
+async def test_stream_bad_transform_run_root_files_from_minio(mocker):
     'Using the async minio interface - fail to transform (like bad DID)'
     mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
     fatal_transform_status = {
@@ -489,7 +516,7 @@ async def test_bad_transform_run_root_files_from_minio(mocker):
 
 
 @pytest.mark.asyncio
-async def test_bad_file_transform_run_root_files_from_minio(mocker):
+async def test_stream_bad_file_transform_run_root_files_from_minio(mocker):
     'Using the async minio interface, some files will fail to translate.'
     mock_cache = build_cache_mock(mocker)
     mock_logger = mocker.MagicMock(spec=log_adaptor)
@@ -514,7 +541,7 @@ async def test_bad_file_transform_run_root_files_from_minio(mocker):
 
 
 @pytest.mark.asyncio
-async def test_good_run_parquet_files_from_minio(mocker):
+async def test_stream_parquet_files_from_minio(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")
@@ -528,9 +555,7 @@ async def test_good_run_parquet_files_from_minio(mocker):
                             cache_adaptor=mock_cache,
                             local_log=mock_logger,
                             data_convert_adaptor=data_adaptor)
-    lst = []
-    async for f_info in ds.get_data_parquet_minio_stream('(valid qastle string)'):
-        lst.append(f_info)
+    lst = [f_info async for f_info in ds.get_data_parquet_url_stream('(valid qastle string)')]
 
     assert len(lst) == 1
     assert lst[0].bucket == '123-456'
@@ -539,6 +564,31 @@ async def test_good_run_parquet_files_from_minio(mocker):
     assert mock_servicex_adaptor.query_json['result-format'] == 'parquet'
 
 
+@pytest.mark.asyncio
+async def test_stream_parquet_files(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)
+
+    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_parquet_stream('(valid qastle string)')]
+
+    assert len(lst) == 1
+    assert lst[0].file == 'one_minio_entry'
+    assert 'foo' in lst[0].path.parts
+    assert 'bar.root' in lst[0].path.parts
+
+    assert mock_servicex_adaptor.query_json['result-format'] == 'parquet'
+
+
 @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'
@@ -715,7 +765,7 @@ async def test_callback_none(mocker):
         "/foo/bar.root")
 
     assert len(r) == 1
-    assert r[0] == '/foo/bar.root'
+    assert r[0] == Path('/foo/bar.root')
 
 
 @pytest.mark.asyncio