omniscidb - workaround to use geopandas output

omniscidb_geopandas_output.py

# workaround to use geopandas output
import pandas as pd
import geopandas as gpd
import geoalchemy2.shape as shape
import shapely.wkt
from pymapd.cursor import Cursor
from ibis.sql.compiler import DDL, DML
from ibis.mapd.client import MapDCursor


class MapDGeoCursor:
    """Cursor to allow the OmniSciDB client to reuse machinery in ibis/client.py
    """

    def __init__(self, cursor):
        self.cursor = cursor

    def to_df(self):
        if isinstance(self.cursor, Cursor):
            col_names = [c.name for c in self.cursor.description]
            result = pd.DataFrame(self.cursor.fetchall(), columns=col_names)
            
            geo_columns = []
            for d in self.cursor.description:
                if d.type_code in (16,):
                    geo_columns.append(d.name)
                    result[d.name] = result[d.name].apply(shapely.wkt.loads)
            
            if len(geo_columns):
                result = gpd.GeoDataFrame(result, geometry=','.join(geo_columns))
            
        elif self.cursor is None:
            result = pd.DataFrame([])
        else:
            result = self.cursor

        return result

    def __enter__(self):
        # For compatibility when constructed from Query.execute()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass


def enable_geopandas_output(self):
    def _f(query=None, results=True):
        """
        query:
        :return:
        """
        if isinstance(query, (DDL, DML)):
            query = query.compile()
        # use EXECUTION_TYPE_CURSOR option
        execute = self.con.cursor().execute

        try:
            df = MapDGeoCursor(execute(query))
        except Exception as e:
            raise Exception('{}: {}'.format(e, query))
        
        if results:
            return df
    setattr(self, '_execute', _f)