Changes for more generic query support, uproot-raw

ponyisi · ponyisi · commit 054776df6f89 · 2024-04-26T03:11:28.000Z
diff --git a/servicex/databinder_models.py b/servicex/databinder_models.py
@@ -30,7 +30,7 @@
 from pydantic import BaseModel, Field, root_validator, constr, validator
 
 from servicex.dataset_identifier import RucioDatasetIdentifier, FileListDataset
-from servicex.func_adl import func_adl_dataset
+from servicex.query import Query as SXQuery, QueryStringGenerator
 
 
 class Sample(BaseModel):
@@ -40,7 +40,7 @@ class Sample(BaseModel):
     XRootDFiles: Optional[Union[str, List[str]]]
     NFiles: Optional[int] = Field(default=None)
     Function: Optional[Union[str, Callable]] = Field(default=None)
-    Query: Optional[Union[str, func_adl_dataset.Query]] = Field(default=None)
+    Query: Optional[Union[str, SXQuery, QueryStringGenerator]] = Field(default=None)
     Tree: Optional[str] = Field(default=None)
     IgnoreLocalCache: bool = False
 
diff --git a/servicex/query.py b/servicex/query.py
@@ -492,3 +492,32 @@ async def as_signed_urls_async(self, display_progress: bool = True,
                                                   dataset_group=dataset_group)
 
     as_signed_urls = make_sync(as_signed_urls_async)
+
+
+class QueryStringGenerator(ABC):
+    '''This abstract class just defines an interface to give the selection string'''
+    @abc.abstractmethod    
+    def generate_selection_string(self) -> str:
+        pass
+
+
+class GenericQueryStringGenerator(QueryStringGenerator):
+    '''Return the string from the initializer'''
+    def __init__(self, query: str):
+        self.query = query
+
+    def generate_selection_string(self) -> str:
+        return self.query
+
+
+class GenericQuery(Query):
+    '''
+    This class gives a "generic" Query object which doesn't require 
+    overloading the constructor
+    '''
+    query_string_generator: Optional[QueryStringGenerator] = None
+
+    def generate_selection_string(self) -> str:
+        if self.query_string_generator is None:
+            raise RuntimeError('query string generator not set')
+        return self.query_string_generator.generate_selection_string()
diff --git a/servicex/servicex_client.py b/servicex/servicex_client.py
@@ -25,13 +25,14 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-from typing import Optional, List, TypeVar, Any, Type
+from typing import Optional, List, TypeVar, Any, Type, Union
 
 from servicex.configuration import Configuration
 from servicex.func_adl.func_adl_dataset import FuncADLQuery
 from servicex.models import ResultFormat, TransformStatus
 from servicex.query_cache import QueryCache
 from servicex.servicex_adapter import ServiceXAdapter
+from servicex.query import Query, GenericQuery, QueryStringGenerator, GenericQueryStringGenerator
 from servicex.types import DID
 from servicex.python_dataset import PythonQuery
 from servicex.dataset_group import DatasetGroup
@@ -57,21 +58,26 @@ def get_codegen(_sample: Sample, _general: General):
     datasets = []
     for sample in config.Sample:
         if sample.Query:
-            if type(sample.Query) is str:
-                qastle_query = qastle.python_ast_to_text_ast(ast.parse(sample.Query)) # NOQA E501
-                sample.Query = FuncADLQuery()
+            # if string or QueryStringGenerator, turn into a Query
+            if isinstance(sample.Query, str) or isinstance(sample.Query, QueryStringGenerator):
+                sample.Query = sx.generic_query(dataset_identifier=sample.dataset_identifier, 
+                                    title=sample.Name,
+                                     codegen=get_codegen(sample, config.General),
+                                     result_format=config.General.OutputFormat,
+                                     ignore_cache=sample.IgnoreLocalCache,
+                                     query=sample.Query)
+            # query._q_ast = sample.Query._q_ast
+            # query._item_type = sample.Query._item_type
+            if isinstance(sample.Query, FuncADLQuery):
+                query = sx.func_adl_dataset(sample.dataset_identifier, sample.Name,
+                                            get_codegen(sample, config.General),
+                                            config.General.OutputFormat)
+                query._q_ast = sample.Query._q_ast
+                query._item_type = sample.Query._item_type
+                if sample.Tree:
+                    query = query.set_tree(sample.Tree)
+                sample.Query = query
 
