From 4707c44a8ed58d5d1c3297f001914185c29045d7 Mon Sep 17 00:00:00 2001 From: kj3 <50238578+kaijennissen@users.noreply.github.com> Date: Wed, 8 May 2024 17:12:38 +0200 Subject: [PATCH] feat(api): isoyear method (#9034) --- ibis/backends/bigquery/compiler.py | 3 +++ ibis/backends/clickhouse/compiler.py | 1 + ibis/backends/duckdb/compiler.py | 1 + ibis/backends/exasol/compiler.py | 3 +++ ibis/backends/oracle/compiler.py | 3 +++ ibis/backends/pandas/kernels.py | 1 + ibis/backends/polars/compiler.py | 1 + ibis/backends/postgres/compiler.py | 3 +++ ibis/backends/snowflake/compiler.py | 1 + ibis/backends/tests/test_temporal.py | 40 ++++++++++++++++++++++++++++ ibis/backends/trino/compiler.py | 1 + ibis/expr/operations/temporal.py | 5 ++++ ibis/expr/types/temporal.py | 4 +++ 13 files changed, 67 insertions(+) diff --git a/ibis/backends/bigquery/compiler.py b/ibis/backends/bigquery/compiler.py index b32a6476c48f..633d12708861 100644 --- a/ibis/backends/bigquery/compiler.py +++ b/ibis/backends/bigquery/compiler.py @@ -391,6 +391,9 @@ def visit_ExtractEpochSeconds(self, op, *, arg): def visit_ExtractWeekOfYear(self, op, *, arg): return self.f.extract(self.v.isoweek, arg) + def visit_ExtractIsoYear(self, op, *, arg): + return self.f.extract(self.v.isoyear, arg) + def visit_ExtractMillisecond(self, op, *, arg): return self.f.extract(self.v.millisecond, arg) diff --git a/ibis/backends/clickhouse/compiler.py b/ibis/backends/clickhouse/compiler.py index c99635b5fac2..9b01beef6dce 100644 --- a/ibis/backends/clickhouse/compiler.py +++ b/ibis/backends/clickhouse/compiler.py @@ -72,6 +72,7 @@ class ClickHouseCompiler(SQLGlotCompiler): ops.ExtractSecond: "toSecond", ops.ExtractWeekOfYear: "toISOWeek", ops.ExtractYear: "toYear", + ops.ExtractIsoYear: "toISOYear", ops.First: "any", ops.IntegerRange: "range", ops.IsInf: "isInfinite", diff --git a/ibis/backends/duckdb/compiler.py b/ibis/backends/duckdb/compiler.py index d17ba9360ba4..9655086cc2d6 100644 --- a/ibis/backends/duckdb/compiler.py +++ b/ibis/backends/duckdb/compiler.py @@ -45,6 +45,7 @@ class DuckDBCompiler(SQLGlotCompiler): ops.BitOr: "bit_or", ops.BitXor: "bit_xor", ops.EndsWith: "suffix", + ops.ExtractIsoYear: "isoyear", ops.Hash: "hash", ops.IntegerRange: "range", ops.TimestampRange: "range", diff --git a/ibis/backends/exasol/compiler.py b/ibis/backends/exasol/compiler.py index de06d99b03cb..0940c80a182e 100644 --- a/ibis/backends/exasol/compiler.py +++ b/ibis/backends/exasol/compiler.py @@ -206,6 +206,9 @@ def visit_ExtractDayOfYear(self, op, *, arg): def visit_ExtractWeekOfYear(self, op, *, arg): return self.cast(self.f.to_char(arg, "IW"), op.dtype) + def visit_ExtractIsoYear(self, op, *, arg): + return self.cast(self.f.to_char(arg, "IYYY"), op.dtype) + def visit_DayOfWeekName(self, op, *, arg): return self.f.concat( self.f.substr(self.f.to_char(arg, "DAY"), 0, 1), diff --git a/ibis/backends/oracle/compiler.py b/ibis/backends/oracle/compiler.py index 4129aab09082..e824ec1d06e3 100644 --- a/ibis/backends/oracle/compiler.py +++ b/ibis/backends/oracle/compiler.py @@ -450,3 +450,6 @@ def visit_WindowFunction(self, op, *, how, func, start, end, group_by, order_by) def visit_StringConcat(self, op, *, arg): any_args_null = (a.is_(NULL) for a in arg) return self.if_(sg.or_(*any_args_null), NULL, self.f.concat(*arg)) + + def visit_ExtractIsoYear(self, op, *, arg): + return self.cast(self.f.to_char(arg, "IYYY"), op.dtype) diff --git a/ibis/backends/pandas/kernels.py b/ibis/backends/pandas/kernels.py index d941eb391eaf..460be6ac4617 100644 --- a/ibis/backends/pandas/kernels.py +++ b/ibis/backends/pandas/kernels.py @@ -440,6 +440,7 @@ def wrapper(*args, **kwargs): ops.ExtractSecond: lambda arg: arg.dt.second, ops.ExtractWeekOfYear: lambda arg: arg.dt.isocalendar().week.astype("int32"), ops.ExtractYear: lambda arg: arg.dt.year, + ops.ExtractIsoYear: lambda arg: arg.dt.isocalendar().year, ops.IsNull: lambda arg: arg.isnull(), ops.NotNull: lambda arg: arg.notnull(), ops.Lowercase: lambda arg: arg.str.lower(), diff --git a/ibis/backends/polars/compiler.py b/ibis/backends/polars/compiler.py index 840738b25c87..a793db0c609b 100644 --- a/ibis/backends/polars/compiler.py +++ b/ibis/backends/polars/compiler.py @@ -1002,6 +1002,7 @@ def array_flatten(op, **kw): ops.ExtractDay: "day", ops.ExtractMonth: "month", ops.ExtractYear: "year", + ops.ExtractIsoYear: "iso_year", ops.ExtractQuarter: "quarter", ops.ExtractDayOfYear: "ordinal_day", ops.ExtractWeekOfYear: "week", diff --git a/ibis/backends/postgres/compiler.py b/ibis/backends/postgres/compiler.py index 7131f061afb2..0a6af4b91b15 100644 --- a/ibis/backends/postgres/compiler.py +++ b/ibis/backends/postgres/compiler.py @@ -508,6 +508,9 @@ def visit_ExtractDayOfYear(self, op, *, arg): def visit_ExtractWeekOfYear(self, op, *, arg): return self.f.extract("week", arg) + def visit_ExtractIsoYear(self, op, *, arg): + return self.f.extract("isoyear", arg) + def visit_ExtractEpochSeconds(self, op, *, arg): return self.f.extract("epoch", arg) diff --git a/ibis/backends/snowflake/compiler.py b/ibis/backends/snowflake/compiler.py index 473918aad66d..72fcd11c00f2 100644 --- a/ibis/backends/snowflake/compiler.py +++ b/ibis/backends/snowflake/compiler.py @@ -75,6 +75,7 @@ class SnowflakeCompiler(SQLGlotCompiler): ops.BitwiseRightShift: "bitshiftright", ops.BitwiseXor: "bitxor", ops.EndsWith: "endswith", + ops.ExtractIsoYear: "yearofweekiso", ops.Hash: "hash", ops.Median: "median", ops.Mode: "mode", diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py index be6b559a3ab9..b3fb08dd914d 100644 --- a/ibis/backends/tests/test_temporal.py +++ b/ibis/backends/tests/test_temporal.py @@ -10,6 +10,7 @@ import pandas as pd import pytest import sqlglot as sg +import toolz from pytest import param import ibis @@ -99,6 +100,45 @@ def test_timestamp_extract(backend, alltypes, df, attr): backend.assert_series_equal(result, expected) +@pytest.mark.parametrize( + "transform", [toolz.identity, methodcaller("date")], ids=["timestamp", "date"] +) +@pytest.mark.notimpl( + ["druid"], + raises=(AttributeError, com.OperationNotDefinedError), + reason="AttributeError: 'StringColumn' object has no attribute 'X'", +) +@pytest.mark.notyet( + ["mysql", "sqlite", "mssql", "impala", "datafusion", "pyspark", "flink"], + raises=com.OperationNotDefinedError, + reason="backend doesn't appear to support this operation directly", +) +def test_extract_iso_year(backend, con, alltypes, df, transform): + value = transform(alltypes.timestamp_col) + name = "iso_year" + expr = value.iso_year().name(name) + result = expr.execute() + expected = backend.default_series_rename( + df.timestamp_col.dt.isocalendar().year.astype("int32") + ).rename(name) + backend.assert_series_equal(result, expected) + + +@pytest.mark.notimpl( + ["druid"], + raises=(AttributeError, com.OperationNotDefinedError), + reason="AttributeError: 'StringColumn' object has no attribute 'X'", +) +@pytest.mark.notyet( + ["mysql", "sqlite", "mssql", "impala", "datafusion", "pyspark", "flink"], + raises=com.OperationNotDefinedError, + reason="backend doesn't appear to support this operation directly", +) +def test_iso_year_does_not_match_date_year(con): + expr = ibis.date("2022-01-01").iso_year() + assert con.execute(expr) == 2021 + + @pytest.mark.parametrize( ("func", "expected"), [ diff --git a/ibis/backends/trino/compiler.py b/ibis/backends/trino/compiler.py index e53fb73454b5..7c0c7634227e 100644 --- a/ibis/backends/trino/compiler.py +++ b/ibis/backends/trino/compiler.py @@ -80,6 +80,7 @@ class TrinoCompiler(SQLGlotCompiler): ops.ExtractPath: "url_extract_path", ops.ExtractFragment: "url_extract_fragment", ops.ArrayPosition: "array_position", + ops.ExtractIsoYear: "year_of_week", } def _aggregate(self, funcname: str, *args, where): diff --git a/ibis/expr/operations/temporal.py b/ibis/expr/operations/temporal.py index a500dd017400..c2635661f15c 100644 --- a/ibis/expr/operations/temporal.py +++ b/ibis/expr/operations/temporal.py @@ -108,6 +108,11 @@ class ExtractYear(ExtractDateField): pass +@public +class ExtractIsoYear(ExtractDateField): + pass + + @public class ExtractMonth(ExtractDateField): pass diff --git a/ibis/expr/types/temporal.py b/ibis/expr/types/temporal.py index 120f59b8e548..a63e393d915d 100644 --- a/ibis/expr/types/temporal.py +++ b/ibis/expr/types/temporal.py @@ -33,6 +33,10 @@ def year(self) -> ir.IntegerValue: """Extract the year component.""" return ops.ExtractYear(self).to_expr() + def iso_year(self) -> ir.IntegerValue: + """Extract the ISO year component.""" + return ops.ExtractIsoYear(self).to_expr() + def month(self) -> ir.IntegerValue: """Extract the month component.""" return ops.ExtractMonth(self).to_expr()