Fix precomputed queries (#208)

* fix return types to work for nglui

* fixing NAs in segment properties

* Update materializationengine/blueprints/client/query.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update materializationengine/blueprints/client/query.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: fcollman

materializationengine/blueprints/client/api2.py

@@ -266,6 +266,9 @@ def fix_dataframe_types(df):
     Returns:
         pd.DataFrame: dataframe with fixed types."""
     for colname in df.columns:
+        if df[colname].isnull().all():
+            df.drop(columns=[colname], inplace=True)
+            continue
         if pd.api.types.is_float_dtype(df[colname]):
             df[colname]=df[colname].astype(np.float32)
         if pd.api.types.is_integer_dtype(df[colname]):
@@ -1151,29 +1154,21 @@ def process_fields(df, fields, column_names, tags, bool_tags, numerical):
             or field_name == "target_id"
         ):
             continue
-
+        if col not in df.columns:
+            continue
         if isinstance(field, mm_fields.String):
-            if df[col].isnull().all():
-                continue
             # check that this column is not all nulls
             tags.append(col)
         elif isinstance(field, mm_fields.Boolean):
-            if df[col].isnull().all():
-                continue
             df[col] = df[col].astype(bool)
             bool_tags.append(col)
         elif isinstance(field, PostGISField):
-            # if all the values are NaNs skip this column
-            if df[col + "_x"].isnull().all():
-                continue
             numerical.append(col + "_x")
             numerical.append(col + "_y")
             numerical.append(col + "_z")
         elif isinstance(field, mm_fields.Number):
-            if df[col].isnull().all():
-                continue
             numerical.append(col)
-
+    return df
 
 def process_view_columns(df, model, column_names, tags, bool_tags, numerical):
     for table_column_name, table_column in model.columns.items():
@@ -1246,10 +1241,10 @@ def preprocess_dataframe(df, table_name, aligned_volume_name, column_names):
         numerical = []
         bool_tags = []
 
-        process_fields(df, fields, column_names[table_name], tags, bool_tags, numerical)
+        df = process_fields(df, fields, column_names[table_name], tags, bool_tags, numerical)
 
         if table_metadata["reference_table"]:
-            process_fields(
+            df = process_fields(
                 df,
                 ref_fields,
                 column_names[ref_table],

materializationengine/blueprints/client/query.py

@@ -16,10 +16,17 @@
 from sqlalchemy.orm.util import AliasedClass
 from sqlalchemy.sql.schema import Table
 from sqlalchemy.sql.selectable import Alias
-from sqlalchemy.sql.sqltypes import Boolean, DateTime, Integer
+from sqlalchemy.sql.sqltypes import Boolean, DateTime, Integer, BigInteger, Float, String
 
 DEFAULT_SUFFIX_LIST = ["x", "y", "z", "xx", "yy", "zz", "xxx", "yyy", "zzz"]
 
+dtype_map = {
+    Boolean: pd.BooleanDtype(),
+    DateTime: np.dtype('datetime64[ns]'),
+    Integer: pd.Int32Dtype(),
+    BigInteger: pd.Int64Dtype(),
+    Float: pd.Float32Dtype(),
+}
 
 def concatenate_position_columns(df):
     grps = itertools.groupby(df.columns, key=lambda x: x[:-2])
@@ -83,7 +90,6 @@ def fix_columns_with_query(
         if n_tables == 1:
             schema_model = query.column_descriptions[0]["type"]
         for colname in df.columns:
-
             if n_tables == 1:
                 coltype = type(getattr(schema_model, colname).type)
             else:
@@ -112,7 +118,6 @@ def fix_columns_with_query(
                 )
             elif isinstance(df[colname].loc[0], Decimal) and fix_decimal is True:
                 df[colname] = _fix_decimal_column(df[colname])
-
     return df
 
 
@@ -329,10 +334,16 @@ def _execute_query(
     else:
         if direct_sql_pandas:
             statement = str(query.statement.compile(engine, compile_kwargs={"literal_binds": True}))
+            dtypes = {}
+            for k in query.statement.columns.keys():
+                coltype = query.statement.columns[k].type
+                if type(coltype) in dtype_map:
+                    dtypes[k] = dtype_map[type(coltype)]
             df = pd.read_sql(statement,
                              session.connection().connection,
                              coerce_float=True,
                              index_col=index_col,
+                             dtype=dtypes,
                              dtype_backend='numpy_nullable')   
         else:
             df = read_sql_tmpfile(