chore: support timestamp subtractions #1346

sycai · 2025-01-31T00:33:42Z

This PR enables subtraction operations for for Timestamp and datetime types.

We don't support mix-match timestamp and datetime values in the same operations. It's not allowed in Ibis anyway.

TrevorBergeron · 2025-02-04T22:41:46Z

bigframes/core/compile/compiler.py

@@ -58,6 +58,7 @@ def compile_sql(
        # TODO: get rid of output_ids arg
        assert len(output_ids) == len(list(node.fields))
        node = set_output_names(node, output_ids)
+        node = nodes.bottom_up(node, rewrites.op_dynamic_dispatch)


I'm pretty sure this will need to be top-down rather than bottom-up

Sure, though I think it shouldn't matter much, because all node schemas are already stable at this point.

TrevorBergeron · 2025-02-05T00:25:16Z

bigframes/core/compile/compiler.py

+        # Need to dispatch op before compilation to keep it consistent with the compile_sql() call
+        return self._compile_node(nodes.bottom_up(node, rewrites.op_dynamic_dispatch))


lets not run this on every node, instead, lets revive the dead _preprocess helper and apply all the pre-transforms there to the entire tree before running compile_node on the root

SG. Moved the code to _preprocess

TrevorBergeron · 2025-02-05T00:27:18Z

bigframes/core/rewrite/__init__.py

 from bigframes.core.rewrite.order import pull_up_order
 from bigframes.core.rewrite.slices import pullup_limit_from_slice, rewrite_slice

 __all__ = [
    "legacy_join_as_projection",
    "try_row_join",
    "rewrite_slice",
+    "op_dynamic_dispatch",


I think something like "convert_duration_to_int" capture the high level intent best

I named it "rewrite_timedelta_ops" to better indicate that we are replacing the operators, not the values.

TrevorBergeron · 2025-02-05T00:27:42Z

bigframes/core/rewrite/operators.py

+    # TODO(b/394354614): FilterByNode and OrderNode also contain expressions. Need to update them too.
+    return root


as long as we get support those nodes before anybody starts using this!

PR soon to follow!

TrevorBergeron · 2025-02-05T00:29:02Z

bigframes/core/rewrite/operators.py

+    if isinstance(expr, ex.OpExpression):
+        updated_inputs = tuple(
+            map(lambda x: _rewrite_expressions(x, schema), expr.inputs)
+        )
+        return _rewrite_op_expr(expr, updated_inputs)


I believe this will also need to be top-down rather than bottom-up.

I don't think it's possible to do this top-down, because we cannot get the input types by first processing the parent node. The parent node output type can only be decided once we have rewrite all the subtrees.

TrevorBergeron · 2025-02-05T00:29:52Z

bigframes/operations/datetime_ops.py

+        if not dtypes.is_datetime_like(input_types[0]):
+            raise TypeError("expected timestamp input")
+
+        return dtypes.TIMEDETLA_DTYPE


Nice catch. I'm glad we haven't officially announced this feature

TrevorBergeron · 2025-02-05T00:31:44Z

bigframes/series.py

+    def sub(
+        self, other: float | int | pandas.Timestamp | datetime.datetime | Series
+    ) -> Series:
        return self._apply_binary_op(other, ops.sub_op)

-    def rsub(self, other: float | int | Series) -> Series:
+    def rsub(
+        self, other: float | int | pandas.Timestamp | datetime.datetime | Series
+    ) -> Series:
        return self._apply_binary_op(other, ops.sub_op, reverse=True)


We might want to consider giving up on annotating other allowed dtypes

It makes sense. The operators themselves will perform type check for us anyway.

TrevorBergeron · 2025-02-05T00:33:23Z

bigframes/series.py

+def _has_timestamp_type(input: typing.Any) -> bool:
+    if isinstance(input, Series):
+        return bigframes.dtypes.is_datetime_like(input.dtype)
+
+    return isinstance(input, (pandas.Timestamp, datetime.datetime))


* chore: support timestamp subtractions * Fix format * use tree rewrites to dispatch timestamp_diff operator * add TODO for more node updates * polish the code and fix typos * fix comment * add rewrites to compile_raw and compile_peek_sql

* feat: add GeoSeries.from_xy * add from_xy test and update ibis types * update geoseries notebook with from_xy * Update docstring example * fix doctstring lint error * return GeometryDtype() for all ibis geo types * chore: support timestamp subtractions (#1346) * chore: support timestamp subtractions * Fix format * use tree rewrites to dispatch timestamp_diff operator * add TODO for more node updates * polish the code and fix typos * fix comment * add rewrites to compile_raw and compile_peek_sql * chore: add a tool to upload tpcds data to bigquery. (#1367) * chore: add a tool to upload tpcds data to bigquery. * update error type * update docstring --------- Co-authored-by: Shenyang Cai <[email protected]> Co-authored-by: Huan Chen <[email protected]>

chore: support timestamp subtractions

0a58e98

product-auto-label bot added size: m Pull request size is medium. api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. labels Jan 31, 2025

Fix format

2a32701

sycai marked this pull request as ready for review January 31, 2025 00:59

sycai requested review from a team as code owners January 31, 2025 00:59

sycai requested a review from shobsi January 31, 2025 00:59

blunderbuss-gcf bot assigned jiaxunwu Jan 31, 2025

sycai requested review from tswast and TrevorBergeron and removed request for shobsi January 31, 2025 00:59

sycai added 2 commits January 30, 2025 19:27

Merge branch 'main' into sycai_timestamp_diff

c6d5635

Merge branch 'main' into sycai_timestamp_diff

6150324

sycai changed the title ~~chore: support timestamp subtractions~~ chore: support timestamp subtractions for series Feb 4, 2025

sycai and others added 2 commits February 4, 2025 19:46

use tree rewrites to dispatch timestamp_diff operator

1e87e0c

Merge branch 'main' into sycai_timestamp_diff

c621cac

sycai changed the title ~~chore: support timestamp subtractions for series~~ chore: support timestamp subtractions Feb 4, 2025

add TODO for more node updates

9e3038c

TrevorBergeron reviewed Feb 5, 2025

View reviewed changes

sycai and others added 2 commits February 5, 2025 01:32

polish the code and fix typos

6ca29ee

Merge branch 'main' into sycai_timestamp_diff

cf0a5b1

sycai requested a review from TrevorBergeron February 5, 2025 01:34

sycai and others added 3 commits February 5, 2025 02:54

fix comment

066cef3

Merge branch 'main' into sycai_timestamp_diff

1e5d9d2

add rewrites to compile_raw and compile_peek_sql

ab1e1b0

TrevorBergeron approved these changes Feb 5, 2025

View reviewed changes

sycai enabled auto-merge (squash) February 5, 2025 18:59

Merge branch 'main' into sycai_timestamp_diff

2a128e9

sycai merged commit 86b7e72 into main Feb 5, 2025
21 of 23 checks passed

sycai deleted the sycai_timestamp_diff branch February 5, 2025 20:34

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chore: support timestamp subtractions #1346

chore: support timestamp subtractions #1346

sycai commented Jan 31, 2025 •

edited

Loading

TrevorBergeron Feb 4, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025 •

edited

Loading

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

TrevorBergeron Feb 5, 2025

sycai Feb 5, 2025

		# Need to dispatch op before compilation to keep it consistent with the compile_sql() call
		return self._compile_node(nodes.bottom_up(node, rewrites.op_dynamic_dispatch))

		# TODO(b/394354614): FilterByNode and OrderNode also contain expressions. Need to update them too.
		return root

chore: support timestamp subtractions #1346

chore: support timestamp subtractions #1346

Conversation

sycai commented Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sycai Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sycai commented Jan 31, 2025 •

edited

Loading

sycai Feb 5, 2025 •

edited

Loading