| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -204,7 +204,12 @@ def filter_by_id(self, predicate_id: str, keep_null: bool = False) -> ArrayValue | |||
| 204 | 204 | return self.filter(predicate) | |
| 205 | 205 | ||
| 206 | 206 | def filter(self, predicate: ex.Expression): | |
| 207 | - return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate)) | ||
| 207 | + if predicate.is_scalar_expr: | ||
| 208 | + return ArrayValue(nodes.FilterNode(child=self.node, predicate=predicate)) | ||
| 209 | + else: | ||
| 210 | + arr, filter_ids = self.compute_general_expression([predicate]) | ||
| 211 | + arr = arr.filter_by_id(filter_ids[0]) | ||
| 212 | + return arr.drop_columns(filter_ids) | ||
| 208 | 213 | ||
| 209 | 214 | def order_by( | |
| 210 | 215 | self, by: Sequence[OrderingExpression], is_total_order: bool = False | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -23,6 +23,7 @@ | |||
| 23 | 23 | import pandas as pd | |
| 24 | 24 | ||
| 25 | 25 | import bigframes.core.blocks | |
| 26 | + import bigframes.core.col | ||
| 26 | 27 | import bigframes.core.expression as ex | |
| 27 | 28 | import bigframes.core.guid as guid | |
| 28 | 29 | import bigframes.core.indexes as indexes | |
@@ -36,7 +37,11 @@ | |||
| 36 | 37 | ||
| 37 | 38 | if typing.TYPE_CHECKING: | |
| 38 | 39 | LocSingleKey = Union[ | |
| 39 | - bigframes.series.Series, indexes.Index, slice, bigframes.core.scalar.Scalar | ||
| 40 | + bigframes.series.Series, | ||
| 41 | + indexes.Index, | ||
| 42 | + slice, | ||
| 43 | + bigframes.core.scalar.Scalar, | ||
| 44 | + bigframes.core.col.Expression, | ||
| 40 | 45 | ] | |
| 41 | 46 | ||
| 42 | 47 | ||
@@ -309,6 +314,15 @@ def _loc_getitem_series_or_dataframe( | |||
| 309 | 314 | raise NotImplementedError( | |
| 310 | 315 | f"loc does not yet support indexing with a slice. {constants.FEEDBACK_LINK}" | |
| 311 | 316 | ) | |
| 317 | + if isinstance(key, bigframes.core.col.Expression): | ||
| 318 | + label_to_col_ref = { | ||
| 319 | + label: ex.deref(id) | ||
| 320 | + for id, label in series_or_dataframe._block.col_id_to_label.items() | ||
| 321 | + } | ||
| 322 | + resolved_expr = key._value.bind_variables(label_to_col_ref) | ||
| 323 | + result = series_or_dataframe.copy() | ||
| 324 | + result._set_block(series_or_dataframe._block.filter(resolved_expr)) | ||
| 325 | + return result | ||
| 312 | 326 | if callable(key): | |
| 313 | 327 | raise NotImplementedError( | |
| 314 | 328 | f"loc does not yet support indexing with a callable. {constants.FEEDBACK_LINK}" | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -623,13 +623,18 @@ def __getitem__( | |||
| 623 | 623 | ): # No return type annotations (like pandas) as type cannot always be determined statically | |
| 624 | 624 | # NOTE: This implements the operations described in | |
| 625 | 625 | # https://pandas.pydata.org/docs/getting_started/intro_tutorials/03_subset_data.html | |
| 626 | + import bigframes.core.col | ||
| 627 | + import bigframes.pandas | ||
| 626 | 628 | ||
| 627 | - if isinstance(key, bigframes.series.Series): | ||
| 629 | + if isinstance(key, bigframes.pandas.Series): | ||
| 628 | 630 | return self._getitem_bool_series(key) | |
| 629 | 631 | ||
| 630 | 632 | if isinstance(key, slice): | |
| 631 | 633 | return self.iloc[key] | |
| 632 | 634 | ||
| 635 | + if isinstance(key, bigframes.core.col.Expression): | ||
| 636 | + return self.loc[key] | ||
| 637 | + | ||
| 633 | 638 | # TODO(tswast): Fix this pylance warning: Class overlaps "Hashable" | |
| 634 | 639 | # unsafely and could produce a match at runtime | |
| 635 | 640 | if isinstance(key, blocks.Label): | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -158,3 +158,21 @@ def test_pd_col_binary_bool_operators(scalars_dfs, op): | |||
| 158 | 158 | pd_result = scalars_pandas_df.assign(**pd_kwargs) | |
| 159 | 159 | ||
| 160 | 160 | assert_frame_equal(bf_result, pd_result) | |
| 161 | + | ||
| 162 | + | ||
| 163 | + def test_loc_with_pd_col(scalars_dfs): | ||
| 164 | + scalars_df, scalars_pandas_df = scalars_dfs | ||
| 165 | + | ||
| 166 | + bf_result = scalars_df.loc[bpd.col("float64_col") > 4].to_pandas() | ||
| 167 | + pd_result = scalars_pandas_df.loc[pd.col("float64_col") > 4] # type: ignore | ||
| 168 | + | ||
| 169 | + assert_frame_equal(bf_result, pd_result) | ||
| 170 | + | ||
| 171 | + | ||
| 172 | + def test_getitem_with_pd_col(scalars_dfs): | ||
| 173 | + scalars_df, scalars_pandas_df = scalars_dfs | ||
| 174 | + | ||
| 175 | + bf_result = scalars_df[bpd.col("float64_col") > 4].to_pandas() | ||
| 176 | + pd_result = scalars_pandas_df[pd.col("float64_col") > 4] # type: ignore | ||
| 177 | + | ||
| 178 | + assert_frame_equal(bf_result, pd_result) | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments