diff --git a/ibis/backends/sqlite/registry.py b/ibis/backends/sqlite/registry.py index bb1d7fe686bad..f37d8ebdfb20c 100644 --- a/ibis/backends/sqlite/registry.py +++ b/ibis/backends/sqlite/registry.py @@ -4,6 +4,7 @@ import sqlalchemy as sa import toolz from multipledispatch import Dispatcher +import functools import ibis import ibis.common.exceptions as com @@ -302,9 +303,16 @@ def _string_concat(t, expr): return functools.reduce(operator.add, map(t.translate, args)) +def _date_from_ymd(t, expr): + y, m, d = map(t.translate, expr.op().args) + ymdstr = sa.func.printf('%04d-%02d-%02d', y, m, d) + return sa.func.date(ymdstr) + + operation_registry.update( { ops.Cast: _cast, + ops.DateFromYMD: _date_from_ymd, ops.Substring: _substr, ops.StrRight: _string_right, ops.StringFind: _string_find, @@ -317,6 +325,7 @@ def _string_concat(t, expr): ops.Date: unary(sa.func.date), ops.TimestampTruncate: _truncate(sa.func.datetime), ops.Strftime: _strftime, + ops.StringConcat: _string_concat, ops.ExtractYear: _strftime_int('%Y'), ops.ExtractMonth: _strftime_int('%m'), ops.ExtractDay: _strftime_int('%d'), diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py index a49aaa48e4de4..f73bd087fb20e 100644 --- a/ibis/backends/tests/test_temporal.py +++ b/ibis/backends/tests/test_temporal.py @@ -622,3 +622,49 @@ def test_now_from_projection(backend, alltypes): now = pd.Timestamp('now') year_expected = pd.Series([now.year] * n, name='ts') tm.assert_series_equal(ts.dt.year, year_expected) + + +def test_date_literal_ex(con): + expr = ibis.date(2022, 2, 4) + result = con.execute(expr) + assert result.strftime('%Y-%m-%d') == '2022-02-04' + + +def test_date_column_from_ymd(con, alltypes, df): + c = alltypes.timestamp_col + expr = ibis.date(c.year(), c.month(), c.day()) + tbl = alltypes[ + expr.name('timestamp_col'), + ] + result = con.execute(tbl) + + golden = df.timestamp_col.dt.date.astype('datetime64[ns]') + tm.assert_series_equal(golden, result.timestamp_col) + + +def test_date_scalar_from_iso(con): + expr = ibis.literal('2022-02-24') + expr2 = ibis.date(expr) + + result = con.execute(expr2) + assert result.strftime('%Y-%m-%d') == '2022-02-24' + + +def test_date_column_from_iso(con, alltypes, df): + expr = ( + alltypes.year.cast('string') + + '-' + + alltypes.month.cast('string').lpad(2, '0') + + '-13' + ) + expr = ibis.date(expr) + + result = con.execute(expr) + golden = ( + df.year.astype(str) + + '-' + + df.month.astype(str).str.rjust(2, '0') + + '-13' + ) + actual = result.dt.strftime('%Y-%m-%d') + tm.assert_series_equal(golden.rename('tmp'), actual) diff --git a/ibis/expr/api.py b/ibis/expr/api.py index 8be92a749c5d4..7878433babebb 100644 --- a/ibis/expr/api.py +++ b/ibis/expr/api.py @@ -405,7 +405,13 @@ def timestamp( return literal(value, type=dt.Timestamp(timezone=timezone)) -def date(value: str) -> ir.DateScalar: +@functools.singledispatch +def date(value) -> DateValue: + raise Exception('notimpl') + + +@date.register(str) +def _(value: str) -> ir.DateScalar: """Return a date literal if `value` is coercible to a date. Parameters @@ -423,6 +429,17 @@ def date(value: str) -> ir.DateScalar: return literal(value, type=dt.date) +@date.register(IntegerColumn) +@date.register(int) +def _(year, month, day) -> ir.DateScalar: + return ops.DateFromYMD(year, month, day).to_expr() + + +@date.register(StringValue) +def _(value: StringValue) -> DateValue: + return value.cast(dt.date) + + def time(value: str) -> ir.TimeScalar: """Return a time literal if `value` is coercible to a time. diff --git a/ibis/expr/operations/temporal.py b/ibis/expr/operations/temporal.py index 692573627b6d3..c993ee0f8839d 100644 --- a/ibis/expr/operations/temporal.py +++ b/ibis/expr/operations/temporal.py @@ -213,6 +213,14 @@ class Date(UnaryOp): output_type = rlz.shape_like('arg', dt.date) +@public +class DateFromYMD(ValueOp): + year = rlz.integer + month = rlz.integer + day = rlz.integer + output_type = rlz.shape_like('args', dt.date) + + @public class TimestampFromUNIX(ValueOp): arg = rlz.any diff --git a/ibis/tests/expr/test_temporal.py b/ibis/tests/expr/test_temporal.py index daf0ec8628007..33f1b634f3cb7 100644 --- a/ibis/tests/expr/test_temporal.py +++ b/ibis/tests/expr/test_temporal.py @@ -705,3 +705,7 @@ def test_time_truncate(table, operand, unit): expr = operand(table).truncate(unit) assert isinstance(expr, ir.TimeValue) assert isinstance(expr.op(), ops.TimeTruncate) + + +def test_date_literal(): + ibis.date(2022, 2, 4)