Moveing forward with typing

Author: amol-

src/mustela/translation/steps/add.py

@@ -4,6 +4,7 @@
 import ibis
 
 from ..translator import Translator
+from ..variables import NumericVariablesGroup, VariablesGroup
 
 
 class AddTranslator(Translator):
@@ -25,33 +26,33 @@ def process(self) -> None:
             raise NotImplementedError("Add: Second input (divisor) must be a constant list.")
 
         type_check_var = first_operand
-        if isinstance(type_check_var, dict):
+        if isinstance(type_check_var, VariablesGroup):
             type_check_var = next(iter(type_check_var.values()), None)        
         if not isinstance(type_check_var, ibis.expr.types.NumericValue):
             raise ValueError("Add: The first operand must be a numeric value.")
 
         add_values = list(second_operand)
-        if isinstance(first_operand, dict):
-            first_operand = typing.cast(dict[str, ibis.expr.types.NumericValue], first_operand)
+        if isinstance(first_operand, VariablesGroup):
+            first_operand = NumericVariablesGroup(first_operand)
             struct_fields = list(first_operand.keys())
             if len(add_values) != len(struct_fields):
                 # TODO: Implement dividing by a single value,
                 #       see Div implementation.
                 raise ValueError(
                     "When the first operand is a group of columns, the second operand must contain the same number of values"
                 )
-            self._variables[self._output_name] = {
+            self.set_output(VariablesGroup({
                 field: (
                     self._optimizer.fold_operation(first_operand[field] + add_values[i])
                 )
                 for i, field in enumerate(struct_fields)
-            }
+            }))
         else:
             if len(add_values) != 1:
                 raise ValueError(
                     "When the first operand is a single column, the second operand must contain exactly 1 value"
                 )
             first_operand = typing.cast(ibis.expr.types.NumericValue, first_operand)
-            self._variables[self._output_name] = self._optimizer.fold_operation(
+            self.set_output(self._optimizer.fold_operation(
                 first_operand + add_values[0]
-            )
+            ))

src/mustela/translation/steps/arrayfeatureextractor.py

@@ -3,6 +3,7 @@
 import ibis.expr.types
 
 from ..translator import Translator
+from ..variables import VariablesGroup
 
 
 class ArrayFeatureExtractorTranslator(Translator):
@@ -34,7 +35,7 @@ def process(self) -> None:
         data = self._variables.consume(self.inputs[0])
         indices = self._variables.consume(self.inputs[1])
 
-        if isinstance(data, dict):
+        if isinstance(data, VariablesGroup):
             # We are selecting a set of columns out of a column group
 
             # This expects that dictionaries are sorted by insertion order
@@ -49,7 +50,7 @@ def process(self) -> None:
                 raise ValueError("Indices requested are more than the available numer of columns.")
 
             # Pick only the columns that are in the list of indicies.
-            result = {data_keys[i]: data[i] for i in indices}
+            result = VariablesGroup({data_keys[i]: data[i] for i in indices})
         elif isinstance(data, (tuple, list)):
             # We are selecting values out of a list of values
             # This is usually used to select "classes" out of a list of

src/mustela/translation/steps/cast.py

@@ -1,15 +1,19 @@
 """Translators for Cast and CastLike operations"""
+import typing
 
 import onnx
 
+import ibis
+
 from ..translator import Translator
+from ..variables import VariablesGroup
 
