Add python table writer to output structed query results (#2832)

* add python table writer to output structed query results

* fix pylint

* update

* update

* update

Author: typhoonzero

go/codegen/experimental/codegen_normal_stmt.go

@@ -24,11 +24,16 @@ var normalStmtStepTmpl = `
 def step_entry_{{.StepIndex}}():
     import runtime
     import runtime.dbapi
+    from runtime.dbapi import table_writer
+
     conn = runtime.dbapi.connect("{{.DataSource}}")
     stmt = """{{.Stmt}}"""
     if conn.is_query(stmt):
         rs = conn.query(stmt)
-        # write rs to stdout using protobuf table writer
+        tw = table_writer.ProtobufWriter(rs)
+        lines = tw.dump_strings()
+        for l in lines:
+            print(l)
     else:
         success = conn.execute(stmt)
         if not success:

go/proto/sqlflow.proto

@@ -90,7 +90,7 @@ message Session {
 //      2. `USE ...`, `DELETE ...`
 //      3. `SELECT ... TO TRAIN/PREDICT ...`
 message Request {
-    string sql = 1;		// The SQL statement to be executed.
+    string sql = 1;      // The SQL statement to be executed.
     Session session = 2;
 }
 

python/runtime/dbapi/mysql_test.py

@@ -15,6 +15,7 @@
 from unittest import TestCase
 
 from runtime import testing
+from runtime.dbapi import table_writer
 from runtime.dbapi.mysql import MySQLConnection
 
 
@@ -70,6 +71,16 @@ def test_get_table_schema(self):
                           ('petal_length', 'FLOAT'), ('petal_width', 'FLOAT'),
                           ('class', 'INT')], col_info)
 
+    def test_proto_table_writer(self):
+        conn = MySQLConnection(testing.get_datasource())
+        rs = conn.query("select * from iris.train limit 10;")
+        self.assertTrue(rs.success())
+        tw = table_writer.ProtobufWriter(rs)
+        lines = tw.dump_strings()
+        self.assertTrue(lines[0].find(
+            "head { column_names: \"sepal_length\" column_names: \"sepal_width\" column_names: \"petal_length\" column_names: \"petal_width\" column_names: \"class\" }"  # noqa: E501
+        ) >= 0)
+
 
 if __name__ == "__main__":
     unittest.main()

python/runtime/dbapi/table_writer/__init__.py

@@ -0,0 +1,16 @@
+# Copyright 2020 The SQLFlow Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from runtime.dbapi.table_writer.protobuf_writer import ProtobufWriter
+
+__all__ = ["ProtobufWriter"]

python/runtime/dbapi/table_writer/protobuf_writer.py

@@ -0,0 +1,54 @@
+# Copyright 2020 The SQLFlow Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from google.protobuf import text_format, wrappers_pb2
+from runtime.dbapi.connection import ResultSet
+from runtime.dbapi.table_writer import sqlflow_pb2
+
+
+class ProtobufWriter:
+    def __init__(self, result_set):
+        assert isinstance(result_set, ResultSet)
+        column_info = result_set.column_info()
+        self.all_responses = []
+        head = sqlflow_pb2.Head()
+        for field_name, _ in column_info:
+            head.column_names.append(field_name)
+        self.all_responses.append(sqlflow_pb2.Response(head=head))
+        for row in result_set:
+            pb_row = sqlflow_pb2.Row()
+            for col in row:
+                any_msg = self.pod_to_pb_any(col)
+                any = pb_row.data.add()
+                any.Pack(any_msg)
+            self.all_responses.append(sqlflow_pb2.Response(row=pb_row))
+
+    @staticmethod
+    def pod_to_pb_any(value):
+        if isinstance(value, bool):
+            v = wrappers_pb2.BoolValue(value=value)
+        elif isinstance(value, int):
+            v = wrappers_pb2.Int32Value(value=value)
+        elif isinstance(value, float):
+            v = wrappers_pb2.FloatValue(value=value)
+        elif isinstance(value, str):
+            v = wrappers_pb2.StringValue(value=value)
+        else:
+            raise ValueError("not supported cell data type: %s" % type(value))
+        return v
+
+    def dump_strings(self):
+        lines = []
+        for resp in self.all_responses:
+            lines.append(text_format.MessageToString(resp, as_one_line=True))
+        return lines

scripts/test/prepare.sh

@@ -48,6 +48,8 @@ git clone https://github.com/sql-machine-learning/models.git
 git checkout v0.0.6 -b v0.0.6 && \
 python setup.py install)
 
+protoc --python_out=python/runtime/dbapi/table_writer/ -I go/proto sqlflow.proto
+
 # 3. install java parser
 echo "Build parser gRPC servers in Java ..."