-                sample.Query.set_provided_qastle(qastle_query)
-
-            query = sx.func_adl_dataset(sample.dataset_identifier, sample.Name,
-                                        get_codegen(sample, config.General),
-                                        config.General.OutputFormat)
-            query._q_ast = sample.Query._q_ast
-            query._item_type = sample.Query._item_type
-            if sample.Tree:
-                query = query.set_tree(sample.Tree)
-
-            sample.Query = query
             sample.Query.ignore_cache = sample.IgnoreLocalCache
 
             datasets.append(sample.Query)
@@ -258,3 +264,49 @@ def python_dataset(
             result_format=result_format,
             ignore_cache=ignore_cache
         )
+
+    def generic_query(
+        self,
+        dataset_identifier: DID,
+        codegen: str,
+        query: Union[str, QueryStringGenerator],
+        title: str = "ServiceX Client",
+        result_format: ResultFormat = ResultFormat.parquet,
+        ignore_cache: bool = False
+    ) -> Query:
+        r"""
+        Generate a Query object for a generic codegen specification
+
+        :param dataset_identifier:  The dataset identifier or filelist to be the source of files
+        :param title: Title to be applied to the transform. This is also useful for
+                      relating transform results.
+        :param codegen: Name of the code generator to use with this transform
+        :param result_format:  Do you want Paqrquet or Root? This can be set later with
+                               the set_result_format method
+        :param ignore_cache: Ignore the query cache and always run the query
+        :return: A Query object
+
+        """
+
+        if codegen not in self.code_generators:
+            raise NameError(
+                f"{codegen} code generator not supported by serviceX "
+                f"deployment at {self.servicex.url}"
+            )
+
+        if isinstance(query, str):
+            query = GenericQueryStringGenerator(query)
+        if not isinstance(query, QueryStringGenerator):
+            raise ValueError("query argument must be string or QueryStringGenerator")
+
+        qobj = GenericQuery(dataset_identifier=dataset_identifier,
+                            sx_adapter=self.servicex,
+                            title=title,
+                            codegen=codegen,
+                            config=self.config,
+                            query_cache=self.query_cache,
+                            result_format=result_format,
+                            ignore_cache=ignore_cache
+                            )
+        qobj.query_string_generator = query
+        return qobj
diff --git a/servicex/uproot_raw/uproot_raw.py b/servicex/uproot_raw/uproot_raw.py
@@ -0,0 +1,55 @@
+# Copyright (c) 2024, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+# pydantic 1.10 API
+
+import pydantic
+from typing import List, Union, Mapping, Optional
+from ..query import QueryStringGenerator
+
+class SubQuery(pydantic.BaseModel):
+    treename: Union[Mapping[str,str], List[str], str]
+    expressions: Optional[Union[List[str], str]]
+    cut: Optional[str]
+    filter_name: Optional[Union[List[str], str]]
+    filter_typename: Optional[Union[List[str], str]]
+    aliases: Optional[Mapping[str,str]]
+    copy_histogram: Optional[bool]
+
+
+@pydantic.dataclasses.dataclass
+class UprootRawQuery(QueryStringGenerator):
+    query: Union[List[SubQuery], SubQuery]
+
+    def generate_selection_string(self):
+        import json
+        if isinstance(self.query, SubQuery):
+            final_query = [self.query]
+        else:
+            final_query = self.query
+        return json.dumps([json.loads(_.json()) for _ in final_query])
