support flipped operands in Sub

.gitignore

@@ -82,6 +82,9 @@ target/
 profile_default/
 ipython_config.py
 
+# Agents
+.pi-subagents/
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:

src/orbital/translation/steps/div.py

@@ -1,11 +1,7 @@
 """Defines the translation step for the Div operation."""
 
-import typing
-
-import ibis
-
 from ..translator import Translator
-from ..variables import NumericVariablesGroup, ValueVariablesGroup, VariablesGroup
+from ..variables import ValueVariablesGroup
 
 
 class DivTranslator(Translator):
@@ -29,8 +25,10 @@ def process(self) -> None:
         """Performs the translation and set the output variable."""
         # https://onnx.ai/onnx/operators/onnx__Div.html
 
-        left_keys, left_values = self._operand_values(self.inputs[0])
-        right_keys, right_values = self._operand_values(self.inputs[1])
+        left_keys, left_values = self._variables.consume_operand_values(self.inputs[0])
+        right_keys, right_values = self._variables.consume_operand_values(
+            self.inputs[1]
+        )
 
         # The first operand that is a group dictates the width of the result
         # and, at the very end, the names of the resulting columns.
@@ -68,30 +66,3 @@ def process(self) -> None:
                 }
             )
         )
-
-    def _operand_values(
-        self, name: str
-    ) -> tuple[typing.Optional[list[str]], list[ibis.expr.types.NumericValue]]:
-        """Consume an operand and return its column names and its values.
-
-        The column names are ``None`` for anything that is not a group of columns,
-        which is what tells apart a group of one column from a plain single value.
-
-        :param name: Name of the variable or constant to consume.
-        """
-        # Classifiers normalize the probabilities of each class with
-        # Div(scores, ReduceSum(Abs(scores))), so an operand can also be a
-        # column computed by a previous node instead of a constant.
-        operand = self._variables.consume(name)
-        if isinstance(operand, VariablesGroup):
-            group = NumericVariablesGroup(operand)
-            return list(group.keys()), list(group.values())
-
-        values = []
-        for value in operand if isinstance(operand, (list, tuple)) else [operand]:
-            if isinstance(value, (int, float)):
-                value = ibis.literal(value)
-            elif not isinstance(value, ibis.expr.types.NumericValue):
-                raise ValueError("Div: the operands must be numeric values.")
-            values.append(value)
-        return None, typing.cast(list[ibis.expr.types.NumericValue], values)

src/orbital/translation/steps/sub.py

@@ -1,70 +1,68 @@
 """Implementation of the Sub operator."""
 
-import typing
-
-import ibis
-
-from orbital.translation.variables import (
-    NumericVariablesGroup,
-    ValueVariablesGroup,
-    VariablesGroup,
-)
-
 from ..translator import Translator
+from ..variables import ValueVariablesGroup
 
 
 class SubTranslator(Translator):
     """Processes a Sub node and updates the variables with the output expression.
 
-    Given the node to translate, the variables and constants available for
-    the translation context, generates a query expression that processes
-    the input variables and produces a new output variable that computes
-    based on the Sub operation.
+    Both operands are treated symmetrically: each one can be a group of columns,
+    a single column, a constant scalar or a constant list of values.
+
+    When any of the two operands is a group of columns, the result is a group of
+    columns too and it borrows its column names from the first operand that is a
+    group, as those names end up being the names of the resulting SQL columns.
+    That group also dictates the width of the result: any other operand must
+    either provide exactly as many values, or a single value that is subtracted
+    from (or subtracts) every column of the group.
+
+    When neither operand is a group, both must be single values and the result
+    is a single column.
     """
 
     def process(self) -> None:
         """Performs the translation and set the output variable."""
         # https://onnx.ai/onnx/operators/onnx__Sub.html
-        assert len(self._inputs) == 2, "The Sub node must have exactly 2 inputs."
-
-        first_operand = self._variables.consume(self._inputs[0])
-        second_operand = self._variables.get_initializer_value(self._inputs[1])
-        if second_operand is None or not isinstance(second_operand, (list, tuple)):
-            raise NotImplementedError(
-                "Sub: Second input (divisor) must be a constant list."
-            )
 
-        type_check_var = first_operand
-        if isinstance(type_check_var, dict):
-            type_check_var = next(iter(type_check_var.values()), None)
-        if not isinstance(type_check_var, ibis.expr.types.NumericValue):
-            raise ValueError("Sub: The first operand must be a numeric value.")
+        left_keys, left_values = self._variables.consume_operand_values(self.inputs[0])
+        right_keys, right_values = self._variables.consume_operand_values(
+            self.inputs[1]
+        )
 
