Skip to content

Commit ce5de57

Browse files
TrevorBergeronTrevorBergeron
authored
feat: Add str, dt accessors to pd.col Expression objects (#2488)
1 parent 47a0feb commit ce5de57

8 files changed

Lines changed: 174 additions & 109 deletions

File tree

‎bigframes/core/col.py‎

Lines changed: 51 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from __future__ importannotations
1515

1616
importdataclasses
17-
fromtypingimportAny, Hashable
17+
fromtypingimportAny, Hashable, Literal, TYPE_CHECKING
1818

1919
importbigframes_vendored.pandas.core.colaspd_col
2020

@@ -23,6 +23,10 @@
2323
importbigframes.operationsasbf_ops
2424
importbigframes.operations.aggregationsasagg_ops
2525

26+
ifTYPE_CHECKING:
27+
importbigframes.operations.datetimesasdatetimes
28+
importbigframes.operations.stringsasstrings
29+
2630

2731
# Not to be confused with the Expression class in `bigframes.core.expressions`
2832
# Name collision unintended
@@ -32,7 +36,7 @@ class Expression:
3236

3337
_value: bf_expression.Expression
3438

35-
def_apply_unary(self, op: bf_ops.UnaryOp) ->Expression:
39+
def_apply_unary_op(self, op: bf_ops.UnaryOp) ->Expression:
3640
returnExpression(op.as_expr(self._value))
3741

3842
def_apply_unary_agg(self, op: agg_ops.UnaryAggregateOp) ->Expression:
@@ -44,7 +48,14 @@ def _apply_unary_agg(self, op: agg_ops.UnaryAggregateOp) -> Expression:
4448
agg_expressions.WindowExpression(agg_expr, window_spec.unbound())
4549
)
4650

47-
def_apply_binary(self, other: Any, op: bf_ops.BinaryOp, reverse: bool=False):
51+
# alignment is purely for series compatibility, and is ignored here
52+
def_apply_binary_op(
53+
self,
54+
other: Any,
55+
op: bf_ops.BinaryOp,
56+
alignment: Literal["outer", "left"] ="outer",
57+
reverse: bool=False,
58+
):
4859
ifisinstance(other, Expression):
4960
other_value=other._value
5061
else:
@@ -55,79 +66,79 @@ def _apply_binary(self, other: Any, op: bf_ops.BinaryOp, reverse: bool = False):
5566
returnExpression(op.as_expr(self._value, other_value))
5667

5768
def__add__(self, other: Any) ->Expression:
58-
returnself._apply_binary(other, bf_ops.add_op)
69+
returnself._apply_binary_op(other, bf_ops.add_op)
5970

6071
def__radd__(self, other: Any) ->Expression:
61-
returnself._apply_binary(other, bf_ops.add_op, reverse=True)
72+
returnself._apply_binary_op(other, bf_ops.add_op, reverse=True)
6273

6374
def__sub__(self, other: Any) ->Expression:
64-
returnself._apply_binary(other, bf_ops.sub_op)
75+
returnself._apply_binary_op(other, bf_ops.sub_op)
6576

6677
def__rsub__(self, other: Any) ->Expression:
67-
returnself._apply_binary(other, bf_ops.sub_op, reverse=True)
78+
returnself._apply_binary_op(other, bf_ops.sub_op, reverse=True)
6879

6980
def__mul__(self, other: Any) ->Expression:
70-
returnself._apply_binary(other, bf_ops.mul_op)
81+
returnself._apply_binary_op(other, bf_ops.mul_op)
7182

7283
def__rmul__(self, other: Any) ->Expression:
73-
returnself._apply_binary(other, bf_ops.mul_op, reverse=True)
84+
returnself._apply_binary_op(other, bf_ops.mul_op, reverse=True)
7485

7586
def__truediv__(self, other: Any) ->Expression:
76-
returnself._apply_binary(other, bf_ops.div_op)
87+
returnself._apply_binary_op(other, bf_ops.div_op)
7788

7889
def__rtruediv__(self, other: Any) ->Expression:
79-
returnself._apply_binary(other, bf_ops.div_op, reverse=True)
90+
returnself._apply_binary_op(other, bf_ops.div_op, reverse=True)
8091

8192
def__floordiv__(self, other: Any) ->Expression:
82-
returnself._apply_binary(other, bf_ops.floordiv_op)
93+
returnself._apply_binary_op(other, bf_ops.floordiv_op)
8394

8495
def__rfloordiv__(self, other: Any) ->Expression:
85-
returnself._apply_binary(other, bf_ops.floordiv_op, reverse=True)
96+
returnself._apply_binary_op(other, bf_ops.floordiv_op, reverse=True)
8697

8798
def__ge__(self, other: Any) ->Expression:
88-
returnself._apply_binary(other, bf_ops.ge_op)
99+
returnself._apply_binary_op(other, bf_ops.ge_op)
89100

90101
def__gt__(self, other: Any) ->Expression:
91-
returnself._apply_binary(other, bf_ops.gt_op)
102+
returnself._apply_binary_op(other, bf_ops.gt_op)
92103

93104
def__le__(self, other: Any) ->Expression:
94-
returnself._apply_binary(other, bf_ops.le_op)
105+
returnself._apply_binary_op(other, bf_ops.le_op)
95106