-ONNX_TYPES_TO_IBIS = {
-    onnx.TensorProto.FLOAT: "float32",  # 1: FLOAT
-    onnx.TensorProto.DOUBLE: "float64",  # 11: DOUBLE
-    onnx.TensorProto.STRING: "string",  # 8: STRING
-    onnx.TensorProto.INT64: "int64",  # 7: INT64
-    onnx.TensorProto.BOOL: "bool",  # 9: BOOL
+ONNX_TYPES_TO_IBIS: dict[int, ibis.expr.datatypes.DataType] = {
+    onnx.TensorProto.FLOAT: ibis.expr.datatypes.float32,   # 1: FLOAT
+    onnx.TensorProto.DOUBLE: ibis.expr.datatypes.float64,  # 11: DOUBLE
+    onnx.TensorProto.STRING: ibis.expr.datatypes.string,   # 8: STRING
+    onnx.TensorProto.INT64: ibis.expr.datatypes.int64,     # 7: INT64
+    onnx.TensorProto.BOOL: ibis.expr.datatypes.boolean,  # 9: BOOL
 }
 
 
@@ -23,19 +27,23 @@ def process(self) -> None:
         """Performs the translation and set the output variable."""
         # https://onnx.ai/onnx/operators/onnx__Cast.html
         expr = self._variables.consume(self.inputs[0])
-        to_type = self._attributes["to"]
-        if to_type in ONNX_TYPES_TO_IBIS:
-            target_type = ONNX_TYPES_TO_IBIS[to_type]
-            if isinstance(expr, dict):
-                casted = {
-                    k: self._optimizer.fold_cast(expr[k].cast(target_type))
-                    for k in expr
-                }
-                self.set_output(casted)
-            else:
-                self.set_output(self._optimizer.fold_cast(expr.cast(target_type)))
-        else:
+        to_type: int = typing.cast(int, self._attributes["to"])
+        if to_type not in ONNX_TYPES_TO_IBIS:
             raise NotImplementedError(f"Cast: type {to_type} not supported")
+    
+        target_type = ONNX_TYPES_TO_IBIS[to_type]
+        if isinstance(expr, VariablesGroup):
+            casted = VariablesGroup({
+                k: self._optimizer.fold_cast(expr.as_value(k).cast(target_type))
+                for k in expr
+            })
+            self.set_output(casted)
+        elif isinstance(expr, ibis.Value):
+            self.set_output(self._optimizer.fold_cast(expr.cast(target_type)))
+        else:
+            raise ValueError(
+                f"Cast: expected a column group or a single column. Got {type(expr)}"
+            )
 
 
 class CastLikeTranslator(Translator):
@@ -56,25 +64,26 @@ def process(self) -> None:
         like_expr = self._variables.consume(self.inputs[1])
 
         # Assert that the first input is a dict (multiple concatenated columns).
-        if not isinstance(expr, dict):
+        if not isinstance(expr, VariablesGroup):
             # TODO: Support single variables as well.
             #       This should be fairly straightforward to implement,
             #       but there hasn't been the need for it yet.
             raise NotImplementedError("CastLike currently only supports casting a group of columns.")
 
         # Assert that the second input is a single expression.
-        if isinstance(like_expr, dict):
+        if isinstance(like_expr, VariablesGroup):
             raise NotImplementedError("CastLike currently only supports casting to a single column type, not a group.")
 
-        assert hasattr(like_expr, "type"), (
-            "CastLike: second input must have a 'type' attribute."
-        )
+        if not isinstance(like_expr, ibis.Value):
+            raise ValueError(
+                f"CastLike: expected a single column. Got {type(like_expr)}"
+            )
 
         # Get the target type from the second input.
-        target_type = like_expr.type()
+        target_type: ibis.DataType = like_expr.type()
 
         # Now cast each field in the dictionary to the target type.
-        casted = {
-            key: self._optimizer.fold_cast(expr[key].cast(target_type)) for key in expr
-        }
+        casted = VariablesGroup({
+            key: self._optimizer.fold_cast(expr.as_value(key).cast(target_type)) for key in expr
+        })
         self.set_output(casted)

src/mustela/translation/steps/concat.py

@@ -1,6 +1,11 @@
 """Translator for Concat and FeatureVectorizer operations."""
 
+import ibis
+
+import typing
+
 from ..translator import Translator
+from ..variables import VariablesGroup
 
 
 class ConcatTranslator(Translator):
@@ -30,13 +35,13 @@ def process(self) -> None:
         self.set_output(self._concatenate_columns(self))
 
     @classmethod
-    def _concatenate_columns(cls, translator: Translator) -> dict:
+    def _concatenate_columns(cls, translator: Translator) -> VariablesGroup:
         """Implement actual operation of concatenating columns.
 
         This is used by both Concat and FeatureVectorizer translators,
         as they both need to concatenate columns.
         """
-        result = {}
+        result = VariablesGroup()
         for col in translator.inputs:
             feature = translator._variables.consume(col)
             if isinstance(feature, dict):
@@ -47,8 +52,12 @@ def _concatenate_columns(cls, translator: Translator) -> dict:
                 for key in feature:
                     varname = col + "." + key
                     result[varname] = feature[key]
-            else:
+            elif isinstance(feature, ibis.Expr):
                 result[col] = feature
+            else:
+                raise ValueError(
+                    f"Concat: expected a column group or a single column. Got {type(feature)}"
+                )
 
         return result
 
@@ -68,7 +77,7 @@ def process(self) -> None:
 
         # We can support this by doing the same as Concat,
         # in most cases it's sufficient
-        ninputdimensions = self._attributes["inputdimensions"]
+        ninputdimensions = typing.cast(list[int], self._attributes["inputdimensions"])
 
         if len(ninputdimensions) != len(self._inputs):
             raise ValueError("Number of input dimensions should be equal to number of inputs.")

src/mustela/translation/steps/div.py

@@ -3,6 +3,8 @@
 
 import ibis
 
+from mustela.translation.variables import VariablesGroup
+
 from ..translator import Translator
 
 
@@ -42,29 +44,29 @@ def process(self) -> None:
                 if not isinstance(second_arg, (int, float)):
                     raise ValueError("Div: The second operand must be a numeric value.")
                 self.set_output(
-                    {
+                    VariablesGroup({
                         field: (
                             self._optimizer.fold_operation(
                                 first_operand[field] / ibis.literal(second_arg)
                             )
                         )
                         for field in struct_fields
-                    }
+                    })
                 )
             else:
                 if len(second_arg) != len(first_operand):
                     raise ValueError(
                         "The number of elements in the second operand must match the number of columns in the first operand."
                     )
                 self.set_output(
-                    {
+                    VariablesGroup({
                         field: (
                             self._optimizer.fold_operation(
                                 first_operand[field] / second_arg[i]
                             )
                         )
                         for i, field in enumerate(struct_fields)
-                    }
+                    })
                 )
         else:
             if not isinstance(first_operand, ibis.expr.types.NumericValue):

src/mustela/translation/steps/gather.py

@@ -1,5 +1,8 @@
 """Defines the translation step for the Gather operation."""
 
+from ibis.common.egraph import Variable
+
+from mustela.translation.variables import VariablesGroup
 from ..translator import Translator
 
 
@@ -38,7 +41,7 @@ def process(self) -> None:
         if not isinstance(idx, int):
             raise ValueError("Gather: index must be an integer constant")
 
-        if isinstance(expr, dict):
+        if isinstance(expr, VariablesGroup):
             keys = list(expr.keys())
             if idx < 0 or idx >= len(keys):
                 raise IndexError("Gather: index out of bounds")
@@ -48,6 +51,6 @@ def process(self) -> None:
             # support axis=1, then the index must be 0.
             if idx != 0:
                 raise NotImplementedError(
-                    f"Gather: index {idx} not supported for non-dict expression of type {type(expr)}"
+                    f"Gather: index {idx} not supported for single columns"
                 )
             self.set_output(expr)

src/mustela/translation/steps/imputer.py

@@ -2,6 +2,7 @@
 import ibis
 
 from ..translator import Translator
+from ..variables import VariablesGroup
 
 
 class ImputerTranslator(Translator):
@@ -23,13 +24,13 @@ def process(self) -> None:
             )
 
         expr = self._variables.consume(self.inputs[0])
