zero323 commented on a change in pull request #35032:
URL: https://github.com/apache/spark/pull/35032#discussion_r775502742
##########
File path: python/pyspark/sql/functions.py
##########
@@ -2200,38 +2200,70 @@ def make_date(year: "ColumnOrName", month:
"ColumnOrName", day: "ColumnOrName")
return Column(jc)
+@overload
+def date_add(start: "ColumnOrName", days: "ColumnOrName") -> Column:
+ ...
+
+
+@overload
def date_add(start: "ColumnOrName", days: int) -> Column:
+ ...
+
+
+def date_add(start: "ColumnOrName", days: Union["ColumnOrName", int]) ->
Column:
"""
Returns the date that is `days` days after `start`
.. versionadded:: 1.5.0
Examples
--------
- >>> df = spark.createDataFrame([('2015-04-08',)], ['dt'])
+ >>> df = spark.createDataFrame([('2015-04-08', 2,)], ['dt', 'add'])
>>> df.select(date_add(df.dt, 1).alias('next_date')).collect()
[Row(next_date=datetime.date(2015, 4, 9))]
+ >>> df.select(date_add(df.dt,
df.add.cast(IntegerType())).alias('next_date')).collect()
+ [Row(next_date=datetime.date(2015, 4, 10))]
"""
sc = SparkContext._active_spark_context
assert sc is not None and sc._jvm is not None
- return Column(sc._jvm.functions.date_add(_to_java_column(start), days))
+ return Column(
+ sc._jvm.functions.date_add(
+ _to_java_column(start), days if isinstance(days, int) else
_to_java_column(days)
Review comment:
Alternatively, with `_create_column_from_literal`
```python
_create_column_from_literal
days = days._jc if isinstance(days, Column) else
_create_column_from_literal(days)
return sc._jvm.functions.date_add(_to_java_column(start), days)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
