Skip to content

chore: support timestamp subtractions - #1346

Merged
sycai merged 13 commits into
mainfrom
sycai_timestamp_diff
Feb 5, 2025
Merged

chore: support timestamp subtractions#1346
sycai merged 13 commits into
mainfrom
sycai_timestamp_diff

Conversation

@sycai

@sycaisycai commented Jan 31, 2025

Copy link
Copy Markdown
Contributor

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.

@product-auto-labelproduct-auto-labelBot added size: m Pull request size is medium. api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. labels Jan 31, 2025
@sycai
sycai marked this pull request as ready for review January 31, 2025 00:59
@sycai
sycai requested review from a team and shobsiJanuary 31, 2025 00:59
@sycai
sycai requested review from TrevorBergeron and tswast and removed request for shobsiJanuary 31, 2025 00:59
@sycaisycai changed the title chore: support timestamp subtractionschore: support timestamp subtractions for seriesFeb 4, 2025
@sycaisycai changed the title chore: support timestamp subtractions for serieschore: support timestamp subtractionsFeb 4, 2025
Comment threadbigframes/core/compile/compiler.py Outdated
# 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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment threadbigframes/core/compile/compiler.py Outdated
Comment on lines +113 to +114
# 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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SG. Moved the code to _preprocess

Comment threadbigframes/core/rewrite/__init__.py Outdated
"legacy_join_as_projection",
"try_row_join",
"rewrite_slice",
"op_dynamic_dispatch",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PR soon to follow!

Comment on lines +51 to +55
if isinstance(expr, ex.OpExpression):
updated_inputs = tuple(
map(lambda x: _rewrite_expressions(x, schema), expr.inputs)
)
return _rewrite_op_expr(expr, updated_inputs)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@sycaisycaiFeb 5, 2025

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment threadbigframes/operations/datetime_ops.py Outdated
if not dtypes.is_datetime_like(input_types[0]):
raise TypeError("expected timestamp input")

return dtypes.TIMEDETLA_DTYPE

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

typo

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment threadbigframes/series.py Outdated
Comment on lines 808 to 816
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment threadbigframes/series.py Outdated
Comment on lines +2089 to +2093
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))

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

dead code?

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed

@sycai
sycai enabled auto-merge (squash) February 5, 2025 18:59
@sycai
sycai merged commit 86b7e72 into mainFeb 5, 2025
@sycai
sycai deleted the sycai_timestamp_diff branch February 5, 2025 20:34
arwas11 pushed a commit that referenced this pull request Feb 6, 2025
* 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
tswast pushed a commit that referenced this pull request Feb 6, 2025
* 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 <sycai@users.noreply.github.com>
Co-authored-by: Huan Chen <142538604+Genesis929@users.noreply.github.com>
shuoweil pushed a commit that referenced this pull request Feb 6, 2025
* 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 <sycai@users.noreply.github.com>
Co-authored-by: Huan Chen <142538604+Genesis929@users.noreply.github.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

api: bigqueryIssues related to the googleapis/python-bigquery-dataframes API.size: mPull request size is medium.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants

@sycai@TrevorBergeron@jiaxunwu