Repository: iris-hep/qastle

Issue:
`Choose` node type
Described briefly in https://github.com/iris-hep/func_adl_uproot/issues/46. In query languages, it's easy to form cartesian products of sets, but it seems generally fairly painful to form combinations (all distinct subsets of a particular size). I propose to add a `Choose` node type that does this latter operation. This is a commonly needed feature in HEP analyses. The format in `qastle` would be `(Choose <source> <n>)`, where `source` is the original sequence and `n` is the number of elements in each combination. Elements are not repeated within each combination.

Comments:

Patch:
diff --git a/README.md b/README.md
index cbf062f..4b09741 100644
--- a/README.md
+++ b/README.md
@@ -79,6 +79,8 @@ All defined s-expressions are listed here, though this specification will be exp
   - OrderBy: `(OrderBy <source> <key_selector>)`
     - `key_selector` must be a `lambda` with one argument
   - CrossJoin: `(CrossJoin <first> <second>)`
+  - Choose: `(Choose <source> <n>)`
+    - `n` must be an integer
 
 
 ## Example
diff --git a/qastle/linq_util.py b/qastle/linq_util.py
index 9fb9251..017dcb5 100644
--- a/qastle/linq_util.py
+++ b/qastle/linq_util.py
@@ -82,6 +82,13 @@ class CrossJoin(ast.AST):
         self.second = second
 
 
+class Choose(ast.AST):
+    def __init__(self, source, n):
+        self._fields = ['source', 'n']
+        self.source = source
+        self.n = n
+
+
 linq_operator_names = ('Where',
                        'Select',
                        'SelectMany',
@@ -93,7 +100,8 @@ linq_operator_names = ('Where',
                        'Sum',
                        'Zip',
                        'OrderBy',
-                       'CrossJoin')
+                       'CrossJoin',
+                       'Choose')
 
 
 class InsertLINQNodesTransformer(ast.NodeTransformer):
@@ -191,6 +199,10 @@ class InsertLINQNodesTransformer(ast.NodeTransformer):
                 raise SyntaxError('CrossJoin() call must have exactly one argument')
             return CrossJoin(first=self.visit(source),
                              second=self.visit(args[0]))
+        elif function_name == 'Choose':
+            if len(args) != 1:
+                raise SyntaxError('Choose() call must have exactly one argument')
+            return Choose(source=self.visit(source), n=self.visit(args[0]))
         else:
             raise NameError('Unhandled LINQ operator: ' + function_name)
 
diff --git a/qastle/transform.py b/qastle/transform.py
index fe9a820..2026c8e 100644
--- a/qastle/transform.py
+++ b/qastle/transform.py
@@ -1,5 +1,5 @@
 from .linq_util import (Where, Select, SelectMany, First, Aggregate, Count, Max,
-                        Min, Sum, Zip, OrderBy, CrossJoin)
+                        Min, Sum, Zip, OrderBy, CrossJoin, Choose)
 from .ast_util import wrap_ast, unwrap_ast
 
 import lark
@@ -214,6 +214,11 @@ class PythonASTToTextASTTransformer(ast.NodeVisitor):
                                                self.visit(node.first),
                                                self.visit(node.second))
 
+    def visit_Choose(self, node):
+        return self.make_composite_node_string('Choose',
+                                               self.visit(node.source),
+                                               self.visit(node.n))
+
     def generic_visit(self, node):
         raise SyntaxError('Unsupported node type: ' + str(type(node)))
 
@@ -465,5 +470,10 @@ class TextASTToPythonASTTransformer(lark.Transformer):
                 raise SyntaxError('CrossJoin node must have two fields; found ' + str(len(fields)))
             return CrossJoin(first=fields[0], second=fields[1])
 
+        elif node_type == 'Choose':
+            if len(fields) != 2:
+                raise SyntaxError('Choose node must have two fields; found ' + str(len(fields)))
+            return Choose(source=fields[0], n=fields[1])
+
         else:
             raise SyntaxError('Unknown composite node type: ' + node_type)


Test patch:
diff --git a/tests/test_ast_language.py b/tests/test_ast_language.py
index 68aa6f8..fc2900d 100644
--- a/tests/test_ast_language.py
+++ b/tests/test_ast_language.py
@@ -230,3 +230,10 @@ def test_CrossJoin():
     crossjoin_node = CrossJoin(first=unwrap_ast(ast.parse('left')),
                                second=unwrap_ast(ast.parse('right')))
     assert_equivalent_python_ast_and_text_ast(wrap_ast(crossjoin_node), '(CrossJoin left right)')
+
+
+def test_Choose():
+    choose_node = Choose(source=unwrap_ast(ast.parse('data_source')),
+                         n=unwrap_ast(ast.parse('2')))
+    assert_equivalent_python_ast_and_text_ast(wrap_ast(choose_node),
+                                              '(Choose data_source 2)')
diff --git a/tests/test_linq_util.py b/tests/test_linq_util.py
index a4576d2..d363e29 100644
--- a/tests/test_linq_util.py
+++ b/tests/test_linq_util.py
@@ -300,3 +300,24 @@ def test_crossjoin_composite():
 def test_crossjoin_bad():
     with pytest.raises(SyntaxError):
         insert_linq_nodes(ast.parse('left.CrossJoin()'))
+
+
+def test_choose():
+    initial_ast = ast.parse("the_source.Choose(2)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Choose(source=unwrap_ast(ast.parse('the_source')),
+                                   n=unwrap_ast(ast.parse('2'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_choose_composite():
+    initial_ast = ast.parse("the_source.First().Choose(2)")
+    final_ast = insert_linq_nodes(initial_ast)
+    expected_ast = wrap_ast(Choose(source=First(source=unwrap_ast(ast.parse('the_source'))),
+                                   n=unwrap_ast(ast.parse('2'))))
+    assert_ast_nodes_are_equal(final_ast, expected_ast)
+
+
+def test_choose_bad():
+    with pytest.raises(SyntaxError):
+        insert_linq_nodes(ast.parse('the_source.Choose()'))