96107
def__lt__(self, other: Any) ->Expression:
97-
returnself._apply_binary(other, bf_ops.lt_op)
108+
returnself._apply_binary_op(other, bf_ops.lt_op)
98109

99110
def__eq__(self, other: object) ->Expression: # type: ignore
100-
returnself._apply_binary(other, bf_ops.eq_op)
111+
returnself._apply_binary_op(other, bf_ops.eq_op)
101112

102113
def__ne__(self, other: object) ->Expression: # type: ignore
103-
returnself._apply_binary(other, bf_ops.ne_op)
114+
returnself._apply_binary_op(other, bf_ops.ne_op)
104115

105116
def__mod__(self, other: Any) ->Expression:
106-
returnself._apply_binary(other, bf_ops.mod_op)
117+
returnself._apply_binary_op(other, bf_ops.mod_op)
107118

108119
def__rmod__(self, other: Any) ->Expression:
109-
returnself._apply_binary(other, bf_ops.mod_op, reverse=True)
120+
returnself._apply_binary_op(other, bf_ops.mod_op, reverse=True)
110121

111122
def__and__(self, other: Any) ->Expression:
112-
returnself._apply_binary(other, bf_ops.and_op)
123+
returnself._apply_binary_op(other, bf_ops.and_op)
113124

114125
def__rand__(self, other: Any) ->Expression:
115-
returnself._apply_binary(other, bf_ops.and_op, reverse=True)
126+
returnself._apply_binary_op(other, bf_ops.and_op, reverse=True)
116127

117128
def__or__(self, other: Any) ->Expression:
118-
returnself._apply_binary(other, bf_ops.or_op)
129+
returnself._apply_binary_op(other, bf_ops.or_op)
119130

120131
def__ror__(self, other: Any) ->Expression:
121-
returnself._apply_binary(other, bf_ops.or_op, reverse=True)
132+
returnself._apply_binary_op(other, bf_ops.or_op, reverse=True)
122133

123134
def__xor__(self, other: Any) ->Expression:
124-
returnself._apply_binary(other, bf_ops.xor_op)
135+
returnself._apply_binary_op(other, bf_ops.xor_op)
125136

126137
def__rxor__(self, other: Any) ->Expression:
127-
returnself._apply_binary(other, bf_ops.xor_op, reverse=True)
138+
returnself._apply_binary_op(other, bf_ops.xor_op, reverse=True)
128139

129140
def__invert__(self) ->Expression:
130-
returnself._apply_unary(bf_ops.invert_op)
141+
returnself._apply_unary_op(bf_ops.invert_op)
131142

132143
defsum(self) ->Expression:
133144
returnself._apply_unary_agg(agg_ops.sum_op)
@@ -147,6 +158,18 @@ def min(self) -> Expression:
147158
defmax(self) ->Expression:
148159
returnself._apply_unary_agg(agg_ops.max_op)
149160

161+
@property
162+
defdt(self) ->datetimes.DatetimeSimpleMethods:
163+
importbigframes.operations.datetimesasdatetimes
164+
165+
returndatetimes.DatetimeSimpleMethods(self)
166+
167+
@property
168+
defstr(self) ->strings.StringMethods:
169+
importbigframes.operations.stringsasstrings
170+
171+
returnstrings.StringMethods(self)
172+
150173

151174
defcol(col_name: Hashable) ->Expression:
152175
returnExpression(bf_expression.free_var(col_name))

‎bigframes/core/compile/ibis_compiler/scalar_op_registry.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -982,7 +982,9 @@ def isin_op_impl(x: ibis_types.Value, op: ops.IsInOp):
982982

983983
@scalar_op_compiler.register_unary_op(ops.ToDatetimeOp, pass_op=True)
984984
defto_datetime_op_impl(x: ibis_types.Value, op: ops.ToDatetimeOp):
985-
ifx.type() in (ibis_dtypes.str, ibis_dtypes.Timestamp("UTC")): # type: ignore
985+
ifx.type() ==ibis_dtypes.Timestamp(None): # type: ignore
986+
returnx# already a timestamp, no-op
987+
elifx.type() in (ibis_dtypes.str, ibis_dtypes.Timestamp("UTC")): # type: ignore
986988
returnx.try_cast(ibis_dtypes.Timestamp(None)) # type: ignore
987989
else:
988990
# Numerical inputs.

‎bigframes/operations/datetime_ops.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT
7474
dtypes.STRING_DTYPE,
7575
dtypes.DATE_DTYPE,
7676
dtypes.TIMESTAMP_DTYPE,
77+
dtypes.DATETIME_DTYPE,
7778
):
7879
raiseTypeError("expected string or numeric input")
7980
returnpd.ArrowDtype(pa.timestamp("us", tz=None))
@@ -87,6 +88,8 @@ class ToTimestampOp(base_ops.UnaryOp):
8788

8889
defoutput_type(self, *input_types: dtypes.ExpressionType) ->dtypes.ExpressionType:
8990
# Must be numeric or string
91+
ifinput_types[0] ==dtypes.TIMESTAMP_DTYPE:
92+
raiseTypeError("Already tz-aware.")
9093
ifinput_types[0] notin (
9194
dtypes.FLOAT_DTYPE,
9295
dtypes.INT_DTYPE,

0 commit comments

Comments
 (0)