Repository: iris-hep/func_adl

Issue:
Add AsParqetFiles
Make sure there is an end-point in `ObjectStream` that will render parquet files


<!-- Edit the body of your new issue then click the ✓ "Create Issue" button in the top right of the editor. The first line will be the issue title. Assignees and Labels follow after a blank line. Leave an empty line before beginning the body of the issue. -->

Discussion:

Patch:
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 14a1a78..1ac2548 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -9,10 +9,13 @@
         "Jupyter",
         "LINQ",
         "ROOTT",
+        "Topo",
         "asts",
         "astunparse",
         "asyncio",
         "codecov",
+        "databanks",
+        "dataframes",
         "elts",
         "gwatts",
         "iscoroutine",
@@ -21,6 +24,8 @@
         "noqa",
         "numpy",
         "pytest",
+        "rucio",
+        "servicex",
         "url's",
         "xaod",
         "xrootd"
@@ -41,8 +46,5 @@
     "python.analysis.openFilesOnly": false,
     "python.pythonPath": ".venv\\Scripts\\python.exe",
     "python.venvPath": ".venv",
-    "python.analysis.typeCheckingMode": "basic",
-    "python.testing.pytestArgs": [
-        "--no-cov"
-    ]
+    "python.analysis.typeCheckingMode": "basic"
 }
\ No newline at end of file
diff --git a/README.md b/README.md
index d579bd6..962cd06 100644
--- a/README.md
+++ b/README.md
@@ -16,3 +16,5 @@ one of the following packages:
 
 - func_adl.xAOD: for running on an ATLAS experiment xAOD file hosted in ServiceX
 - func_adl.xAOD.backend: for running on a local file using docker
+
+See the documentation for more information.
diff --git a/documentation.md b/documentation.md
index f555583..ad47ae0 100644
--- a/documentation.md
+++ b/documentation.md
@@ -1,10 +1,10 @@
 # Documentation for func_adl
 
