Exporting Python UDFs Using the Python Client (#26)

* Add support for locally debugging Python UDFs.

* Exporting UDFs as py and bin files to be imported in IDEs

* Fixing problem with udf export

* Removing traces option from udf debugging

* Exporting input data as dict

* Correcting Functions Name to meet the paper

* Adding Sampling option

* Fix some PEP8 violations.

* Working testcase and additional formatting.

* More PEP8 violations.

* Fixing bug

Author: pdet

pymonetdb/sql/cursors.py

@@ -6,6 +6,11 @@
 
 import logging
 from collections import namedtuple
+import tempfile
+import re
+import pickle
+import pdb
+
 from pymonetdb.sql import monetize, pythonize
 from pymonetdb.exceptions import ProgrammingError, InterfaceError
 from pymonetdb import mapi
@@ -24,6 +29,7 @@ class Cursor(object):
     connection are not isolated, i.e., any changes done to the
     database by a cursor are immediately visible by the other
     cursors"""
+
     def __init__(self, connection):
         """This read-only attribute return a reference to the Connection
         object on which the cursor was created."""
@@ -153,10 +159,10 @@ def execute(self, operation, parameters=None):
         if parameters:
             if isinstance(parameters, dict):
                 query = operation % dict([(k, monetize.convert(v))
-                                         for (k, v) in parameters.items()])
+                                          for (k, v) in parameters.items()])
             elif type(parameters) == list or type(parameters) == tuple:
-                query = operation % tuple([monetize.convert(item) for item
-                                           in parameters])
+                query = operation % tuple(
+                    [monetize.convert(item) for item in parameters])
             elif isinstance(parameters, str):
                 query = operation % monetize.convert(parameters)
             else:
@@ -185,6 +191,244 @@ def executemany(self, operation, seq_of_parameters):
         self.rowcount = count
         return count
 
+    def __exportparameters(self, ftype, fname, query, quantity_parameters,
+                           sample):
+        """ Exports the input parameters of a given UDF execution
+            to the Python process. Used internally for .debug() and
+            .export() functions.
+        """
+
+        # create a dummy function that only exports its parameters
+        # using the pickle module
+        if ftype == 5:
+            # table producing function
+            return_type = "TABLE(s STRING)"
+        else:
+            return_type = "STRING"
+        if sample == -1:
+            export_function = """
+                CREATE OR REPLACE FUNCTION export_parameters(*)
+                RETURNS %s LANGUAGE PYTHON
+                {
+                    import inspect
+                    import pickle
+                    frame = inspect.currentframe();
+                    args, _, _, values = inspect.getargvalues(frame);
+                    dd = {x: values[x] for x in args};
+                    del dd['_conn']
+                    return pickle.dumps(dd);
+                };""" % return_type
+        else:
+            export_function = """
+                CREATE OR REPLACE FUNCTION export_parameters(*)
+                RETURNS %s LANGUAGE PYTHON
+                {
+                import inspect
+                import pickle
+                import numpy
+                frame = inspect.currentframe();
+                args, _, _, values = inspect.getargvalues(frame);
+                dd = {x: values[x] for x in args};
+                del dd['_conn']
+                result = dict()
+                argname = "arg1"
+                x = numpy.arange(len(dd[argname]))
+                x = numpy.random.choice(x,%s,replace=False)
+                for i in range(len(dd)-2):
+                    argname = "arg" + str(i + 1)
+                    result = dd[argname]
+                    aux = []
+                    for j in range(len(x)):
+                        aux.append(result[x[j]])
+                    dd[argname] = aux
+                    print(dd[argname])
+                print(x)
+                return pickle.dumps(dd);
+                };
+                """ % (return_type, str(sample))
+
+        if fname not in query:
+            raise Exception("Function %s not found in query!" % fname)
+
+        query = query.replace(fname, 'export_parameters')
+        query = query.replace(';', ' sample 1;')
+
+        self.execute(export_function)
+        self.execute(query)
+        input_data = self.fetchall()
+        self.execute('DROP FUNCTION export_parameters;')
+        if len(input_data) <= 0:
+            raise Exception("Could not load input data!")
+        arguments = pickle.loads(str(input_data[0][0]))
+
+        if len(arguments) != quantity_parameters + 2:
+            raise Exception("Incorrect amount of input arguments found!")
+
+        return arguments
+
+    def export(self, query, fname, sample=-1, filespath='./'):
+        """ Exports a Python UDF and its input parameters to a given
+            file so it can be called locally in an IDE environment.
+        """
+
+        # first retrieve UDF information from the server
+        self.execute("""
+            SELECT func,type
+            FROM functions
+            WHERE language >= 6 AND language <= 11 AND name='%s';""" % fname)
+        data = self.fetchall()
+        self.execute("""
+            SELECT args.name
+            FROM args INNER JOIN functions ON args.func_id=functions.id
+            WHERE functions.name='%s' AND args.inout=1
+            ORDER BY args.number;""" % fname)
+        input_names = self.fetchall()
+        quantity_parameters = len(input_names)
+        fcode = data[0][0]
+        ftype = data[0][1]
+        parameter_list = []
+        # exporting Python UDF Function
+        if len(data) == 0:
+            raise Exception("Function not found!")
+        else:
+            parameters = '('
+            for x in range(0, len(input_names)):
+                parameter = str(input_names[x]).split('\'')
+                if x < len(input_names) - 1:
+                    parameter_list.append(parameter[1])
+                    parameters = parameters + parameter[1] + ','
+                else:
+                    parameter_list.append(parameter[1])
+                    parameters = parameters + parameter[1] + '): \n'
+
+            data = str(data[0]).replace('\\t', '\t').split('\\n')
+
+            python_udf = 'import pickle \n \n \ndef ' + fname + parameters
+            for x in range(1, len(data) - 1):
+                python_udf = python_udf + '\t' + str(data[x]) + '\n'
+
+        # exporting Columns as Binary Files
+        arguments = self.__exportparameters(ftype, fname, query,
+                                            quantity_parameters, sample)
+        result = dict()
+        for i in range(len(arguments) - 2):
+            argname = "arg%d" % (i + 1)
+            result[parameter_list[i]] = arguments[argname]
+        pickle.dump(result, open(filespath + 'input_data.bin', 'wb'))
+
+        # loading Columns in Python & Call Function
+        python_udf += '\n' + 'input_parameters = pickle.load(open(\''
+        python_udf += filespath + 'input_data.bin\',\'rb\'))' + '\n'
+        python_udf += fname + '('
+        for i in range(0, quantity_parameters):
+            if i < quantity_parameters - 1:
+                python_udf += 'input_parameters[\''
+                python_udf += parameter_list[i] + '\'],'
+            else:
+                python_udf += 'input_parameters[\''
+                python_udf += parameter_list[i] + '\'])'
+
+        file = open(filespath + fname + '.py', 'w')
+        file.write(python_udf)
+        file.close()
+
+    def debug(self, query, fname, sample=-1):
+        """ Locally debug a given Python UDF function in a SQL query
+            using the PDB debugger. Optionally can run on only a
+            sample of the input data, for faster data export.
+        """
+
+        # first gather information about the function
+        self.execute("""
+            SELECT func, type
+            FROM functions
+            WHERE language>=6 AND language <= 11 AND name='%s';""" % fname)
+        data = self.fetchall()
+        if len(data) == 0:
+            raise Exception("Function not found!")
+
+        # then gather the input arguments of the function
+        self.execute("""
+            SELECT args.name, args.type
+            FROM args
+            INNER JOIN functions ON args.func_id=functions.id
+            WHERE functions.name='%s' AND args.inout=1
+            ORDER BY args.number;""" % fname)
+        input_types = self.fetchall()
+
+        fcode = data[0][0]
+        ftype = data[0][1]
+
+        # now obtain the input columns
+        arguments = self.__exportparameters(ftype, fname, query,
+                                            len(input_types), sample)
+
+        arglist = "_columns, _column_types, _conn"
+        cleaned_arguments = dict()
+        for i in range(len(input_types)):
+            argname = "arg%d" % (i + 1)
+            if argname not in arguments:
+                raise Exception("Argument %d not found!" % (i + 1))
+            input_name = str(input_types[i][0])
+            cleaned_arguments[input_name] = arguments[argname]
+            arglist += ", %s" % input_name
+        cleaned_arguments['_columns'] = arguments['_columns']
+        cleaned_arguments['_column_types'] = arguments['_column_types']
+
+        # create a temporary file for the function execution and run it
+        with tempfile.NamedTemporaryFile() as f:
+            fcode = fcode.strip()
+            fcode = re.sub('^{', '', fcode)
+            fcode = re.sub('};$', '', fcode)
+            fcode = re.sub('^\n', '', fcode)
+            function_definition = "def pyfun(%s):\n %s\n" % (
+                arglist, fcode.replace("\n", "\n "))
+            f.write(function_definition)
+            f.flush()
+            execfile(f.name, globals(), locals())
+
+            class LoopbackObject(object):
+                def __init__(self, connection):
+                    self.__conn = connection
+
+                def execute(self, query):
+                    self.__conn.execute("""
+                        CREATE OR REPLACE FUNCTION export_parameters(*)
+                        RETURNS TABLE(s STRING) LANGUAGE PYTHON
+                        {
+                            import inspect
+                            import pickle
+                            frame = inspect.currentframe();
+                            args, _, _, values = inspect.getargvalues(frame);
+                            dd = {x: values[x] for x in args};
+                            del dd['_conn']
+                            del dd['_columns']
+                            del dd['_column_types']
+                            return pickle.dumps(dd);
+                        };""")
+                    self.__conn.execute("""
+                        SELECT *
+                        FROM (%s) AS xx
+                        LIMIT 1""" % query)
+                    query_description = self.__conn.description
+                    self.__conn.execute("""
+                        SELECT *
+                        FROM export_parameters ( (%s) );""" % query)
+                    data = self.__conn.fetchall()
+                    arguments = pickle.loads(str(data[0][0]))
+                    self.__conn.execute('DROP FUNCTION export_parameters;')
+                    if len(arguments) != len(query_description):
+                        raise Exception("Incorrect number of parameters!")
+                    result = dict()
+                    for i in range(len(arguments)):
+                        argname = "arg%d" % (i + 1)
+                        result[query_description[i][0]] = arguments[argname]
+                    return result
+
+            cleaned_arguments['_conn'] = LoopbackObject(self)
+            pdb.set_trace()
+            return locals()['pyfun'](*[], **cleaned_arguments)
+
     def fetchone(self):
         """Fetch the next row of a query result set, returning a
         single sequence, or None when no more data is available."""

test/test_capabilities.py

@@ -15,6 +15,7 @@
 import os
 from time import time
 import unittest
+import tempfile
 
 from pymonetdb.exceptions import ProgrammingError
 import pymonetdb.sql
@@ -366,7 +367,7 @@ def test_bigresult(self):
         self.cursor.execute('select count(*) from types')
         r = self.cursor.fetchone()
         n = r[0]
-        self.cursor.arraysize=100000
+        self.cursor.arraysize = 100000
         self.cursor.execute('select * from types, types')
         r = self.cursor.fetchall()
         self.assertEqual(len(r), n**2)
@@ -393,3 +394,32 @@ def test_multiple_queries(self):
                             (table1, table2))
         result = self.cursor.fetchall()
         self.assertEqual(result, [(50, 50)])
+
+    def test_debug_udf(self):
+        self.cursor.execute("""
+            CREATE FUNCTION test_python_udf(i INTEGER)
+            RETURNS INTEGER
+            LANGUAGE PYTHON {
+                return i * 2;
+            };
+            """)
+        # test if Python UDFs are enabled on the server
+        try:
+            self.cursor.execute('SELECT test_python_udf(1);')
+        except:
+            # python UDFs are disabled or not compiled in, skip this test
+            return
+        result = self.cursor.fetchall()
+        self.assertEqual(result, [(2,)])
+        # test python debugging capabilities
+        with tempfile.NamedTemporaryFile(delete=False) as f:
+            fname = f.name
+        self.cursor.export('SELECT test_python_udf(1)', 'test_python_udf', filespath=fname)
+        fname += "test_python_udf.py"
+        with open(fname) as f:
+            code = f.read()
+            self.assertEqual('test_python_udf(i)' in code, True)
+
+
+
+