feat(checker): handle scalar ops on frames and columns

midas/checker/frames/column_methods.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import ast
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import midas.ast.python as p
 from midas.ast.location import Location
@@ -12,12 +12,10 @@ from midas.checker.types import (
     ColumnGroupBy,
     ColumnType,
     Function,
-    GenericType,
     OverloadedFunction,
     ParamSpec,
     TopType,
     Type,
-    TypeVar,
     UnitType,
     UnknownType,
     unfold_type,
@@ -65,7 +63,7 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         )
         return result.result
 
-    def _element_binary_op(self, call: Call, method: str) -> ColumnType:
+    def _element_binary_op(self, call: Call, method: str) -> Type:
         """Compute the result of an element-wise binary operation
 
         This function delegates to the inner types for computing the resulting
@@ -76,28 +74,31 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
             method (str): the method name
 
         Returns:
-            ColumnType: the resulting column type
+            Type: the resulting type
         """
-        column2: Optional[ColumnType] = None
+        if len(call.positional) == 0:
+            return UnknownType()
 
         col_type1: Type = call.column.type
-        new_column: Type = ColumnType(type=UnknownType())
+        operand: TypedExpr = call.positional[0]
-        if len(call.positional) != 0:
+        unfolded_operand: Type = unfold_type(operand[1])
-            other: Type = call.positional[0][1]
+        col_type2: Type
-            unfolded_other: Type = unfold_type(other)
-            if isinstance(unfolded_other, ColumnType):
-                column2 = unfolded_other
-                col_type2: Type = column2.type
 
-                new_inner_type = self.typer.result_of_binary_op(
+        # Operand is a column -> get the inner type
-                    location=call.location,
+        if isinstance(unfolded_operand, ColumnType):
-                    expr=call.call_expr,
+            col_type2 = unfolded_operand.type
-                    left=(call.column_expr, col_type1),
+        # Otherwise use the operand type itself
-                    right=(call.positional[0][0], col_type2),
+        else:
-                    method=method,
+            col_type2 = operand[1]
-                )
+
-                new_column = ColumnType(type=new_inner_type)
+        new_inner_type = self.typer.result_of_binary_op(
-        return new_column
+            location=call.location,
+            expr=call.call_expr,
+            left=(call.column_expr, col_type1),
+            right=(operand[0], col_type2),
+            method=method,
+        )
+        return ColumnType(type=new_inner_type)
 
     def _element_wise(self, call: Call, method: str) -> Type:
         """Compute the result of an element-wise method call
@@ -112,26 +113,20 @@ class ColumnMethodRegistry(MethodRegistry[Call]):
         Returns:
             Type: the result type
         """
-        # TODO: support add with scalar
 
         # Build signature with new column type and generic operand
-        param_type: TypeVar = TypeVar(name="T", bound=None)
+        signature = Function(
-        signature = GenericType(
+            params=ParamSpec(
-            name=method,
+                mixed=[
-            params=[param_type],
+                    Function.Parameter(
-            body=Function(
+                        pos=0,
-                params=ParamSpec(
+                        name="other",
-                    mixed=[
+                        type=TopType(),
-                        Function.Parameter(
+                        required=True,
-                            pos=0,
+                    ),
-                            name="other",
+                ],
-                            type=ColumnType(type=param_type),
-                            required=True,
-                        ),
-                    ],
-                ),
-                returns=self._element_binary_op(call, method),
             ),
+            returns=self._element_binary_op(call, method),
         )
 
         # Map arguments and compute result type

midas/checker/frames/frame_methods.py

@@ -102,7 +102,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
             return ColumnType(type=UnknownType())
         return result
 
-    def _element_binary_op(self, call: Call, method: str) -> DataFrameType:
+    def _element_binary_op(self, call: Call, method: str) -> Type:
         """Compute the result of an element-wise binary operation
 
         This function delegates to the matching columns for computing resulting
@@ -115,21 +115,22 @@ class FrameMethodRegistry(MethodRegistry[Call]):
             method (str): the method name
 
         Returns:
-            DataFrameType: the resulting frame type
+            Type: the resulting type
         """
+
+        if len(call.positional) == 0:
+            return UnknownType()
+
+        operand: TypedExpr = call.positional[0]
         new_columns: list[DataFrameType.Column] = []
 
         by_name: dict[str, DataFrameType.Column] = {}
         frame2: Optional[DataFrameType] = None
-        # Get map of operand's columns by name, if there is at least 1 operand, which is a dataframe
+        # Get map of operand's columns by name, if the operand is a dataframe
-        if len(call.positional) != 0:
+        unfolded_other: Type = unfold_type(operand[1])
-            operand: TypedExpr = call.positional[0]
+        if isinstance(unfolded_other, DataFrameType):
-            unfolded_other: Type = unfold_type(operand[1])
+            frame2 = unfolded_other
-            if isinstance(unfolded_other, DataFrameType):
+            by_name = {col.name: col for col in frame2.columns if col.name is not None}
-                frame2 = unfolded_other
-                by_name = {
-                    col.name: col for col in frame2.columns if col.name is not None
-                }
 
         # Compute new schema:
         # Step 1: for all columns in frame1:
@@ -142,10 +143,20 @@ class FrameMethodRegistry(MethodRegistry[Call]):
 
             col_type1: ColumnType = column.type
             col_type: ColumnType = ColumnType(type=UnknownType())
-            if column.name in by_name:
-                column2 = by_name[column.name]
-                col_type2: ColumnType = column2.type
 
+            col_type2: Optional[ColumnType] = None
+
+            # Operand is a frame -> lookup column with the same name
+            if frame2 is not None:
+                if column.name in by_name:
+                    column2 = by_name[column.name]
+                    col_type2 = column2.type
+
+            # Operand is not a frame -> scalar operation -> ad-hoc column
+            else:
+                col_type2 = ColumnType(type=operand[1])
+
+            if col_type2 is not None:
                 col_type = self._get_method_result(call, col_type1, col_type2, method)
 
             new_column = DataFrameType.Column(
@@ -184,7 +195,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
         Returns:
             Type: the result type
         """
-        # TODO: support scalar, sequence, Series, dict operand
+        # TODO: support sequence, Series, dict operand
         # Build signature with new schema and generic operand
         signature = Function(
             params=ParamSpec(
@@ -192,7 +203,7 @@ class FrameMethodRegistry(MethodRegistry[Call]):
                     Function.Parameter(
                         pos=0,
                         name="other",
-                        type=DataFrameType(columns=[]),
+                        type=TopType(),
                         required=True,
                     ),
                 ],