Update PyLagoon (#4)

Updated PyLagoon's query methods to allow them to support
cross-source queries and added a PyLagoon example which corresponds
to the annoucement blog post. Also fixed the image name
in the docker-compose file and some other references in the documentation.

Author: dorranh

.gitignore

@@ -50,3 +50,6 @@ RubyLagoon/vendor
 .vscode
 
 clients/RLagoon/_libs
+
+# Extra notebook files
+.ipynb_checkpoints

clients/PyLagoon/PyLagoon/datalake.py renamed to clients/PyLagoon/PyLagoon/lagoon.py

@@ -15,6 +15,7 @@
 import pandas as pd
 import json
 import os.path
+from typing import List
 
 from PyLagoon.config import LagoonConfig
 from PyLagoon.postgresql import PGMeta, build_sql_query
@@ -36,9 +37,10 @@ def __init__(self, config, host=None, port=None):
             self.__cookies = self.__connect()
 
     def __connect(self):
-        reply = req.post(self.conn_str + "user/login",
-                         json={"user": self.__cfg.USER,
-                               "pass": self.__cfg.PASSWORD})
+        reply = req.post(
+            self.conn_str + "user/login",
+            json={"user": self.__cfg.USER, "pass": self.__cfg.PASSWORD},
+        )
         if reply.ok:
             return reply.cookies
         else:
@@ -61,9 +63,7 @@ def sources(self, ontoClass=None, tags=None, columns=None, **kwargs):
             kwargs["tag"] = tags
         if columns:
             kwargs["columns"] = columns
-        reply = req.get(self.conn_str + "sources",
-                        params=kwargs,
-                        cookies=self.__cookies)
+        reply = req.get(self.conn_str + "sources", params=kwargs, cookies=self.__cookies)
         return [Source(j) for j in reply.json()]
 
     def ingest(self, file_path, name, ontoClass=None, tags=None, **kwargs):
@@ -79,12 +79,14 @@ def ingest(self, file_path, name, ontoClass=None, tags=None, **kwargs):
             kwargs["tag"] = tags
         kwargs["name"] = name
         kwargs["input"] = os.path.split(file_path)[1]
-                          # So the server can guess the fileType
-        reply = req.post(self.conn_str + "sources",
-                         data=open(file_path, "rb"),
-                         params=kwargs,
-                         stream=True,
-                         cookies=self.__cookies)
+        # So the server can guess the fileType
+        reply = req.post(
+            self.conn_str + "sources",
+            data=open(file_path, "rb"),
+            params=kwargs,
+            stream=True,
+            cookies=self.__cookies,
+        )
         report = (json.loads(line.decode("utf-8")) for line in reply.raw)
         stack = []
         last = None
@@ -116,38 +118,67 @@ def users(self):
         reply = req.get(self.conn_str + "users")
         return reply.json()
 
-    def tbl(self, source=None, query=None):
-        """tbl() in RLagoon
-
-        Give one of source or query.
-
-        source is a Source, query is an sqlalchemy.orm.query.Query created
-        through use of PyLagoon.postgresql.PGMeta and the sqlalchemy EDSL.
+    def download_source(self, source):
+        """Constructs a DataFrame containing an entire source
         """
+        is_json = any(c["type"][0] == "JSON" for c in source.columns)
+        if is_json:
+            # We need a JSON document in that case
+            # the sql endpoint will return one
+            meta = PGMeta([source])
+            table = meta[source]
+            return self.__tbl_from_raw_sql(build_sql_query(meta.query(table)))
+        else:
+            reply = req.get(
+                self.conn_str + "source/" + str(source.ix) + "/download",
+                stream=True,
+                cookies=self.__cookies,
+            )
+            if reply.ok:
+                return pd.read_csv(reply.text)
+
+    def download_query(self, query, sources):
+        """Constructs a DataFrame from a SQLAlchemy query and corresponding sources
+
+        Note that this method will sequentially search for each columns type in the list
+        of sources and take the first match. This is necessary since query results only
+        include column names and not data source identifiers.
+        """
+        return self.__tbl_from_raw_sql(build_sql_query(query), sources)
 
-        if source:
-            is_json = any(c["type"][0] == "JSON" for c in source.columns)
-            if is_json:
-                # We need a JSON document in that case
-                # the sql endpoint will return one
-                meta = PGMeta([source])
-                table = meta[source]
-                return self.__tbl_from_raw_sql(build_sql_query(meta.query(table)))
-            else:
-                reply = req.get(self.conn_str + "source/" + str(source.ix) + "/download",
-                                stream=True,
-                                cookies=self.__cookies)
-                if reply.ok:
-                    return pd.read_csv(reply.text)
-        elif query:
-            return self.__tbl_from_raw_sql(build_sql_query(query))
-
-    def __tbl_from_raw_sql(self, query):
+    def __tbl_from_raw_sql(self, query, sources):
         reply = req.post(
-            self.conn_str + "sql",
-            json={"sql": query},
-            stream=True,
-            cookies=self.__cookies)
-        if reply.ok:
-            # return pd.read_json(reply.raw)
-            return pd.DataFrame(reply.json())
+            self.conn_str + "sql", json={"sql": query}, stream=True, cookies=self.__cookies
+        )
+        reply.raise_for_status()
+        return _query_to_df(reply.json(), sources)
+
+
+def _group_rows(rows: List[dict]):
+    columns = {}
+    for row in rows:
+        for c, v in row.items():
+            if c in columns:
+                columns[c].append(v)
+            else:
+                columns[c] = [v]
+    return columns
+
+
+def _get_dtype(col_name, sources):
+    for source in sources:
+        if col_name in source.col_types:
+            return source.col_types[col_name]
+    return object
+
+
+def _query_to_df(rows, sources):
+    grouped = _group_rows(rows)
+    col_names = list(grouped.keys())
+    series = []
+    for name in col_names:
+        vals = grouped.pop(name)
+        series.append(pd.Series(vals, name=name, dtype=_get_dtype(name, sources)))
+    df = pd.concat(series, axis=1)
+    df.columns = col_names
+    return df

clients/PyLagoon/PyLagoon/postgresql.py

@@ -24,6 +24,7 @@
 from PyLagoon.config import LagoonConfig
 from PyLagoon.source import Source
 
+
 class PGMeta:
     """Given a list of Sources, builds the classes necessary
     to the querying EDSL.
@@ -34,7 +35,7 @@ def __init__(self, sources):
         self.__md = sa.MetaData()
         for s in sources:
             self.__add_source_to_md(s)
-        self.__base = automap_base(metadata = self.__md)
+        self.__base = automap_base(metadata=self.__md)
         self.__base.prepare()
 
     @property
@@ -48,8 +49,9 @@ def __getitem__(self, key):
         return self.__base.classes[key]
 
     def __sql_column_from_json_column(self, col):
-        typ = col["type"] \
-              .replace(" ", "_").replace("DOCUMENT", "TEXT") # "DOUBLE PRECISION" -> "DOUBLE_PRECISION"
+        typ = (
+            col["type"].replace(" ", "_").replace("DOCUMENT", "TEXT")
+        )  # "DOUBLE PRECISION" -> "DOUBLE_PRECISION"
         subtyp = None
         sql_typ = None
         if isinstance(typ, list):
@@ -66,10 +68,13 @@ def __sql_column_from_json_column(self, col):
             return sa.Column(col["inView"], sql_typ)
 
     def __add_source_to_md(self, source):
-        sa.Table(source.view_name, self.__md,
-                 sa.Column("ix", sa.Integer, primary_key=True),
-                 *(self.__sql_column_from_json_column(c)
-                   for c in source.columns))
+        sa.Table(
+            source.view_name,
+            self.__md,
+            sa.Column("ix", sa.Integer, primary_key=True),
+            *(self.__sql_column_from_json_column(c) for c in source.columns.values()),
+            schema=source.schema
+        )
 
     def query(self, *sources):
         """Starts a query on the given sources, which can
@@ -82,9 +87,9 @@ def query(self, *sources):
 
         if len(sources) == 0:
             sources = self.__view_names
-        return Session().query(*((self[s] if isinstance(s,str) or isinstance(s,Source)
-                                          else s)
-                                 for s in sources))
+        return Session().query(
+            *((self[s] if isinstance(s, str) or isinstance(s, Source) else s) for s in sources)
+        )
 
 
 def build_sql_query(query):
@@ -93,8 +98,8 @@ def build_sql_query(query):
     to the lagoon-server"""
 
     d = psql.dialect()
-    q = query.statement.compile(dialect = d)
-    #The following is not ideal, as q.params and str(q) should
+    q = query.statement.compile(dialect=d)
+    # The following is not ideal, as q.params and str(q) should
     # normally be passed separately to the PostgreSQL database:
     ps = {}
     for k, v in q.params.items():

clients/PyLagoon/PyLagoon/source.py

@@ -11,16 +11,34 @@
 # 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.
+import pandas as pd
+import numpy as np
+
+LAGOON_TYPES = {
+    "BOOLEAN": bool,
+    "INTEGER": int,
+    "BIGINT": pd.Int64Dtype,
+    "DOUBLE PRECISION": np.float64,
+    "TEXT": str,
+    "ARR": object,
+    "DOCUMENT": object,
+    "JSON": object,
+}
+UNKNOWN_COLUMN_TYPE = object
+
+
 class Source:
     """A wrapper for the json description of sources
     returned by lagoon-server"""
-    
+
+    INDEX_COL = "ix"
+    INDEX_COL_TYPE = int
+
     def __init__(self, json):
         self.__json = json
 
     def __str__(self):
-        return "<Source: ix={0}, view_name={1}>".format(
-            self.ix, self.view_name)
+        return "<Source: ix={0}, view_name={1}>".format(self.ix, self.view_name)
 
     def __repr__(self):
         return str(self)
@@ -35,8 +53,30 @@ def view_name(self):
 
     @property
     def columns(self):
-        return self.__json["columns"]
+        return {get_column_name(c): c for c in self.__json["columns"]}
 
     @property
     def _json(self):
         return self.__json
+
+    @property
+    def schema(self):
+        return self.__json["schema"]
+
+    @property
+    def col_types(self):
+        types = {name: get_column_type(col) for name, col in self.columns.items()}
+        types[self.INDEX_COL] = self.INDEX_COL_TYPE
+        return types
+
+
+def get_column_type(column: dict):
+    column_type = column["type"]
+    if column_type in LAGOON_TYPES:
+        return LAGOON_TYPES[column_type]
+    else:
+        return UNKNOWN_COLUMN_TYPE
+
+
+def get_column_name(column: dict):
+    return column["inView"]