|
25 | 25 | from vgi.scalar_function import ScalarFunction |
26 | 26 |
|
27 | 27 | __all__ = [ |
28 | | - "DoubleColumnFunction", |
29 | 28 | "AddNumericColumnsFunction", |
| 29 | + "DoubleColumnFunction", |
| 30 | + "SumColumnsFunction", |
30 | 31 | "UpperCaseFunction", |
31 | 32 | ] |
32 | 33 |
|
@@ -225,3 +226,65 @@ def catalog_output_type(cls) -> pa.DataType | type[AnyArrow]: |
225 | 226 | def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]: |
226 | 227 | """Convert the column values to uppercase.""" |
227 | 228 | return pc.utf8_upper(batch.column(self.column)) # type: ignore[no-matching-overload] |
| 229 | + |
| 230 | + |
| 231 | +class SumColumnsFunction(ScalarFunction): |
| 232 | + """Sums values from multiple numeric columns. |
| 233 | +
|
| 234 | + Uses varargs with type_bound to accept any number of numeric columns |
| 235 | + and validates that all columns are addable types at bind time. |
| 236 | +
|
| 237 | + Example: |
| 238 | + Input: a=[1, 2], b=[10, 20], c=[100, 200] |
| 239 | + Args: columns=('a', 'b', 'c') |
| 240 | + Output: result=[111, 222] |
| 241 | +
|
| 242 | + """ |
| 243 | + |
| 244 | + class Meta: |
| 245 | + """Function metadata.""" |
| 246 | + |
| 247 | + name = "sum_columns" |
| 248 | + description = "Sum values from multiple numeric columns" |
| 249 | + examples = [ |
| 250 | + FunctionExample( |
| 251 | + sql="SELECT sum_columns(price, tax, shipping) FROM orders", |
| 252 | + description="Calculate total cost from multiple columns", |
| 253 | + ), |
| 254 | + ] |
| 255 | + |
| 256 | + # Varargs with type_bound validates all columns are numeric |
| 257 | + columns = Arg[AnyArrow]( |
| 258 | + 0, |
| 259 | + varargs=True, |
| 260 | + type_bound=_is_addable_type, |
| 261 | + doc="Columns to sum (must be numeric)", |
| 262 | + ) |
| 263 | + |
| 264 | + _output_type: pa.DataType |
| 265 | + |
| 266 | + def bind(self) -> None: |
| 267 | + """Compute output type from first column, promoted for overflow safety.""" |
| 268 | + # With varargs=True, self.columns is a tuple of column names |
| 269 | + first_col = self.columns[0] # type: ignore[index] |
| 270 | + first_type = self.input_schema.field(first_col).type |
| 271 | + self._output_type = _promote_for_addition(first_type) |
| 272 | + |
| 273 | + @classmethod |
| 274 | + def catalog_output_type(cls) -> pa.DataType | type[AnyArrow]: |
| 275 | + """Output type depends on input column types.""" |
| 276 | + return AnyArrow |
| 277 | + |
| 278 | + @property |
| 279 | + def output_type(self) -> pa.DataType: |
| 280 | + """Return the computed output type based on first column.""" |
| 281 | + return self._output_type |
| 282 | + |
| 283 | + def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]: |
| 284 | + """Sum values from all specified columns.""" |
| 285 | + # With varargs=True, self.columns is a tuple of column names at runtime |
| 286 | + columns: tuple[str, ...] = self.columns # type: ignore[assignment] |
| 287 | + result = batch.column(columns[0]) |
| 288 | + for col_name in columns[1:]: |
| 289 | + result = pc.add(result, batch.column(col_name)) |
| 290 | + return result |
0 commit comments