Merge pull request #7 from citrusleaf/query_changes

Lua parameters to a query are a list

Author: rbotzer

src/main/query/apply.c

@@ -26,28 +26,35 @@
 
 AerospikeQuery * AerospikeQuery_Apply(AerospikeQuery * self, PyObject * args, PyObject * kwds)
 {
-	as_error err;
 
-	// Initialize error
+	// Python function arguments
+	PyObject * py_module = NULL;
+	PyObject * py_function = NULL;
+	PyObject * py_args = NULL;
+	PyObject * py_policy = NULL;
+
+	// Python function keyword arguments
+	static char * kwlist[] = {"module", "function", "arguments", "policy", NULL};
+
+	if ( PyArg_ParseTupleAndKeywords(args, kwds, "OO|OO:apply", kwlist, &py_module, &py_function, &py_args, &py_policy) == false ){
+		return NULL;
+	}
+
+	// Aerospike error object
+	as_error err;
+	// Initialize error object
 	as_error_init(&err);
 
-	int nargs = (int) PyTuple_Size(args);
+	if ( !self || !self->client->as ){
+		as_error_update(&err, AEROSPIKE_ERR_PARAM, "Invalid query object");
+		goto CLEANUP;
+	}
 
 	// Aerospike API Arguments
 	char * module = NULL;
 	char * function = NULL;
 	as_arraylist * arglist = NULL;
 
-	// too few args
-	if ( nargs < 2 ) {
-		as_error_update(&err, AEROSPIKE_ERR_CLIENT, "udf module and function names are required.");
-		goto CLEANUP;
-	}
-
-	// Python Arguments
-	PyObject * py_module = PyTuple_GetItem(args, 0);
-	PyObject * py_function = PyTuple_GetItem(args, 1);
-
 	if ( PyString_Check(py_module) ) {
 		module = PyString_AsString(py_module);
 	}
@@ -64,10 +71,13 @@ AerospikeQuery * AerospikeQuery_Apply(AerospikeQuery * self, PyObject * args, Py
 		goto CLEANUP;
 	}
 
-	if ( nargs > 2 ) {
-		arglist = as_arraylist_new(nargs-2, 0);
-		for ( int i = 2; i < nargs; i++ ) {
-			PyObject * py_val = PyTuple_GetItem(args, i);
+	if ( py_args && PyList_Check(py_args) ){
+		Py_ssize_t size = PyList_Size(py_args);
+
+		arglist = as_arraylist_new(size, 0);
+
+		for ( int i = 0; i < size; i++ ) {
+			PyObject * py_val = PyList_GetItem(py_args, (Py_ssize_t)i);
 			as_val * val = NULL;
 			pyobject_to_val(&err, py_val, &val);
 			if ( err.code != AEROSPIKE_OK ) {

test/stream_example.lua

@@ -17,3 +17,40 @@ end
 function count_less()
     return stream : map(one) : reduce(add);
 end
+
+local function having_ge_threshold(bin_having, ge_threshold)
+    return function(rec)
+        if rec[bin_having] < ge_threshold then
+            return false
+        end
+        return true
+    end
+end
+
+local function count_bins(group_by_bin)
+  return function(group, rec)
+    if rec[group_by_bin] then
+      local bin_name = rec[group_by_bin]
+      group[bin_name] = (group[bin_name] or 0) + 1
+    end
+    return group
+  end
+end
+
+local function add_values(val1, val2)
+  return val1 + val2
+end
+
+local function reduce_groups(a, b)
+  return map.merge(a, b, add_values)
+end
+
+function group_count(stream, group_by_bin, bin_having, ge_threshold)
+  if bin_having and ge_threshold then
+    local myfilter = having_ge_threshold(bin_having, ge_threshold)
+    return stream : filter(myfilter) : aggregate(map{}, count_bins(group_by_bin)) : reduce(reduce_groups)
+  else
+    return stream : aggregate(map{}, count_bins(group_by_bin)) : reduce(reduce_groups)
+  end
+end
+

test/test_aggregate.py

@@ -348,3 +348,19 @@ def user_callback(value):
             query.foreach(user_callback)
 
         assert exception.value[0] == 1L
+
+    def test_aggregate_with_arguments_to_lua_function(self):
+        """
+            Invoke apply() with parameter's list for lua function.
+        """
+        query = self.client.query('test', 'demo')
+        query.where(p.between('test_age', 0, 5))
+        query.apply('stream_example', 'group_count', ["name", "addr"])
+
+        rec = []
+        def callback(value):
+            rec.append(value)
+
+        query.foreach(callback)
+        assert rec == [{u'name4': 1, u'name2': 1, u'name3': 1, u'name0': 1, u'name1': 1}]
+