From 1098e33d07cbeba51d85997edd20060b191f4b79 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Tue, 7 Jul 2026 10:27:20 +0200 Subject: [PATCH] feat(checker): handle scalar ops on frames and columns --- midas/checker/frames/column_methods.py | 73 ++++++++++++-------------- midas/checker/frames/frame_methods.py | 43 +++++++++------ 2 files changed, 61 insertions(+), 55 deletions(-) diff --git a/midas/checker/frames/column_methods.py b/midas/checker/frames/column_methods.py index d370398..80792d0 100644 --- a/midas/checker/frames/column_methods.py +++ b/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()) - if len(call.positional) != 0: - other: Type = call.positional[0][1] - unfolded_other: Type = unfold_type(other) - if isinstance(unfolded_other, ColumnType): - column2 = unfolded_other - col_type2: Type = column2.type + operand: TypedExpr = call.positional[0] + unfolded_operand: Type = unfold_type(operand[1]) + col_type2: Type - new_inner_type = self.typer.result_of_binary_op( - location=call.location, - expr=call.call_expr, - left=(call.column_expr, col_type1), - right=(call.positional[0][0], col_type2), - method=method, - ) - new_column = ColumnType(type=new_inner_type) - return new_column + # Operand is a column -> get the inner type + if isinstance(unfolded_operand, ColumnType): + col_type2 = unfolded_operand.type + # Otherwise use the operand type itself + else: + col_type2 = operand[1] + + new_inner_type = self.typer.result_of_binary_op( + 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 = GenericType( - name=method, - params=[param_type], - body=Function( - params=ParamSpec( - mixed=[ - Function.Parameter( - pos=0, - name="other", - type=ColumnType(type=param_type), - required=True, - ), - ], - ), - returns=self._element_binary_op(call, method), + signature = Function( + params=ParamSpec( + mixed=[ + Function.Parameter( + pos=0, + name="other", + type=TopType(), + required=True, + ), + ], ), + returns=self._element_binary_op(call, method), ) # Map arguments and compute result type diff --git a/midas/checker/frames/frame_methods.py b/midas/checker/frames/frame_methods.py index c5be389..7d3d991 100644 --- a/midas/checker/frames/frame_methods.py +++ b/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 - if len(call.positional) != 0: - operand: TypedExpr = call.positional[0] - unfolded_other: Type = unfold_type(operand[1]) - if isinstance(unfolded_other, DataFrameType): - frame2 = unfolded_other - by_name = { - col.name: col for col in frame2.columns if col.name is not None - } + # Get map of operand's columns by name, if the operand is a dataframe + unfolded_other: Type = unfold_type(operand[1]) + if isinstance(unfolded_other, DataFrameType): + 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, ), ],