Skip to content

Commit

Permalink
feat(api): support deferred expressions in ibis.date
Browse files Browse the repository at this point in the history
  • Loading branch information
jcrist authored and cpcloud committed Sep 27, 2023
1 parent f2b3744 commit f454a71
Show file tree
Hide file tree
Showing 2 changed files with 104 additions and 18 deletions.
96 changes: 79 additions & 17 deletions ibis/expr/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import functools
import numbers
import operator
from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar, overload

import ibis.expr.builders as bl
import ibis.expr.datatypes as dt
Expand All @@ -19,7 +19,7 @@
from ibis.common.exceptions import IbisInputError
from ibis.common.temporal import normalize_datetime, normalize_timezone
from ibis.expr.decompile import decompile
from ibis.expr.deferred import Deferred
from ibis.expr.deferred import Deferred, deferred_apply
from ibis.expr.schema import Schema
from ibis.expr.sql import parse_sql, show_sql, to_sql
from ibis.expr.streaming import Watermark
Expand Down Expand Up @@ -648,30 +648,92 @@ def timestamp(value, *args, timezone: str | None = None) -> ir.TimestampScalar:
return literal(value, type=dt.Timestamp(timezone=timezone))


def date(value, *args) -> DateValue:
"""Return a date literal if `value` is coercible to a date.
@overload
def date(
year: int | ir.IntegerValue | Deferred,
month: int | ir.IntegerValue | Deferred,
day: int | ir.IntegerValu | Deferred,
/,
) -> DateValue:
...


@overload
def date(value: Any, /) -> DateValue:
...


def date(value_or_year, month=None, day=None, /):
"""Construct a date scalar or column.
Parameters
----------
value
Date string, datetime object or numeric value
args
Month and day if `value` is a year
value_or_year
Either a string value or `datetime.date` to coerce to a date, or
an integral value representing the date year component.
month
The date month component; required if `value_or_year` is a year.
day
The date day component; required if `value_or_year` is a year.
Returns
-------
DateScalar
DateValue
A date expression
Examples
--------
>>> import ibis
>>> ibis.options.interactive = True
Create a date scalar from a string
>>> ibis.date("2023-01-02")
Timestamp('2023-01-02 00:00:00')
Create a date scalar from year, month, and day
>>> ibis.date(2023, 1, 2)
Timestamp('2023-01-02 00:00:00')
Create a date column from year, month, and day
>>> t = ibis.examples.airquality.fetch()
>>> ibis.date(1973, t.month, t.day).name("date")
┏━━━━━━━━━━━━┓
┃ date ┃
┡━━━━━━━━━━━━┩
│ date │
├────────────┤
│ 1973-05-01 │
│ 1973-05-02 │
│ 1973-05-03 │
│ 1973-05-04 │
│ 1973-05-05 │
│ 1973-05-06 │
│ 1973-05-07 │
│ 1973-05-08 │
│ 1973-05-09 │
│ 1973-05-10 │
│ … │
└────────────┘
"""
if isinstance(value, (numbers.Real, ir.IntegerValue)):
year, month, day = value, *args
return ops.DateFromYMD(year, month, day).to_expr()
elif isinstance(value, ir.StringValue):
return value.cast(dt.date)
elif isinstance(value, Deferred):
return value.date()
is_ymd = month is not None and day is not None
args = (value_or_year, month, day)

if any(isinstance(a, Deferred) for a in args):
return (
deferred_apply(date, *args)
if is_ymd
else deferred_apply(date, value_or_year)
)

if is_ymd:
return ops.DateFromYMD(value_or_year, month, day).to_expr()
elif isinstance(value_or_year, ir.StringValue):
return value_or_year.cast(dt.date)
else:
return literal(value, type=dt.date)
return literal(value_or_year, type=dt.date)


def time(value, *args) -> TimeValue:
Expand Down
26 changes: 25 additions & 1 deletion ibis/tests/expr/test_temporal.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import ibis.expr.datatypes as dt
import ibis.expr.operations as ops
import ibis.expr.types as ir
from ibis import _
from ibis.common.temporal import IntervalUnit
from ibis.expr import api
from ibis.tests.util import assert_equal
Expand Down Expand Up @@ -809,7 +810,30 @@ def test_time_truncate(table, operand, unit):
assert isinstance(expr.op(), ops.TimeTruncate)


def test_date_literal():
expr = ibis.date(2022, 2, 4)
sol = ops.DateFromYMD(2022, 2, 4).to_expr()
assert expr.equals(sol)

expr = ibis.date("2022-02-04")
sol = ibis.literal("2022-02-04", type=dt.date)
assert expr.equals(sol)


def test_date_expression():
t = ibis.table({"x": "int", "y": "int", "z": "int", "s": "string"})
deferred = ibis.date(_.x, _.y, _.z)
expr = ibis.date(t.x, t.y, t.z)
assert isinstance(expr.op(), ops.DateFromYMD)
assert deferred.resolve(t).equals(expr)
assert repr(deferred) == "date(_.x, _.y, _.z)"

deferred = ibis.date(_.s)
expr = ibis.date(t.s)
assert deferred.resolve(t).equals(expr)
assert repr(deferred) == "date(_.s)"


def test_date_time_literals():
assert ibis.date(2022, 2, 4).type() == dt.date
assert ibis.time(16, 20, 00).type() == dt.time
assert ibis.timestamp(2022, 2, 4, 16, 20, 00).type() == dt.timestamp

0 comments on commit f454a71

Please sign in to comment.