Expose a simple way to perform projections

Author: amol-

examples/minimal.py

@@ -41,11 +41,14 @@
 })
 print(orbitalml_pipeline)
 
-sql = orbitalml.export_sql("DATA_TABLE", orbitalml_pipeline, dialect="duckdb")
+sql = orbitalml.export_sql("DATA_TABLE", orbitalml_pipeline, projection=orbitalml.ResultsProjection(["sepal_width"]), dialect="duckdb")
 print("\nGenerated Query for DuckDB:")
 print(sql)
 print("\nPrediction with SQL")
 duckdb.register("DATA_TABLE", X_test)
-print(duckdb.sql(sql).df()["variable"][:5].to_numpy())
+result = duckdb.sql(sql).df()
+print(result.head())
+print("---")
+print(result["variable"][:5].to_numpy())
 print("\nPrediction with SciKit-Learn")
 print(pipeline.predict(X_test)[:5])

src/orbitalml/__init__.py

@@ -9,6 +9,6 @@
 
 from .ast import parse_pipeline
 from .sql import export_sql
-from .translate import translate
+from .translate import ResultsProjection, translate
 
-__all__ = ["parse_pipeline", "translate", "export_sql"]
+__all__ = ["parse_pipeline", "translate", "export_sql", "ResultsProjection"]

src/orbitalml/sql.py

@@ -12,7 +12,7 @@
 from ibis.expr.sql import Catalog
 
 from .ast import ParsedPipeline
-from .translate import translate
+from .translate import ResultsProjection, translate
 
 OPTIMIZER_RULES = (
     sqlglot.optimizer.optimizer.qualify,
@@ -36,6 +36,7 @@ def export_sql(
     table_name: str,
     pipeline: ParsedPipeline,
     dialect: str = "duckdb",
+    projection: ResultsProjection = ResultsProjection(),
     optimize: bool = True,
 ) -> str:
     """Export SQL for a given pipeline.
@@ -58,7 +59,12 @@ def export_sql(
         name=table_name,
     )
 
-    ibis_expr = translate(unbound_table, pipeline)
+    if projection.is_empty():
+        raise ValueError(
+            "Projection is empty. Please provide a projection to export SQL."
+        )
+
+    ibis_expr = translate(unbound_table, pipeline, projection=projection)
     sqlglot_expr = getattr(sc, dialect).compiler.to_sqlglot(ibis_expr)
 
     if optimize:

src/orbitalml/translate.py

@@ -76,7 +76,57 @@
 LOG_SQL = False
 
 
-def translate(table: ibis.Table, pipeline: ParsedPipeline) -> ibis.Table:
+class ResultsProjection:
+    """Projection of the results of the pipeline.
+
+    This class is used to select the columns to be returned
+    from the pipeline. It can be used to select specific
+    columns or to omit the results of the pipeline.
+
+    It can also be used to skip the select step of columns
+    from the pipeline.
+    """
+
+    RESULTS = object()
+    OMIT = object()
+
+    def __init__(self, select: typing.Optional[list[str]] = None) -> None:
+        """
+        :param select: A list of additional columns to be selected from the pipeline.
+                       or ResultsProjection.OMIT to skip the selection.
+        """
+        if select is self.OMIT:
+            self._select = None
+        else:
+            self._select = [self.RESULTS]
+            if select:
+                self._select.extend(select)
+
+    def is_empty(self) -> bool:
+        """Check if the projection step should be skipped."""
+        return self._select is None
+
+    def _expand(self, results: typing.Iterable[str]) -> typing.Optional[list[str]]:
+        if self._select is None:
+            return None
+
+        selected_columns = self._select
+
+        def _emit_projection() -> typing.Generator[str, None, None]:
+            for item in selected_columns:
+                if item is self.RESULTS:
+                    yield from results
+                elif isinstance(item, str):
+                    yield item
+
+        return list(_emit_projection())
+
+
+def translate(
+    table: ibis.Table,
+    pipeline: ParsedPipeline,
+    projection: ResultsProjection = ResultsProjection(),
+) -> ibis.Table:
     """Translate a pipeline into an Ibis expression.
 
     This function takes a pipeline and a table and translates the pipeline
@@ -98,14 +148,18 @@ def translate(table: ibis.Table, pipeline: ParsedPipeline) -> ibis.Table:
         translator.process()
         table = translator.mutated_table  # Translator might return a new table.
         _log_debug_end(translator, variables)
-    return _projection_results(table, variables)
+    return _projection_results(table, variables, projection)
 
 
-def _projection_results(table: ibis.Table, variables: GraphVariables) -> ibis.Table:
+def _projection_results(
+    table: ibis.Table,
+    variables: GraphVariables,
+    projection: ResultsProjection = ResultsProjection(),
+) -> ibis.Table:
     # As we pop out the variables as we use them
     # the remaining ones are the values resulting from all
     # graph branches.
-    final_projections = {}
+    final_projections: dict[str, typing.Any] = {}
     for key, value in variables.remaining().items():
         if isinstance(value, dict):
             for field in value:
@@ -116,7 +170,11 @@ def _projection_results(table: ibis.Table, variables: GraphVariables) -> ibis.Ta
                 final_projections[colkey] = colvalue
         else:
             final_projections[key] = value
-    return table.mutate(**final_projections).select(final_projections.keys())
+    query = table.mutate(**final_projections)
+    selection = projection._expand(final_projections.keys())
+    if selection is not None:
+        query = query.select(*selection)
+    return query
 
 
 def _log_debug_start(translator: Translator, variables: GraphVariables) -> None:

tests/test_pipeline_e2e.py

@@ -94,6 +94,30 @@ def test_simple_linear_regression(self, iris_data, db_connection):
             sql_results.values.flatten(), sklearn_preds.flatten(), rtol=1e-4, atol=1e-4
         )
 
+    def test_simple_linear_with_projection(self, iris_data, db_connection):
+        df, feature_names = iris_data
+        conn, dialect = db_connection
+
+        sklearn_pipeline = Pipeline(
+            [("scaler", StandardScaler()), ("regression", LinearRegression())]
+        )
+        X = df[feature_names]
+        y = df["target"]
+        sklearn_pipeline.fit(X, y)
+        sklearn_preds = sklearn_pipeline.predict(X)
+
+        features = {fname: types.FloatColumnType() for fname in feature_names}
+        parsed_pipeline = orbitalml.parse_pipeline(sklearn_pipeline, features=features)
+
+        sql = orbitalml.export_sql("data", parsed_pipeline, projection=orbitalml.ResultsProjection(["sepal_length"]), dialect=dialect)
+
+        sql_results = self.execute_sql(sql, conn, dialect, df)
+        print(sql_results)
+        assert set(sql_results.columns) == {"sepal_length", "variable.target_0"}
+        np.testing.assert_allclose(
+            sql_results["variable.target_0"].values.flatten(), sklearn_preds.flatten(), rtol=1e-4, atol=1e-4
+        )
+
     def test_feature_selection_pipeline(self, diabetes_data, db_connection):
         df, feature_names = diabetes_data
         conn, dialect = db_connection