Changes towards an engine that could help in Neurolang

.github/workflows/test.yml

@@ -84,7 +84,7 @@ jobs:
         run: |
           # explicitly install docker, fugue and sqlalchemy package
           conda install sqlalchemy psycopg2 -c conda-forge
-          pip install docker fugue
+          pip install docker "fugue<=0.5.0"
         if: matrix.os == 'ubuntu-latest'
       - name: Install Java (again) and test with pytest
         shell: bash -l {0}

dask_sql/context.py

@@ -68,6 +68,8 @@ def __init__(self):
         """
         # Storage for the registered tables
         self.tables = {}
+        # Storage for the registered views
+        self.views = {}
         # Storage for the registered functions
         self.functions: Dict[str, Callable] = {}
         self.function_list: List[FunctionDescription] = []
@@ -86,6 +88,8 @@ def __init__(self):
         RelConverter.add_plugin_class(logical.LogicalSortPlugin, replace=False)
         RelConverter.add_plugin_class(logical.LogicalTableScanPlugin, replace=False)
         RelConverter.add_plugin_class(logical.LogicalUnionPlugin, replace=False)
+        RelConverter.add_plugin_class(logical.LogicalIntersectPlugin, replace=False)
+        RelConverter.add_plugin_class(logical.LogicalMinusPlugin, replace=False)
         RelConverter.add_plugin_class(logical.LogicalValuesPlugin, replace=False)
         RelConverter.add_plugin_class(logical.SamplePlugin, replace=False)
         RelConverter.add_plugin_class(custom.AnalyzeTablePlugin, replace=False)
@@ -116,6 +120,7 @@ def create_table(
         input_table: InputType,
         format: str = None,
         persist: bool = True,
+        sql: str = None,
         **kwargs,
     ):
         """
@@ -191,6 +196,8 @@ def create_table(
             **kwargs,
         )
         self.tables[table_name.lower()] = dc
+        if sql is not None:
+            self.views[table_name.lower()] = sql
 
     def register_dask_table(self, df: dd.DataFrame, name: str):
         """
@@ -448,11 +455,17 @@ def _prepare_schema(self):
             logger.warning("No tables are registered.")
 
         for name, dc in self.tables.items():
-            table = DaskTable(name)
             df = dc.df
-            logger.debug(
-                f"Adding table '{name}' to schema with columns: {list(df.columns)}"
-            )
+            if name in self.views:
+                table = DaskTable(name, self.views[name])
+                logger.debug(
+                    f"Adding materialied table '{name}' to schema with columns: {list(df.columns)}"
+                )
+            else:
+                table = DaskTable(name)
+                logger.debug(
+                    f"Adding table '{name}' to schema with columns: {list(df.columns)}"
+                )
             for column in df.columns:
                 data_type = df[column].dtype
                 sql_data_type = python_to_sql_type(data_type)
@@ -513,8 +526,14 @@ def _get_ral(self, sql):
             else:
                 validatedSqlNode = generator.getValidatedNode(sqlNode)
                 nonOptimizedRelNode = generator.getRelationalAlgebra(validatedSqlNode)
+                rel_string_non_op = str(generator.getRelationalAlgebraString(nonOptimizedRelNode))
+                rel_non_op_count = rel_string_non_op.count('\n')
                 rel = generator.getOptimizedRelationalAlgebra(nonOptimizedRelNode)
                 rel_string = str(generator.getRelationalAlgebraString(rel))
+                logger.debug(
+                    f"Non optimised query plan: {rel_non_op_count} ops\n "
+                    f"{rel_string_non_op}"
+                )
         except (ValidationException, SqlParseException) as e:
             logger.debug(f"Original exception raised by Java:\n {e}")
             # We do not want to re-raise an exception here
@@ -544,8 +563,8 @@ def _get_ral(self, sql):
                 "Not extracting output column names as the SQL is not a SELECT call"
             )
             select_names = None
-
-        logger.debug(f"Extracted relational algebra:\n {rel_string}")
+        br = '\n'
+        logger.debug(f"Extracted relational algebra {rel_string.count(br)} ops:\n {rel_string}")
         return rel, select_names, rel_string
 
     def _to_sql_string(self, s: "org.apache.calcite.sql.SqlNode", default_dialect=None):

dask_sql/datacontainer.py

@@ -42,7 +42,8 @@ def _copy(self) -> ColumnContainer:
         Internal function to copy this container
         """
         return ColumnContainer(
-            self._frontend_columns.copy(), self._frontend_backend_mapping.copy()
+            self._frontend_columns.copy(),
+            self._frontend_backend_mapping.copy(),
         )
 
     def limit_to(self, fields: List[str]) -> ColumnContainer:
@@ -137,7 +138,9 @@ def make_unique(self, prefix="col"):
         where <number> is the column index.
         """
         return self.rename(
-            columns={str(col): f"{prefix}_{i}" for i, col in enumerate(self.columns)}
+            columns={
+                str(col): f"{prefix}_{i}" for i, col in enumerate(self.columns)
+            }
         )
 
 
@@ -166,11 +169,22 @@ def assign(self) -> dd.DataFrame:
         a dataframe which has the the columns specified in the
         stored ColumnContainer.
         """
-        df = self.df.assign(
-            **{
-                col_from: self.df[col_to]
-                for col_from, col_to in self.column_container.mapping()
-                if col_from in self.column_container.columns
-            }
-        )
+        # We rename as many cols as possible because renaming is much more
+        # efficient than assigning.
+
+        renames = {}
+        assigns = {}
+        for col_from, col_to in self.column_container.mapping():
+            if col_from in self.column_container.columns:
+                if (
+                    len(renames) < len(self.df.columns)
+                    and col_to not in renames
+                    and (col_from not in self.df.columns or col_from == col_to)
+                ):
+                    renames[col_to] = col_from
+                else:
+                    assigns[col_from] = self.df[col_to]
+        df = self.df.rename(columns=renames)
+        if len(assigns) > 0:
+            df = df.assign(**assigns)
         return df[self.column_container.columns]

dask_sql/physical/rel/convert.py

@@ -1,4 +1,5 @@
 import logging
+import time
 
 import dask.dataframe as dd
 
@@ -53,6 +54,8 @@ def convert(
         logger.debug(
             f"Processing REL {rel} using {plugin_instance.__class__.__name__}..."
         )
+        start_time = time.perf_counter()
         df = plugin_instance.convert(rel, context=context)
-        logger.debug(f"Processed REL {rel} into {LoggableDataFrame(df)}")
+        elapsed_time = time.perf_counter() - start_time
+        logger.debug(f"Processed REL {rel} into {LoggableDataFrame(df)} ({elapsed_time}s)")
         return df

dask_sql/physical/rel/logical/__init__.py

@@ -6,6 +6,8 @@
 from .sort import LogicalSortPlugin
 from .table_scan import LogicalTableScanPlugin
 from .union import LogicalUnionPlugin
+from .intersect import LogicalIntersectPlugin
+from .minus import LogicalMinusPlugin
 from .values import LogicalValuesPlugin
 
 __all__ = [
@@ -16,6 +18,8 @@
     LogicalSortPlugin,
     LogicalTableScanPlugin,
     LogicalUnionPlugin,
+    LogicalIntersectPlugin,
+    LogicalMinusPlugin,
     LogicalValuesPlugin,
     SamplePlugin,
 ]