-func_adl is an SQL-like language used to construct hierarchical data queries for `servicex` (https://pypi.org/project/servicex/). func_adl is written to work with both `servicex` service modes, and requires either `func_adl_xAOD` (https://pypi.org/project/func-adl-xAOD/) or `func_adl_uproot` (https://pypi.org/project/func-adl-uproot/) for full functionality.
+func_adl is an SQL-like language used to construct hierarchical data queries for [`servicex`](https://pypi.org/project/servicex/). func_adl is written to work with both `servicex` service modes, and requires either [`func_adl_xAOD`](https://pypi.org/project/func-adl-xAOD/) or [`func_adl_uproot`](https://pypi.org/project/func-adl-uproot/) for full functionality.
 
-### Example of a func_adl_xAOD query
+## Example of a func_adl_xAOD query
 
-```
+```python
 data = ServiceXDatasetSource(dataset) \
     .Select('lambda e: (e.Electrons("Electrons"), e.Muons("Muons"))') \
     .Select('lambda ls: (ls[0].Select(lambda e: e.pt()), ls[0].Select(lambda e: e.eta()), \
@@ -23,7 +23,7 @@ The following lines are each func_adl transformations on the dataset. func_adl f
 
 It should be noted that all commands except for `.value()` are used only to construct the query, and do not send any information or request anything from `ServiceX`. The entire query is sent using the `.value()` command, and there is no data exchange until that command is sent.
 
-### Commands in a func_adl_xAOD query:
+## Commands in a func_adl_xAOD query
 
 There are a variety of useful operators that can be submitted in a func_adl query. Here, they are listed and explained:
 
@@ -41,7 +41,10 @@ This command replaces the list with its length, counting the number of elements.
 
 `.AsAwkwardArray(...)`
 `.AsPandasDF(...)`
+`.AsROOTTTree(...)`
+`.AsParquetFile(...)`
 These commands specifies how the query should return the selection. Currently, the xAOD backend can return selections as either awkward arrays or pandas dataframes.
 
 `.value()`
-This command actually executes the query built by the other commands, sending it to `ServiceX` and then returning the data. 
\ No newline at end of file
+`.value_async()`
+This command actually executes the query built by the other commands, sending it to `ServiceX` and then returning the data.
diff --git a/func_adl/__init__.py b/func_adl/__init__.py
index 4392042..859e66e 100644
--- a/func_adl/__init__.py
+++ b/func_adl/__init__.py
@@ -6,7 +6,7 @@
 # just turn off checking here.
 
 # Main streaming object
-from .object_stream import ObjectStream, ObjectStreamException  # NOQA
+from .object_stream import ObjectStream  # NOQA
 
 # Dataset accessors
 from .event_dataset import EventDataset, find_EventDataset  # NOQA
diff --git a/func_adl/object_stream.py b/func_adl/object_stream.py
index ef44ece..a1ae73e 100644
--- a/func_adl/object_stream.py
+++ b/func_adl/object_stream.py
@@ -1,16 +1,16 @@
 # An Object stream represents a stream of objects, floats, integers, etc.
 import ast
-from typing import Any, Callable, Union, cast, Awaitable
+from typing import Any, Awaitable, Callable, List, Union, cast
 
 from make_it_sync import make_sync
 
 from .util_ast import as_ast, function_call, parse_as_ast
 
 
-class ObjectStreamException(Exception):
-    'Exception thrown by the ObjectStream object.'
-    def __init__(self, msg):
-        Exception.__init__(self, msg)
+# class ObjectStreamException(Exception):
+#     'Exception thrown by the ObjectStream object.'
+#     def __init__(self, msg):
+#         Exception.__init__(self, msg)
 
 
 class ObjectStream:
@@ -117,10 +117,35 @@ class ObjectStream:
 
             A new ObjectStream with type [(filename, treename)]. This is because multiple tree's
             may be written by the back end, and need to be concatenated together to get the full
+            dataset.  The order of the files back is consistent for different queries on the same
             dataset.
         """
         return ObjectStream(function_call("ResultTTree", [self._ast, as_ast(columns), as_ast(treename), as_ast(filename)]))
 
+    def AsParquetFiles(self, filename: str, columns: Union[str, List[str]] = []) -> 'ObjectStream':
+        '''Returns the sequence of items as a `parquet` file. Each item in the ObjectStream gets a separate
+        entry in the file. The times must be of types that the infrastructure can work with:
+
+            Float               A tree with a single float in each entry will be written.
+            vector<float>       A tree with a list of floats in each entry will be written.
+            (<tuple>)           A tree with multiple items (leaves) will be written. Each leaf
+                                must have one of the above types. Nested tuples are not supported.
+            {k:v, }             A dictionary with named columns. v is either a float or a vector of floats.
+
+        Arguments:
+
+            filename            Name of a file in which the data will be written. Depending on where the data comes
+                                from this may not be used - consider it a suggestion.
+            columns             If the data does not arrive by dictionary, then these are the column names.
+
+        Returns:
+
+            A new `ObjectStream` with type `[filename]`. This is because multiple files may be written by
+            the backend - the data should be concatinated together to get a final result. The order of the files back
+            is consistent for different queries on the same dataset.
+        '''
+        return ObjectStream(function_call("ResultParquet", [self._ast, as_ast(columns), as_ast(filename)]))
+
     def AsAwkwardArray(self, columns=[]) -> 'ObjectStream':
         r'''
         Return a pandas stream that contains one item, an `awkward` array, or dictionary of `awkward` arrays.


Test patch:
diff --git a/tests/test_object_stream.py b/tests/test_object_stream.py
index 2a76e78..9831fec 100644
--- a/tests/test_object_stream.py
+++ b/tests/test_object_stream.py
@@ -1,5 +1,6 @@
 # Test the object stream
 import sys
+from typing import Any, Optional
 sys.path += ['.']
 from func_adl import EventDataset
 import ast
@@ -29,6 +30,16 @@ def test_simple_query():
         .value()
     assert isinstance(r, ast.AST)
 
+
+def test_simple_query_parquet():
+    r = my_event() \
+        .SelectMany("lambda e: e.jets()") \
+        .Select("lambda j: j.pT()") \
+        .AsParquetFiles("junk.root", 'jetPT') \
+        .value()
+    assert isinstance(r, ast.AST)
+
+
 def test_simple_query_panda():
     r = my_event() \
         .SelectMany("lambda e: e.jets()") \
@@ -88,3 +99,22 @@ async def test_2await_exe_from_coroutine():
     rpair = await asyncio.gather(r1, r2)
     assert isinstance(rpair[0], ast.AST)
     assert isinstance(rpair[1], ast.AST)
+
+
+@pytest.mark.asyncio
+async def test_passed_in_executor():
+    logged_ast: Optional[ast.AST] = None
+    
+    async def my_exe(a: ast.AST) -> Any:
+        nonlocal logged_ast
+        logged_ast = a
+        return 1
+        
+    r = my_event() \
+        .SelectMany("lambda e: e.jets()") \
+        .Select("lambda j: j.pT()") \
+        .AsROOTTTree("junk.root", "analysis", "jetPT") \
+        .value_async(executor=my_exe)
+
+    assert (await r) == 1
+    assert logged_ast is not None