diff options
Diffstat (limited to 'sqlglot/dataframe')
| -rw-r--r-- | sqlglot/dataframe/sql/column.py | 6 | ||||
| -rw-r--r-- | sqlglot/dataframe/sql/functions.py | 20 |
2 files changed, 14 insertions, 12 deletions
diff --git a/sqlglot/dataframe/sql/column.py b/sqlglot/dataframe/sql/column.py index 40ffe3e..f5b0974 100644 --- a/sqlglot/dataframe/sql/column.py +++ b/sqlglot/dataframe/sql/column.py @@ -260,11 +260,7 @@ class Column: """ if isinstance(dataType, DataType): dataType = dataType.simpleString() - new_expression = exp.Cast( - this=self.column_expression, - to=sqlglot.parse_one(dataType, into=exp.DataType, read="spark"), # type: ignore - ) - return Column(new_expression) + return Column(exp.cast(self.column_expression, dataType, dialect="spark")) def startswith(self, value: t.Union[str, Column]) -> Column: value = self._lit(value) if not isinstance(value, Column) else value diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py index a141fe4..47d5e7b 100644 --- a/sqlglot/dataframe/sql/functions.py +++ b/sqlglot/dataframe/sql/functions.py @@ -536,15 +536,15 @@ def month(col: ColumnOrName) -> Column: def dayofweek(col: ColumnOrName) -> Column: - return Column.invoke_anonymous_function(col, "DAYOFWEEK") + return Column.invoke_expression_over_column(col, glotexp.DayOfWeek) def dayofmonth(col: ColumnOrName) -> Column: - return Column.invoke_anonymous_function(col, "DAYOFMONTH") + return Column.invoke_expression_over_column(col, glotexp.DayOfMonth) def dayofyear(col: ColumnOrName) -> Column: - return Column.invoke_anonymous_function(col, "DAYOFYEAR") + return Column.invoke_expression_over_column(col, glotexp.DayOfYear) def hour(col: ColumnOrName) -> Column: @@ -560,7 +560,7 @@ def second(col: ColumnOrName) -> Column: def weekofyear(col: ColumnOrName) -> Column: - return Column.invoke_anonymous_function(col, "WEEKOFYEAR") + return Column.invoke_expression_over_column(col, glotexp.WeekOfYear) def make_date(year: ColumnOrName, month: ColumnOrName, day: ColumnOrName) -> Column: @@ -1144,10 +1144,16 @@ def aggregate( merge_exp = _get_lambda_from_func(merge) if finish is not None: finish_exp = _get_lambda_from_func(finish) - return Column.invoke_anonymous_function( - col, "AGGREGATE", initialValue, Column(merge_exp), Column(finish_exp) + return Column.invoke_expression_over_column( + col, + glotexp.Reduce, + initial=initialValue, + merge=Column(merge_exp), + finish=Column(finish_exp), ) - return Column.invoke_anonymous_function(col, "AGGREGATE", initialValue, Column(merge_exp)) + return Column.invoke_expression_over_column( + col, glotexp.Reduce, initial=initialValue, merge=Column(merge_exp) + ) def transform( |