Skip to content

Commit c3f27a9

Browse files
kaijennissencpcloud
authored andcommitted
feat(api): add day_of_week.iso_index method for date and timestamp types
1 parent a9e4bbf commit c3f27a9

File tree

17 files changed

+126
-15
lines changed

17 files changed

+126
-15
lines changed

ibis/backends/polars/compiler.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1103,6 +1103,7 @@ def extract_epoch_seconds(op, **kw):
11031103
ops.Cos: operator.methodcaller("cos"),
11041104
ops.Cot: lambda arg: 1.0 / arg.tan(),
11051105
ops.DayOfWeekIndex: lambda arg: arg.dt.weekday() - 1,
1106+
ops.IsoDayOfWeekIndex: (lambda arg: arg.dt.weekday().cast(pl.Int16)),
11061107
ops.Exp: operator.methodcaller("exp"),
11071108
ops.Floor: lambda arg: arg.floor().cast(pl.Int64),
11081109
ops.IsInf: operator.methodcaller("is_infinite"),

ibis/backends/sql/compilers/base.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -965,6 +965,9 @@ def visit_TimeTruncate(self, op, *, arg, unit):
965965
def visit_DayOfWeekIndex(self, op, *, arg):
966966
return (self.f.dayofweek(arg) + 6) % 7
967967

968+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
969+
return ((self.f.dayofweek(arg) + 6) % 7) + 1
970+
968971
def visit_DayOfWeekName(self, op, *, arg):
969972
# day of week number is 0-indexed
970973
# Sunday == 0

ibis/backends/sql/compilers/bigquery/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,9 @@ def visit_StringJoin(self, op, *, arg, sep):
464464
def visit_DayOfWeekIndex(self, op, *, arg):
465465
return self.f.mod(self.f.extract(self.v.dayofweek, arg) + 5, 7)
466466

467+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
468+
return self.f.mod(self.f.extract(self.v.dayofweek, arg) + 5, 7) + 1
469+
467470
def visit_DayOfWeekName(self, op, *, arg):
468471
return self.f.initcap(sge.Cast(this=arg, to="STRING FORMAT 'DAY'"))
469472

ibis/backends/sql/compilers/clickhouse.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -504,6 +504,9 @@ def visit_DayOfWeekIndex(self, op, *, arg):
504504
weekdays = len(calendar.day_name)
505505
return (((self.f.toDayOfWeek(arg) - 1) % weekdays) + weekdays) % weekdays
506506

507+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
508+
return self.f.toDayOfWeek(arg)
509+
507510
def visit_DayOfWeekName(self, op, *, arg):
508511
# ClickHouse 20 doesn't support dateName
509512
#

ibis/backends/sql/compilers/datafusion.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,9 @@ def visit_ExtractDayOfYear(self, op, *, arg):
258258
def visit_DayOfWeekIndex(self, op, *, arg):
259259
return (self.f.date_part("dow", arg) + 6) % 7
260260