-        sub_values = list(second_operand)
-        if isinstance(first_operand, VariablesGroup):
-            first_operand = NumericVariablesGroup(first_operand)
-            struct_fields = list(first_operand.keys())
-            assert len(sub_values) == len(struct_fields), (
-                f"The number of values in the initializer ({len(sub_values)}) must match the number of fields ({len(struct_fields)}"
-            )
-            self.set_output(
-                ValueVariablesGroup(
-                    {
-                        field: (
-                            self._optimizer.fold_operation(
-                                first_operand[field] - sub_values[i]
-                            )
-                        )
-                        for i, field in enumerate(struct_fields)
-                    }
+        # The first operand that is a group dictates the width of the result
+        # and, at the very end, the names of the resulting columns.
+        keys = left_keys if left_keys is not None else right_keys
+        if keys is None:
+            # Simple case, no columns group involved, so we just subtract the two values.
+            if len(left_values) != 1 or len(right_values) != 1:
+                raise ValueError(
+                    "Sub: when no operand is a group of columns, each operand must contain only one value."
                 )
+            self.set_output(
+                self._optimizer.fold_operation(left_values[0] - right_values[0])
             )
-        else:
-            if len(sub_values) != 1:
+            return
+
+        for values in (left_values, right_values):
+            if len(values) not in (1, len(keys)):
                 raise ValueError(
-                    "When the first operand is a single column, the second operand must contain exactly 1 value"
+                    "Sub: the number of values of each operand must match the number of columns of the resulting group."
                 )
-            first_operand = typing.cast(ibis.expr.types.NumericValue, first_operand)
-            self.set_output(
-                self._optimizer.fold_operation(first_operand - sub_values[0])
+
+        # A single value is shared by all columns of the resulting group.
+        if len(left_values) == 1:
+            left_values = left_values * len(keys)
+        if len(right_values) == 1:
+            right_values = right_values * len(keys)
+
+        self.set_output(
+            ValueVariablesGroup(
+                {
+                    key: self._optimizer.fold_operation(left_value - right_value)
+                    for key, left_value, right_value in zip(
+                        keys, left_values, right_values
+                    )
+                }
             )
+        )

src/orbital/translation/variables.py

@@ -144,6 +144,33 @@ def consume(
         self._consumed.add(name)
         return self._variables[name]
 
+    def consume_operand_values(
+        self, name: str
+    ) -> tuple[typing.Optional[list[str]], list[ibis.expr.types.NumericValue]]:
+        """Consume an operand of a mathematical operation as column names and values.
+
+        The column names are ``None`` for anything that is not a group of columns,
+        which is what tells apart a group of one column from a plain single value.
+
+        :param name: Name of the variable or constant to consume.
+        """
+        # Classifiers normalize the probabilities of each class with
+        # Div(scores, ReduceSum(Abs(scores))) for example, so an operand can
+        # also be a column computed by a previous node instead of a constant.
+        operand = self.consume(name)
+        if isinstance(operand, VariablesGroup):
+            group = NumericVariablesGroup(operand)
+            return list(group.keys()), list(group.values())
+
+        values = []
+        for value in operand if isinstance(operand, (list, tuple)) else [operand]:
+            if isinstance(value, (int, float)):
+                value = ibis.literal(value)
+            elif not isinstance(value, ibis.expr.types.NumericValue):
+                raise ValueError("The operands must be numeric values.")
+            values.append(value)
+        return None, typing.cast(list[ibis.expr.types.NumericValue], values)
+
     def peek_variable(
         self, name: str, default: typing.Optional[ibis.Expr] = None
     ) -> typing.Union[ibis.Expr, VariablesGroup, None]:

tests/test_pipeline_e2e.py

@@ -3,7 +3,7 @@
 from sklearn.compose import ColumnTransformer
 from sklearn.feature_selection import SelectKBest, f_regression
 from sklearn.linear_model import ElasticNet, LinearRegression, LogisticRegression
-from sklearn.neural_network import MLPRegressor
+from sklearn.neural_network import MLPClassifier, MLPRegressor
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import MinMaxScaler, OneHotEncoder, StandardScaler
 
@@ -272,3 +272,50 @@ def test_tanh_mlp_regression(self, iris_data, db_connection):
             rtol=1e-4,
             atol=1e-4,
         )
+
+    def test_binary_mlp_classifier(self, iris_data, db_connection):
+        """Test a binary MLP classifier, whose negative class is Sub(1.0, probability)."""
+        df, feature_names = iris_data
+        conn, dialect = db_connection
+
+        binary_df = df[df["target"].isin([0, 1])].copy()
+
+        sklearn_pipeline = Pipeline(
+            [
+                ("scaler", StandardScaler()),
+                (
+                    "mlp",
+                    MLPClassifier(
+                        hidden_layer_sizes=(8,),
+                        # lbfgs converges on this small dataset,
+                        # avoiding ConvergenceWarning noise from adam.
+                        solver="lbfgs",
+                        max_iter=500,
+                        random_state=0,
+                    ),
+                ),
+            ]
+        )
+
+        X = binary_df[feature_names]
+        y = binary_df["target"]
+        sklearn_pipeline.fit(X, y)
+        sklearn_proba = pd.DataFrame(
+            sklearn_pipeline.predict_proba(X),
+            columns=sklearn_pipeline.classes_,
+            index=binary_df.index,
+        )
+
+        features = {fname: types.FloatColumnType() for fname in feature_names}
+        parsed_pipeline = orbital.parse_pipeline(sklearn_pipeline, features=features)
+
+        sql = orbital.export_sql("data", parsed_pipeline, dialect=dialect)
+        sql_results = execute_sql(sql, conn, dialect, binary_df)
+
+        for class_label in sklearn_pipeline.classes_:
+            np.testing.assert_allclose(
+                sql_results[f"output_probability.{class_label}"].values.flatten(),
+                sklearn_proba[class_label].values.flatten(),
+                rtol=1e-4,
+                atol=1e-4,
+            )