-        if isinstance(expr, dict):
+        if isinstance(expr, VariablesGroup):
             keys = list(expr.keys())
             if len(keys) != len(imputed_values):
                 raise ValueError(
                     "Imputer: number of imputed values does not match number of columns"
                 )
-            new_expr = {}
+            new_expr = VariablesGroup()
             for i, key in enumerate(keys):
                 new_expr[key] = ibis.coalesce(expr[key], imputed_values[i])
             self.set_output(new_expr)

src/mustela/translation/steps/matmul.py

@@ -2,6 +2,7 @@
 import ibis
 
 from ..translator import Translator
+from ..variables import VariablesGroup
 
 
 class MatMulTranslator(Translator):
@@ -86,9 +87,9 @@ def process(self) -> None:
                     result = result_list[0]
                 else:
                     # Return a dict of output expressions if there are multiple output columns.
-                    result = {
+                    result = VariablesGroup({
                         f"out_{j}": result_list[j] for j in range(output_dim)
-                    }
+                    })
                 self.set_output(result)
             else:
                 raise NotImplementedError(
@@ -118,9 +119,9 @@ def process(self) -> None:
                         result = result_list[0]
                         self._variables[self._output_name] = result_list[0]
                     else:
-                        result = {
+                        result = VariablesGroup({
                             f"out_{j}": result_list[j] for j in range(output_dim)
-                        }
+                        })
                     self.set_output(result)
                 elif coef_shape[1] == 1:
                     # This case implies the left operand is a vector of length matching coef_shape[0],

src/mustela/translation/steps/onehotencoder.py

@@ -4,6 +4,7 @@
 import ibis
 
 from ..translator import Translator
+from ..variables import VariablesGroup
 
 
 class OneHotEncoderTranslator(Translator):
@@ -39,4 +40,4 @@ def process(self) -> None:
         # OneHot encoded features are usually consumed multiple times
         # by subsequent operations, so preserving them makes sense.
         casted_variables = self.preserve(*casted_variables)
-        self.set_output({cat: casted_variables[i] for i, cat in enumerate(cats)})
+        self.set_output(VariablesGroup({cat: casted_variables[i] for i, cat in enumerate(cats)}))

src/mustela/translation/steps/softmax.py

@@ -46,7 +46,7 @@ def compute_softmax(cls, data: ibis.expr.types.NumericValue | dict[str, ibis.exp
                 sum_exp = expr if sum_exp is None else sum_exp + expr
 
             # Multi columns case: softmax = exp(column_exp) / (exponents_sum)
-            softmax_result = {k: exp_dict[k] / sum_exp for k in data.keys()}
+            softmax_result = VariablesGroup({k: exp_dict[k] / sum_exp for k in data.keys()})
         else:
             # Single column case: softmax(x) = exp(x) / exp(x) = 1
             softmax_result = ibis.literal(1.0)

src/mustela/translation/steps/trees/__init__.py

@@ -1,3 +1,5 @@
+"""Translators for trees based models."""
+
 from .classifier import TreeEnsembleClassifierTranslator
 from .regressor import TreeEnsembleRegressorTranslator