261+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
262+
return ((self.f.date_part("dow", arg) + 6) % 7) + 1
263+
261264
def visit_DayOfWeekName(self, op, *, arg):
262265
return sg.exp.Case(
263266
this=sge.paren(self.f.date_part("dow", arg) + 6, copy=False) % 7,

ibis/backends/sql/compilers/exasol.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@ class ExasolCompiler(SQLGlotCompiler):
5050
ops.DateAdd,
5151
ops.DateSub,
5252
ops.DateFromYMD,
53-
ops.DayOfWeekIndex,
5453
ops.ElementWiseVectorizedUDF,
5554
ops.IntervalFromInteger,
5655
ops.IsInf,
@@ -257,6 +256,12 @@ def visit_ExtractDayOfYear(self, op, *, arg):
257256
def visit_ExtractWeekOfYear(self, op, *, arg):
258257
return self.cast(self.f.to_char(arg, "IW"), op.dtype)
259258

259+
def visit_DayOfWeekIndex(self, op, *, arg):
260+
return self.cast(self.f.to_char(arg, "ID"), op.dtype) - 1
261+
262+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
263+
return self.cast(self.f.to_char(arg, "ID"), op.dtype)
264+
260265
def visit_ExtractIsoYear(self, op, *, arg):
261266
return self.cast(self.f.to_char(arg, "IYYY"), op.dtype)
262267

ibis/backends/sql/compilers/flink.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,9 @@ def visit_ExtractMicrosecond(self, op, *, arg):
451451
def visit_DayOfWeekIndex(self, op, *, arg):
452452
return (self.f.anon.dayofweek(arg) + 5) % 7
453453

454+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
455+
return ((self.f.dayofweek(arg) + 5) % 7) + 1
456+
454457
def visit_DayOfWeekName(self, op, *, arg):
455458
index = self.cast(self.f.anon.dayofweek(self.cast(arg, dt.date)), op.dtype)
456459
lookup_table = self.f.anon.str_to_map(

ibis/backends/sql/compilers/impala.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,9 @@ def visit_RandomScalar(self, op):
124124
def visit_DayOfWeekIndex(self, op, *, arg):
125125
return self.f.pmod(self.f.dayofweek(arg) - 2, 7)
126126

127+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
128+
return self.f.pmod(self.f.dayofweek(arg) - 2, 7) + 1
129+
127130
def visit_ExtractMillisecond(self, op, *, arg):
128131
return self.f.extract(self.v.millisecond, arg) % 1_000
129132

ibis/backends/sql/compilers/mssql.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,9 @@ def visit_ApproxQuantile(self, op, *, arg, quantile, where):
240240
def visit_DayOfWeekIndex(self, op, *, arg):
241241
return self.f.datepart(self.v.weekday, arg) - 1
242242

243+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
244+
return self.f.datepart(self.v.weekday, arg)
245+
243246
def visit_DayOfWeekName(self, op, *, arg):
244247
days = calendar.day_name
245248
return sge.Case(

ibis/backends/sql/compilers/mysql.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,9 @@ def visit_GroupConcat(self, op, *, arg, sep, where, order_by):
166166
def visit_DayOfWeekIndex(self, op, *, arg):
167167
return (self.f.dayofweek(arg) + 5) % 7
168168

169+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
170+
return ((self.f.dayofweek(arg) + 5) % 7) + 1
171+
169172
def visit_Literal(self, op, *, value, dtype):
170173
# avoid casting NULL: the set of types allowed by MySQL and
171174
# MariaDB when casting is a strict subset of allowed types in other

ibis/backends/sql/compilers/postgres.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -623,6 +623,9 @@ def visit_TimestampBucket(self, op, *, arg, interval, offset):
623623
def visit_DayOfWeekIndex(self, op, *, arg):
624624
return self.cast(self.f.extract("dow", arg) + 6, dt.int16) % 7
625625

626+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
627+
return self.cast(self.f.extract("isodow", arg), dt.int16)
628+
626629
def visit_DayOfWeekName(self, op, *, arg):
627630
return self.f.trim(self.f.to_char(arg, "Day"), string.whitespace)
628631

ibis/backends/sql/compilers/pyspark.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,9 @@ def visit_IntervalFromInteger(self, op, *, arg, unit):
177177
def visit_DayOfWeekIndex(self, op, *, arg):
178178
return (self.f.dayofweek(arg) + 5) % 7
179179

180+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
181+
return ((self.f.dayofweek(arg) + 5) % 7) + 1
182+
180183
def visit_DayOfWeekName(self, op, *, arg):
181184
return sge.Case(
182185
this=(self.f.dayofweek(arg) + 5) % 7,

ibis/backends/sql/compilers/sqlite.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -473,6 +473,16 @@ def visit_DayOfWeekIndex(self, op, *, arg):
473473
self.f.mod(self.cast(self.f.strftime("%w", arg) + 6, dt.int64), 7), dt.int64
474474
)
475475

476+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
477+
# return self.cast(self.f.strftime("%u", arg), dt.int64)
478+
return (
479+
self.cast(
480+
self.f.mod(self.cast(self.f.strftime("%w", arg) + 6, dt.int64), 7),
481+
dt.int64,
482+
)
483+
+ 1
484+
)
485+
476486
def visit_DayOfWeekName(self, op, *, arg):
477487
return sge.Case(
478488
this=self.f.strftime("%w", arg),

ibis/backends/sql/compilers/trino.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,9 @@ def visit_UnwrapJSONBoolean(self, op, *, arg):
276276
def visit_DayOfWeekIndex(self, op, *, arg):
277277
return self.cast(sge.paren(self.f.anon.dow(arg) + 6, copy=False) % 7, op.dtype)
278278

279+
def visit_IsoDayOfWeekIndex(self, op, *, arg):
280+
return self.cast(sge.paren(self.f.day_of_week(arg), copy=False), op.dtype)
281+
279282
def visit_DayOfWeekName(self, op, *, arg):
280283
return self.f.date_format(arg, "%W")
281284

ibis/backends/tests/test_temporal.py

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ def test_iso_year_does_not_match_date_year(con):
200200
id="day_of_week_index",
201201
marks=[
202202
pytest.mark.notimpl(
203-
["druid", "oracle", "exasol"], raises=com.OperationNotDefinedError
203+
["druid", "oracle"], raises=com.OperationNotDefinedError
204204
),
205205
],
206206
),
@@ -1331,25 +1331,35 @@ def test_string_as_time(backend, alltypes):
13311331

13321332

13331333
@pytest.mark.parametrize(
1334-
("date", "expected_index", "expected_day"),
1334+
("date", "expected_index", "expected_iso_index", "expected_day"),
13351335
[
1336-
param("2017-01-01", 6, "Sunday", id="sunday"),
1337-
param("2017-01-02", 0, "Monday", id="monday"),
1338-
param("2017-01-03", 1, "Tuesday", id="tuesday"),
1339-
param("2017-01-04", 2, "Wednesday", id="wednesday"),
1340-
param("2017-01-05", 3, "Thursday", id="thursday"),
1341-
param("2017-01-06", 4, "Friday", id="friday"),
1342-
param("2017-01-07", 5, "Saturday", id="saturday"),
1336+
param("2017-01-01", 6, 7, "Sunday", id="sunday"),
1337+
param("2017-01-02", 0, 1, "Monday", id="monday"),
1338+
param("2017-01-03", 1, 2, "Tuesday", id="tuesday"),
1339+
param("2017-01-04", 2, 3, "Wednesday", id="wednesday"),
1340+
param("2017-01-05", 3, 4, "Thursday", id="thursday"),
1341+
param("2017-01-06", 4, 5, "Friday", id="friday"),
1342+
param("2017-01-07", 5, 6, "Saturday", id="saturday"),
13431343
],
13441344
)
13451345
@pytest.mark.notimpl(["druid", "oracle"], raises=com.OperationNotDefinedError)
13461346
@pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
13471347
@mark_notyet_risingwave_14670
1348-
def test_day_of_week_scalar(con, date, expected_index, expected_day):
1348+
@pytest.mark.broken(
1349+
["risingwave"],
1350+
raises=AssertionError,
1351+
reason="Refer to https://github.com/risingwavelabs/risingwave/issues/14670",
1352+
)
1353+
def test_day_of_week_scalar(
1354+
con, date, expected_index, expected_iso_index, expected_day
1355+
):
13491356
expr = ibis.literal(date).cast(dt.date)
13501357
result_index = con.execute(expr.day_of_week.index().name("tmp"))
13511358
assert result_index == expected_index
13521359

1360+
result_iso_index = con.execute(expr.day_of_week.iso_index().name("tmp"))
1361+
assert result_iso_index == expected_iso_index
1362+
13531363
result_day = con.execute(expr.day_of_week.full_name().name("tmp"))
13541364
assert result_day.lower() == expected_day.lower()
13551365

@@ -1364,6 +1374,11 @@ def test_day_of_week_column(backend, alltypes, df):
13641374

13651375
backend.assert_series_equal(result_index, expected_index, check_names=False)
13661376

1377+
result_iso_index = expr.iso_index().name("tmp").execute()
1378+
expected_iso_index = df.timestamp_col.dt.isocalendar().day.astype("int16")
1379+
1380+
backend.assert_series_equal(result_iso_index, expected_iso_index, check_names=False)
1381+
13671382
result_day = expr.full_name().name("tmp").execute()
13681383
expected_day = df.timestamp_col.dt.day_name()
13691384

@@ -1377,9 +1392,6 @@ def test_day_of_week_column(backend, alltypes, df):
13771392
lambda t: t.timestamp_col.day_of_week.index().count(),
13781393
lambda s: s.dt.dayofweek.count(),
13791394
id="day_of_week_index",
1380-
marks=[
1381-
pytest.mark.notimpl(["exasol"], raises=com.OperationNotDefinedError)
1382-
],
13831395
),
13841396
param(
13851397
lambda t: t.timestamp_col.day_of_week.full_name().length().sum(),

ibis/expr/operations/temporal.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -201,6 +201,13 @@ class DayOfWeekIndex(Unary):
201201
dtype = dt.int16
202202

203203

204+
@public
205+
class IsoDayOfWeekIndex(Unary):
206+
arg: Value[dt.Date | dt.Timestamp]
207+
208+
dtype = dt.int16
209+
210+
204211
@public
205212
class DayOfWeekName(Unary):
206213
"""Extract the name of the day of the week from a date or timestamp."""
@@ -210,6 +217,13 @@ class DayOfWeekName(Unary):
210217
dtype = dt.string
211218

212219

220+
@public
221+
class IsoDayOfWeekName(Unary):
222+
arg: Value[dt.Date | dt.Timestamp]
223+
224+
dtype = dt.string
225+
226+
213227
@public
214228
class Time(Unary):
215229
"""Extract the time from a timestamp."""

ibis/expr/types/temporal.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,34 @@ def day_of_week(self) -> DayOfWeek:
103103
Returns
104104
-------
105105
DayOfWeek
106-
An namespace expression containing methods to use to extract
106+
A namespace expression with methods for extracting day-of-week
107107
information.
108+
109+
Examples
110+
--------
111+
>>> import ibis
112+
>>> import datetime as dt
113+
>>> from ibis import _
114+
>>> ibis.options.interactive = True
115+
>>> t = ibis.memtable({"date": [dt.datetime(2024, 4, x) for x in range(14, 21)]})
116+
>>> t.mutate(
117+
... day_of_week=_.date.day_of_week.index(),
118+
... day_of_week_name=_.date.day_of_week.full_name(),
119+
... iso_day_of_week=_.date.day_of_week.iso_index(),
120+
... )
121+
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
122+
┃ date ┃ day_of_week ┃ day_of_week_name ┃ iso_day_of_week ┃
123+
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
124+
│ timestamp │ int16 │ string │ int16 │
125+
├─────────────────────┼─────────────┼──────────────────┼─────────────────┤
126+
│ 2024-04-14 00:00:00 │ 6 │ Sunday │ 7 │
127+
│ 2024-04-15 00:00:00 │ 0 │ Monday │ 1 │
128+
│ 2024-04-16 00:00:00 │ 1 │ Tuesday │ 2 │
129+
│ 2024-04-17 00:00:00 │ 2 │ Wednesday │ 3 │
130+
│ 2024-04-18 00:00:00 │ 3 │ Thursday │ 4 │
131+
│ 2024-04-19 00:00:00 │ 4 │ Friday │ 5 │
132+
│ 2024-04-20 00:00:00 │ 5 │ Saturday │ 6 │
133+
└─────────────────────┴─────────────┴──────────────────┴─────────────────┘
108134
"""
109135
return DayOfWeek(self)
110136

@@ -1497,3 +1523,13 @@ def full_name(self):
14971523
└─────────────────────────┘
14981524
"""
14991525
return ops.DayOfWeekName(self._expr).to_expr()
1526+
1527+
def iso_index(self):
1528+
"""Get the index of the day of the week in iso-format (1=Monday, 7=Sunday).
1529+
1530+
Returns
1531+
-------
1532+
IntegerValue
1533+
The index of the day of the week in iso-format (1=Monday, 7=Sunday).
1534+
"""
1535+
return ops.IsoDayOfWeekIndex(self._expr).to_expr()

0 commit comments

Comments
 (0)