|
34 | 34 | import bigframes.operations.array_ops as arr_ops |
35 | 35 | import bigframes.operations.bool_ops as bool_ops |
36 | 36 | import bigframes.operations.comparison_ops as comp_ops |
| 37 | +import bigframes.operations.date_ops as date_ops |
37 | 38 | import bigframes.operations.datetime_ops as dt_ops |
| 39 | +import bigframes.operations.frequency_ops as freq_ops |
38 | 40 | import bigframes.operations.generic_ops as gen_ops |
39 | 41 | import bigframes.operations.json_ops as json_ops |
40 | 42 | import bigframes.operations.numeric_ops as num_ops |
@@ -75,6 +77,20 @@ def decorator(func): |
75 | 77 |
|
76 | 78 |
|
77 | 79 | if polars_installed: |
| 80 | + _FREQ_MAPPING = { |
| 81 | + "Y": "1y", |
| 82 | + "Q": "1q", |
| 83 | + "M": "1mo", |
| 84 | + "W": "1w", |
| 85 | + "D": "1d", |
| 86 | + "h": "1h", |
| 87 | + "min": "1m", |
| 88 | + "s": "1s", |
| 89 | + "ms": "1ms", |
| 90 | + "us": "1us", |
| 91 | + "ns": "1ns", |
| 92 | + } |
| 93 | + |
78 | 94 | _DTYPE_MAPPING = { |
79 | 95 | # Direct mappings |
80 | 96 | bigframes.dtypes.INT_DTYPE: pl.Int64(), |
@@ -330,11 +346,48 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
330 | 346 | else: |
331 | 347 | return pl.any_horizontal(*(input.str.ends_with(pat) for pat in op.pat)) |
332 | 348 |
|
| 349 | + @compile_op.register(freq_ops.FloorDtOp) |
| 350 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 351 | + assert isinstance(op, freq_ops.FloorDtOp) |
| 352 | + return input.dt.truncate(every=_FREQ_MAPPING[op.freq]) |
| 353 | + |
333 | 354 | @compile_op.register(dt_ops.StrftimeOp) |
334 | 355 | def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
335 | 356 | assert isinstance(op, dt_ops.StrftimeOp) |
336 | 357 | return input.dt.strftime(op.date_format) |
337 | 358 |
|
| 359 | + @compile_op.register(date_ops.YearOp) |
| 360 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 361 | + return input.dt.year() |
| 362 | + |
| 363 | + @compile_op.register(date_ops.QuarterOp) |
| 364 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 365 | + return input.dt.quarter() |
| 366 | + |
| 367 | + @compile_op.register(date_ops.MonthOp) |
| 368 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 369 | + return input.dt.month() |
| 370 | + |
| 371 | + @compile_op.register(date_ops.DayOfWeekOp) |
| 372 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 373 | + return input.dt.weekday() - 1 |
| 374 | + |
| 375 | + @compile_op.register(date_ops.DayOp) |
| 376 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 377 | + return input.dt.day() |
| 378 | + |
| 379 | + @compile_op.register(date_ops.IsoYearOp) |
| 380 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 381 | + return input.dt.iso_year() |
| 382 | + |
| 383 | + @compile_op.register(date_ops.IsoWeekOp) |
| 384 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 385 | + return input.dt.week() |
| 386 | + |
| 387 | + @compile_op.register(date_ops.IsoDayOp) |
| 388 | + def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
| 389 | + return input.dt.weekday() |
| 390 | + |
338 | 391 | @compile_op.register(dt_ops.ParseDatetimeOp) |
339 | 392 | def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr: |
340 | 393 | assert isinstance(op, dt_ops.ParseDatetimeOp) |
|
0 commit comments