mirror of
https://github.com/apache/superset.git
synced 2026-09-09 00:34:49 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ff9a52ab15 |
+128
-11
@@ -58,7 +58,7 @@ from flask_appbuilder.security.sqla.models import User
|
||||
from flask_babel import get_locale, lazy_gettext as _
|
||||
from jinja2.exceptions import TemplateError, UndefinedError
|
||||
from markupsafe import escape, Markup
|
||||
from pandas import DateOffset
|
||||
from pandas import DateOffset, Timedelta
|
||||
from sqlalchemy import and_, Column, or_, UniqueConstraint
|
||||
from sqlalchemy.exc import MultipleResultsFound
|
||||
from sqlalchemy.ext.hybrid import hybrid_property
|
||||
@@ -1758,6 +1758,15 @@ class SqlaQuery(NamedTuple):
|
||||
sql_shifted_temporal_labels: set[str]
|
||||
|
||||
|
||||
WEEK_GRAINS = (
|
||||
TimeGrain.WEEK_STARTING_SUNDAY,
|
||||
TimeGrain.WEEK_ENDING_SATURDAY,
|
||||
TimeGrain.WEEK,
|
||||
TimeGrain.WEEK_STARTING_MONDAY,
|
||||
TimeGrain.WEEK_ENDING_SUNDAY,
|
||||
)
|
||||
|
||||
|
||||
class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
"""
|
||||
Allows any flask_appbuilder.Model (Query, Table, etc.)
|
||||
@@ -3141,6 +3150,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
x_axis_label: str | None = None,
|
||||
x_axis_is_temporal: bool = False,
|
||||
x_axis_datetime_format: str | None = None,
|
||||
resolved_week_offset: DateOffset | None = None,
|
||||
) -> tuple[pd.DataFrame, list[str]]:
|
||||
"""Determine appropriate join keys and modify DataFrames if needed."""
|
||||
if time_grain and not is_date_range_offset:
|
||||
@@ -3158,7 +3168,12 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
|
||||
# Add offset join columns for relative time offsets
|
||||
self.add_offset_join_column(
|
||||
df, column_name, time_grain, offset, join_column_producer
|
||||
df,
|
||||
column_name,
|
||||
time_grain,
|
||||
offset,
|
||||
join_column_producer,
|
||||
resolved_week_offset,
|
||||
)
|
||||
self.add_offset_join_column(
|
||||
offset_df, column_name, time_grain, None, join_column_producer
|
||||
@@ -3388,6 +3403,19 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
"DATE_RANGE_TIMESHIFTS_ENABLED"
|
||||
)
|
||||
|
||||
# Resolved once per offset, from the main series, and reused for
|
||||
# both the join column and (if needed) the full-range coalesce
|
||||
# below so the two cannot drift apart. Skipped entirely when a
|
||||
# custom join_column_producer is configured: that path bypasses
|
||||
# all built-in offset parsing (including normalize_time_delta),
|
||||
# so resolving here could raise on an offset the producer itself
|
||||
# never needs to parse.
|
||||
resolved_week_offset = (
|
||||
None
|
||||
if join_column_producer
|
||||
else self._resolve_week_grain_offset(df, time_grain, offset)
|
||||
)
|
||||
|
||||
offset_df, actual_join_keys = self._determine_join_keys(
|
||||
df,
|
||||
offset_df,
|
||||
@@ -3399,6 +3427,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
x_axis_label,
|
||||
x_axis_is_temporal,
|
||||
x_axis_datetime_format,
|
||||
resolved_week_offset,
|
||||
)
|
||||
|
||||
# The full-range option is only meaningful for relative offsets aligned
|
||||
@@ -3415,7 +3444,9 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
df = self._perform_join(df, offset_df, actual_join_keys, how=how)
|
||||
|
||||
if use_outer_join:
|
||||
df = self._coalesce_offset_index(df, offset, join_keys)
|
||||
df = self._coalesce_offset_index(
|
||||
df, offset, join_keys, resolved_week_offset
|
||||
)
|
||||
|
||||
df = self._apply_cleanup_logic(
|
||||
df, offset, time_grain, join_keys, is_date_range_offset
|
||||
@@ -3437,6 +3468,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
df: pd.DataFrame,
|
||||
offset: str,
|
||||
join_keys: list[str],
|
||||
resolved_week_offset: DateOffset | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Rebuild the temporal x-axis after an outer join with an offset DataFrame.
|
||||
@@ -3447,23 +3479,90 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
right-hand column, expressed in the offset's own time range (e.g. "yesterday
|
||||
15:00"). Shifting it forward by the offset places it on the main series'
|
||||
axis (e.g. "today 15:00") so the comparison line spans the full period.
|
||||
|
||||
Under a Week grain, ``resolved_week_offset`` is the same whole-week shift
|
||||
used to build the join column (see ``_resolve_week_grain_offset``); reusing
|
||||
it here instead of the raw calendar offset keeps this reconstructed axis
|
||||
value aligned to the same weekday the join matched on.
|
||||
"""
|
||||
x_axis = join_keys[0]
|
||||
offset_x_axis = f"{x_axis}{R_SUFFIX}"
|
||||
if x_axis not in df.columns or offset_x_axis not in df.columns:
|
||||
return df
|
||||
|
||||
# normalize_time_delta returns a negative delta for "... ago" offsets, so
|
||||
# subtracting it shifts the historical timestamp forward onto the main axis.
|
||||
try:
|
||||
forward_shift = DateOffset(**normalize_time_delta(offset))
|
||||
except (ValueError, TimeDeltaAmbiguousError):
|
||||
return df
|
||||
if resolved_week_offset is not None:
|
||||
forward_shift = resolved_week_offset
|
||||
else:
|
||||
# normalize_time_delta returns a negative delta for "... ago"
|
||||
# offsets, so subtracting it shifts the historical timestamp
|
||||
# forward onto the main axis.
|
||||
try:
|
||||
forward_shift = DateOffset(**normalize_time_delta(offset))
|
||||
except (ValueError, TimeDeltaAmbiguousError):
|
||||
return df
|
||||
|
||||
shifted = df[offset_x_axis] - forward_shift
|
||||
df[x_axis] = df[x_axis].fillna(shifted)
|
||||
return df
|
||||
|
||||
@staticmethod
|
||||
def _resolve_week_grain_offset(
|
||||
df: pd.DataFrame,
|
||||
time_grain: str | None,
|
||||
time_offset: str | None,
|
||||
) -> DateOffset | None:
|
||||
"""
|
||||
Resolve a relative time offset applied under a Week grain to a single
|
||||
whole-week ``DateOffset`` shared by every row of ``df``.
|
||||
|
||||
A calendar month/quarter/year is not a whole number of weeks, so
|
||||
applying the raw calendar shift independently to each row rounds to a
|
||||
different number of weeks depending on how many leap days or
|
||||
month-length differences happen to fall inside that particular row's
|
||||
span. Two main-series rows exactly one grain apart can then round to
|
||||
*different* whole-week counts, colliding onto the same shifted date
|
||||
(or skipping one). Resolving the shift once, from a single reference
|
||||
date, and reusing that constant for every row keeps rows exactly as
|
||||
many whole weeks apart as they started -- matching the offset
|
||||
series' own real week-start dates, which are always aligned to the
|
||||
grain's weekday.
|
||||
|
||||
Returns ``None`` when the offset does not apply (no offset, a date
|
||||
range, or a non-Week grain), in which case callers fall back to the
|
||||
original per-call calendar-offset behavior.
|
||||
"""
|
||||
if (
|
||||
not time_grain
|
||||
or time_grain not in WEEK_GRAINS
|
||||
or not time_offset
|
||||
or ExploreMixin.is_valid_date_range_static(time_offset)
|
||||
or df.empty
|
||||
):
|
||||
return None
|
||||
|
||||
reference_column = df.iloc[:, 0]
|
||||
reference_values = reference_column[
|
||||
reference_column.apply(lambda value: hasattr(value, "strftime"))
|
||||
]
|
||||
if reference_values.empty:
|
||||
return None
|
||||
|
||||
# The reference must be picked by value, not row position: two rows
|
||||
# exactly one grain apart can shift by calendar spans that differ by
|
||||
# up to a whole week (depending on how many leap days fall inside
|
||||
# each row's own span), so whichever row happened to land first
|
||||
# would make the resolved constant depend on DataFrame row order.
|
||||
# The minimum is deterministic for a given set of dates regardless
|
||||
# of ordering.
|
||||
reference = reference_values.min()
|
||||
calendar_offset = DateOffset(**normalize_time_delta(time_offset))
|
||||
calendar_shifted = reference + calendar_offset
|
||||
# Timedelta.days floors toward negative infinity, which would round
|
||||
# e.g. an 83-hour ("< half a week") shift down to a full week instead
|
||||
# of zero; dividing by a one-day Timedelta keeps the exact fraction.
|
||||
exact_days = (calendar_shifted - reference) / Timedelta(days=1)
|
||||
return DateOffset(days=round(exact_days / 7) * 7)
|
||||
|
||||
def add_offset_join_column(
|
||||
self,
|
||||
df: pd.DataFrame,
|
||||
@@ -3471,6 +3570,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
time_grain: str,
|
||||
time_offset: str | None = None,
|
||||
join_column_producer: Any = None,
|
||||
resolved_week_offset: DateOffset | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Adds an offset join column to the provided DataFrame.
|
||||
@@ -3482,12 +3582,25 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
:param time_grain: The time grain used to calculate the new column.
|
||||
:param time_offset: The time offset used to calculate the new column.
|
||||
:param join_column_producer: A function to generate the join column.
|
||||
:param resolved_week_offset: Under a Week grain, the single whole-week
|
||||
``DateOffset`` to apply to every row (see
|
||||
``_resolve_week_grain_offset``). Computed from ``df`` when not
|
||||
supplied, so callers that already resolved it for this same
|
||||
``df`` and ``time_offset`` (e.g. to also reuse it in
|
||||
``_coalesce_offset_index``) can pass it through instead of
|
||||
recomputing it.
|
||||
"""
|
||||
if join_column_producer:
|
||||
df[name] = df.apply(lambda row: join_column_producer(row, 0), axis=1)
|
||||
else:
|
||||
if resolved_week_offset is None:
|
||||
resolved_week_offset = self._resolve_week_grain_offset(
|
||||
df, time_grain, time_offset
|
||||
)
|
||||
df[name] = df.apply(
|
||||
lambda row: self.generate_join_column(row, 0, time_grain, time_offset),
|
||||
lambda row: self.generate_join_column(
|
||||
row, 0, time_grain, time_offset, resolved_week_offset
|
||||
),
|
||||
axis=1,
|
||||
)
|
||||
|
||||
@@ -3497,12 +3610,16 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
|
||||
column_index: int,
|
||||
time_grain: str,
|
||||
time_offset: str | None = None,
|
||||
resolved_week_offset: DateOffset | None = None,
|
||||
) -> str:
|
||||
value = row.iloc[column_index]
|
||||
|
||||
if hasattr(value, "strftime"):
|
||||
if time_offset and not ExploreMixin.is_valid_date_range_static(time_offset):
|
||||
value = value + DateOffset(**normalize_time_delta(time_offset))
|
||||
if resolved_week_offset is not None:
|
||||
value = value + resolved_week_offset
|
||||
else:
|
||||
value = value + DateOffset(**normalize_time_delta(time_offset))
|
||||
|
||||
if time_grain in (
|
||||
TimeGrain.WEEK_STARTING_SUNDAY,
|
||||
|
||||
@@ -14,7 +14,8 @@
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
from pandas import DataFrame, Series, Timestamp
|
||||
from flask import current_app
|
||||
from pandas import DataFrame, DateOffset, Series, Timestamp
|
||||
from pandas.testing import assert_frame_equal
|
||||
from pytest import fixture, mark, raises # noqa: PT013
|
||||
|
||||
@@ -67,6 +68,7 @@ _datasource._coalesce_offset_index = ExploreMixin._coalesce_offset_index.__get__
|
||||
# Static methods don't need binding - assign directly
|
||||
_datasource.generate_join_column = ExploreMixin.generate_join_column
|
||||
_datasource.is_valid_date_range_static = ExploreMixin.is_valid_date_range_static
|
||||
_datasource._resolve_week_grain_offset = ExploreMixin._resolve_week_grain_offset
|
||||
|
||||
# Convenience reference for backward compatibility in tests
|
||||
query_context_processor = _datasource
|
||||
@@ -109,6 +111,265 @@ def test_join_column_producer(make_join_column_producer):
|
||||
assert_frame_equal(df, result)
|
||||
|
||||
|
||||
def test_join_column_producer_week_grain_bypasses_offset_parsing(
|
||||
make_join_column_producer,
|
||||
):
|
||||
"""
|
||||
A configured join_column_producer must bypass ALL built-in offset
|
||||
handling, including the Week-grain whole-week resolution added to fix
|
||||
the weekday-drift bug. "one year ago" is outside normalize_time_delta's
|
||||
grammar and would raise TimeDeltaAmbiguousError if the built-in
|
||||
resolution path ran; with a producer configured it must never be
|
||||
reached, and the producer's own output must be used untouched.
|
||||
"""
|
||||
df = DataFrame({"ds": [Timestamp("2020-01-07")]})
|
||||
column_name = "join_column"
|
||||
query_context_processor.add_offset_join_column(
|
||||
df, column_name, TimeGrain.WEEK, "one year ago", make_join_column_producer
|
||||
)
|
||||
result = DataFrame(
|
||||
{"ds": [Timestamp("2020-01-07")], column_name: ["CUSTOM_FORMAT"]}
|
||||
)
|
||||
assert_frame_equal(df, result)
|
||||
|
||||
|
||||
def test_join_offset_dfs_custom_producer_week_grain_bypasses_offset_parsing(
|
||||
monkeypatch, make_join_column_producer
|
||||
) -> None:
|
||||
"""
|
||||
Regression guard for join_offset_dfs itself: it resolves the Week-grain
|
||||
whole-week shift once per offset, before delegating to
|
||||
_determine_join_keys / add_offset_join_column. That resolution must be
|
||||
skipped whenever a join_column_producer is configured for the grain --
|
||||
otherwise a free-form offset like "one year ago" (outside
|
||||
normalize_time_delta's grammar) raises TimeDeltaAmbiguousError before
|
||||
the producer ever runs, even though the producer never needed the
|
||||
built-in offset parsing at all.
|
||||
"""
|
||||
monkeypatch.setitem(
|
||||
current_app.config,
|
||||
"TIME_GRAIN_JOIN_COLUMN_PRODUCERS",
|
||||
{TimeGrain.WEEK: make_join_column_producer},
|
||||
)
|
||||
|
||||
df = DataFrame({"ds": [Timestamp("2020-01-07")], "D": [1]})
|
||||
offset_df = DataFrame({"ds": [Timestamp("2019-01-07")], "B": [5]})
|
||||
offset_dfs = {"one year ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"]
|
||||
)
|
||||
|
||||
assert result["B"].tolist() == [5]
|
||||
|
||||
|
||||
def test_join_offset_dfs_week_grain_year_offset_aligns_to_whole_weeks() -> None:
|
||||
"""
|
||||
A calendar-year ``DateOffset`` shift lands a Monday week-start on a
|
||||
non-Monday (a year is not a whole number of weeks), so ``%Y-W%W`` reports
|
||||
a different week number than the offset query's own real week-start
|
||||
bucket and the join drops every row. The join key must instead be shifted
|
||||
by the nearest whole number of weeks so both week-start dates line up.
|
||||
"""
|
||||
df = DataFrame({"ds": [Timestamp("2026-06-15")], "D": [1]}) # Monday
|
||||
# The real historical week-grain query result for "1 year ago": also a
|
||||
# Monday, but not exactly 365 days back.
|
||||
offset_df = DataFrame({"ds": [Timestamp("2025-06-16")], "B": [5]})
|
||||
offset_dfs = {"1 year ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"]
|
||||
)
|
||||
|
||||
assert result["B"].tolist() == [5]
|
||||
|
||||
|
||||
def test_join_offset_dfs_week_grain_month_offset_aligns_to_whole_weeks() -> None:
|
||||
"""A calendar-month shift has the same weekday-drift issue as a year."""
|
||||
df = DataFrame({"ds": [Timestamp("2026-06-15")], "D": [1]}) # Monday
|
||||
offset_df = DataFrame({"ds": [Timestamp("2026-05-18")], "B": [5]}) # Monday
|
||||
offset_dfs = {"1 month ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"]
|
||||
)
|
||||
|
||||
assert result["B"].tolist() == [5]
|
||||
|
||||
|
||||
def test_join_offset_dfs_week_grain_week_multiple_offset_still_aligns() -> None:
|
||||
"""
|
||||
A week-multiple offset (already a whole number of weeks) must keep
|
||||
aligning correctly; the whole-week rounding is a no-op for it.
|
||||
"""
|
||||
df = DataFrame({"ds": [Timestamp("2026-06-15")], "D": [1]}) # Monday
|
||||
offset_df = DataFrame({"ds": [Timestamp("2026-06-01")], "B": [5]}) # Monday
|
||||
offset_dfs = {"2 weeks ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"]
|
||||
)
|
||||
|
||||
assert result["B"].tolist() == [5]
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
("time_grain", "main_dates", "offset_dates"),
|
||||
[
|
||||
# %W grains: bucket rows are Mondays (WEEK/WEEK_STARTING_MONDAY use
|
||||
# the week-start date; WEEK_ENDING_SUNDAY's week also starts Monday).
|
||||
(
|
||||
TimeGrain.WEEK,
|
||||
["2026-06-15", "2026-06-22"],
|
||||
["2025-06-16", "2025-06-23"],
|
||||
),
|
||||
(
|
||||
TimeGrain.WEEK_STARTING_MONDAY,
|
||||
["2026-06-15", "2026-06-22"],
|
||||
["2025-06-16", "2025-06-23"],
|
||||
),
|
||||
# WEEK_ENDING_SUNDAY is represented by its week-end date (Sunday).
|
||||
(
|
||||
TimeGrain.WEEK_ENDING_SUNDAY,
|
||||
["2026-06-21", "2026-06-28"],
|
||||
["2025-06-22", "2025-06-29"],
|
||||
),
|
||||
# %U grains: WEEK_STARTING_SUNDAY is represented by its week-start
|
||||
# date (Sunday); WEEK_ENDING_SATURDAY by its week-end date (Saturday).
|
||||
(
|
||||
TimeGrain.WEEK_STARTING_SUNDAY,
|
||||
["2026-06-14", "2026-06-21"],
|
||||
["2025-06-15", "2025-06-22"],
|
||||
),
|
||||
(
|
||||
TimeGrain.WEEK_ENDING_SATURDAY,
|
||||
["2026-06-20", "2026-06-27"],
|
||||
["2025-06-21", "2025-06-28"],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_join_offset_dfs_week_grain_variants_align_to_whole_weeks(
|
||||
time_grain: str, main_dates: list[str], offset_dates: list[str]
|
||||
) -> None:
|
||||
"""
|
||||
Every Week-grain variant (Monday- and Sunday-starting) is subject to the
|
||||
same weekday-drift bug and must be fixed the same way.
|
||||
"""
|
||||
df = DataFrame(
|
||||
{"ds": [Timestamp(d) for d in main_dates], "D": [1, 2]},
|
||||
)
|
||||
offset_df = DataFrame(
|
||||
{"ds": [Timestamp(d) for d in offset_dates], "B": [5, 6]},
|
||||
)
|
||||
offset_dfs = {"1 year ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, time_grain, join_keys=["ds"]
|
||||
)
|
||||
|
||||
assert result["B"].tolist() == [5, 6]
|
||||
|
||||
|
||||
def test_join_offset_dfs_week_grain_multi_year_offset_is_injective() -> None:
|
||||
"""
|
||||
Rounding each row's calendar-shift span independently is not injective:
|
||||
with a "3 years ago" offset, 2024-02-26 has a raw calendar span of -1095
|
||||
days (rounds to -1092), while 2024-03-04 -- exactly one week later -- has
|
||||
a raw span of -1096 days (rounds to -1099); both would land on the same
|
||||
historical date 2021-03-01, colliding, while 2021-02-22 is skipped
|
||||
entirely. The whole-week shift must be resolved once for the series and
|
||||
applied uniformly so that main-series rows exactly one grain apart stay
|
||||
exactly one grain apart after the shift, matching the offset query's own
|
||||
(also one-week-apart) real week-start dates.
|
||||
"""
|
||||
df = DataFrame(
|
||||
{
|
||||
"ds": [Timestamp("2024-02-26"), Timestamp("2024-03-04")], # Mondays
|
||||
"D": [1, 2],
|
||||
}
|
||||
)
|
||||
offset_df = DataFrame(
|
||||
{
|
||||
"ds": [Timestamp("2021-03-01"), Timestamp("2021-03-08")], # Mondays
|
||||
"B": [5, 6],
|
||||
}
|
||||
)
|
||||
offset_dfs = {"3 years ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"]
|
||||
)
|
||||
|
||||
# Neither a collision (both rows joining to the same offset row) nor a
|
||||
# gap (one row failing to join): each main row must match its own
|
||||
# distinct, correctly-shifted offset row.
|
||||
assert result["B"].tolist() == [5, 6]
|
||||
|
||||
|
||||
@mark.parametrize(
|
||||
"dates",
|
||||
[
|
||||
[Timestamp("2022-02-28"), Timestamp("2022-03-07")],
|
||||
[Timestamp("2022-03-07"), Timestamp("2022-02-28")],
|
||||
],
|
||||
ids=["ascending", "descending"],
|
||||
)
|
||||
def test_resolve_week_grain_offset_is_order_independent(
|
||||
dates: list[Timestamp],
|
||||
) -> None:
|
||||
"""
|
||||
The resolved whole-week displacement must depend only on the SET of
|
||||
dates in the main series, not on which row happens to be first.
|
||||
2022-02-28 and 2022-03-07 are exactly one week apart, but their raw
|
||||
"14 years ago" calendar spans differ by a full week (-5114 vs -5113
|
||||
days, rounding independently to -5117 vs -5110) because of how many
|
||||
Feb 29ths fall inside each date's own 14-year window. Picking the
|
||||
reference by row position would make the resolved shift -- and every
|
||||
row's join key -- depend on DataFrame row order; picking it by value
|
||||
(the minimum date) does not.
|
||||
"""
|
||||
df = DataFrame({"ds": dates})
|
||||
|
||||
resolved = query_context_processor._resolve_week_grain_offset(
|
||||
df, TimeGrain.WEEK, "14 years ago"
|
||||
)
|
||||
|
||||
assert resolved == DateOffset(days=-5117)
|
||||
|
||||
|
||||
def test_join_offset_dfs_week_grain_full_range_uses_resolved_whole_week_shift() -> None:
|
||||
"""
|
||||
With ``full_range=True``, offset-only rows are projected back onto the
|
||||
main axis by shifting them forward by the offset. That reconstruction
|
||||
must reuse the same resolved whole-week shift as the join, not the raw
|
||||
calendar offset -- otherwise a Monday-aligned main series gets an
|
||||
offset-only row projected onto a Tuesday instead of the following
|
||||
Monday.
|
||||
"""
|
||||
df = DataFrame({"ds": [Timestamp("2026-06-15")], "V": [1.0]}) # Monday
|
||||
offset_df = DataFrame(
|
||||
{
|
||||
"ds": [Timestamp("2025-06-16"), Timestamp("2025-06-23")], # Mondays
|
||||
"B": [10.0, 20.0],
|
||||
}
|
||||
)
|
||||
offset_dfs = {"1 year ago": offset_df}
|
||||
|
||||
result = query_context_processor.join_offset_dfs(
|
||||
df, offset_dfs, TimeGrain.WEEK, join_keys=["ds"], full_range=True
|
||||
)
|
||||
|
||||
expected = DataFrame(
|
||||
{
|
||||
"ds": [Timestamp("2026-06-15"), Timestamp("2026-06-22")], # both Mondays
|
||||
"V": [1.0, None],
|
||||
"B": [10.0, 20.0],
|
||||
}
|
||||
)
|
||||
|
||||
assert_frame_equal(expected, result)
|
||||
|
||||
|
||||
def test_join_offset_dfs_no_offsets():
|
||||
df = DataFrame({"A": ["2021-01-01", "2021-02-01", "2021-03-01"]})
|
||||
offset_dfs = {}
|
||||
|
||||
Reference in New Issue
Block a user