HEL/midas
1
0
Fork 0
Code Issues Pull Requests Actions Releases Activity

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

Browse Source
This commit is contained in:
Louis Heredero
2026-07-07 10:27:20 +02:00
parent 9277bd2cd0
commit 1098e33d07
2 changed files with 61 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
midas/checker/frames/column_methods.py
View File

@@ -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

43
midas/checker/frames/frame_methods.py
View File

@@ -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,
),
],