Merge branch 'main' into ashleyxu-add-kbins-discretizer

Author: ashleyxuu

.github/.OwlBot.lock.yaml

@@ -13,5 +13,5 @@
 # limitations under the License.
 docker:
   image: gcr.io/cloud-devrel-public-resources/owlbot-python:latest
-  digest: sha256:3e3800bb100af5d7f9e810d48212b37812c1856d20ffeafb99ebe66461b61fc7
-# created: 2023-08-02T10:53:29.114535628Z
+  digest: sha256:fac304457974bb530cc5396abd4ab25d26a469cd3bc97cbfb18c8d4324c584eb
+# created: 2023-10-02T21:31:03.517640371Z

.gitignore

@@ -50,6 +50,7 @@ docs.metadata
 
 # Virtual environment
 env/
+venv/
 
 # Test logs
 coverage.xml

.kokoro/requirements.txt

@@ -113,30 +113,30 @@ commonmark==0.9.1 \
     --hash=sha256:452f9dc859be7f06631ddcb328b6919c67984aca654e5fefb3914d54691aed60 \
     --hash=sha256:da2f38c92590f83de410ba1a3cbceafbc74fee9def35f9251ba9a971d6d66fd9
     # via rich
-cryptography==41.0.3 \
-    --hash=sha256:0d09fb5356f975974dbcb595ad2d178305e5050656affb7890a1583f5e02a306 \
-    --hash=sha256:23c2d778cf829f7d0ae180600b17e9fceea3c2ef8b31a99e3c694cbbf3a24b84 \
-    --hash=sha256:3fb248989b6363906827284cd20cca63bb1a757e0a2864d4c1682a985e3dca47 \
-    --hash=sha256:41d7aa7cdfded09b3d73a47f429c298e80796c8e825ddfadc84c8a7f12df212d \
-    --hash=sha256:42cb413e01a5d36da9929baa9d70ca90d90b969269e5a12d39c1e0d475010116 \
-    --hash=sha256:4c2f0d35703d61002a2bbdcf15548ebb701cfdd83cdc12471d2bae80878a4207 \
-    --hash=sha256:4fd871184321100fb400d759ad0cddddf284c4b696568204d281c902fc7b0d81 \
-    --hash=sha256:5259cb659aa43005eb55a0e4ff2c825ca111a0da1814202c64d28a985d33b087 \
-    --hash=sha256:57a51b89f954f216a81c9d057bf1a24e2f36e764a1ca9a501a6964eb4a6800dd \
-    --hash=sha256:652627a055cb52a84f8c448185922241dd5217443ca194d5739b44612c5e6507 \
-    --hash=sha256:67e120e9a577c64fe1f611e53b30b3e69744e5910ff3b6e97e935aeb96005858 \
-    --hash=sha256:6af1c6387c531cd364b72c28daa29232162010d952ceb7e5ca8e2827526aceae \
-    --hash=sha256:6d192741113ef5e30d89dcb5b956ef4e1578f304708701b8b73d38e3e1461f34 \
-    --hash=sha256:7efe8041897fe7a50863e51b77789b657a133c75c3b094e51b5e4b5cec7bf906 \
-    --hash=sha256:84537453d57f55a50a5b6835622ee405816999a7113267739a1b4581f83535bd \
-    --hash=sha256:8f09daa483aedea50d249ef98ed500569841d6498aa9c9f4b0531b9964658922 \
-    --hash=sha256:95dd7f261bb76948b52a5330ba5202b91a26fbac13ad0e9fc8a3ac04752058c7 \
-    --hash=sha256:a74fbcdb2a0d46fe00504f571a2a540532f4c188e6ccf26f1f178480117b33c4 \
-    --hash=sha256:a983e441a00a9d57a4d7c91b3116a37ae602907a7618b882c8013b5762e80574 \
-    --hash=sha256:ab8de0d091acbf778f74286f4989cf3d1528336af1b59f3e5d2ebca8b5fe49e1 \
-    --hash=sha256:aeb57c421b34af8f9fe830e1955bf493a86a7996cc1338fe41b30047d16e962c \
-    --hash=sha256:ce785cf81a7bdade534297ef9e490ddff800d956625020ab2ec2780a556c313e \
-    --hash=sha256:d0d651aa754ef58d75cec6edfbd21259d93810b73f6ec246436a21b7841908de
+cryptography==41.0.4 \
+    --hash=sha256:004b6ccc95943f6a9ad3142cfabcc769d7ee38a3f60fb0dddbfb431f818c3a67 \
+    --hash=sha256:047c4603aeb4bbd8db2756e38f5b8bd7e94318c047cfe4efeb5d715e08b49311 \
+    --hash=sha256:0d9409894f495d465fe6fda92cb70e8323e9648af912d5b9141d616df40a87b8 \
+    --hash=sha256:23a25c09dfd0d9f28da2352503b23e086f8e78096b9fd585d1d14eca01613e13 \
+    --hash=sha256:2ed09183922d66c4ec5fdaa59b4d14e105c084dd0febd27452de8f6f74704143 \
+    --hash=sha256:35c00f637cd0b9d5b6c6bd11b6c3359194a8eba9c46d4e875a3660e3b400005f \
+    --hash=sha256:37480760ae08065437e6573d14be973112c9e6dcaf5f11d00147ee74f37a3829 \
+    --hash=sha256:3b224890962a2d7b57cf5eeb16ccaafba6083f7b811829f00476309bce2fe0fd \
+    --hash=sha256:5a0f09cefded00e648a127048119f77bc2b2ec61e736660b5789e638f43cc397 \
+    --hash=sha256:5b72205a360f3b6176485a333256b9bcd48700fc755fef51c8e7e67c4b63e3ac \
+    --hash=sha256:7e53db173370dea832190870e975a1e09c86a879b613948f09eb49324218c14d \
+    --hash=sha256:7febc3094125fc126a7f6fb1f420d0da639f3f32cb15c8ff0dc3997c4549f51a \
+    --hash=sha256:80907d3faa55dc5434a16579952ac6da800935cd98d14dbd62f6f042c7f5e839 \
+    --hash=sha256:86defa8d248c3fa029da68ce61fe735432b047e32179883bdb1e79ed9bb8195e \
+    --hash=sha256:8ac4f9ead4bbd0bc8ab2d318f97d85147167a488be0e08814a37eb2f439d5cf6 \
+    --hash=sha256:93530900d14c37a46ce3d6c9e6fd35dbe5f5601bf6b3a5c325c7bffc030344d9 \
+    --hash=sha256:9eeb77214afae972a00dee47382d2591abe77bdae166bda672fb1e24702a3860 \
+    --hash=sha256:b5f4dfe950ff0479f1f00eda09c18798d4f49b98f4e2006d644b3301682ebdca \
+    --hash=sha256:c3391bd8e6de35f6f1140e50aaeb3e2b3d6a9012536ca23ab0d9c35ec18c8a91 \
+    --hash=sha256:c880eba5175f4307129784eca96f4e70b88e57aa3f680aeba3bab0e980b0f37d \
+    --hash=sha256:cecfefa17042941f94ab54f769c8ce0fe14beff2694e9ac684176a2535bf9714 \
+    --hash=sha256:e40211b4923ba5a6dc9769eab704bdb3fbb58d56c5b336d30996c24fcf12aadb \
+    --hash=sha256:efc8ad4e6fc4f1752ebfb58aefece8b4e3c4cae940b0994d43649bdfce8d0d4f
     # via
     #   gcp-releasetool
     #   secretstorage
@@ -382,6 +382,7 @@ protobuf==3.20.3 \
     #   gcp-docuploader
     #   gcp-releasetool
     #   google-api-core
+    #   googleapis-common-protos
 pyasn1==0.4.8 \
     --hash=sha256:39c7e2ec30515947ff4e87fb6f456dfc6e84857d34be479c9d4a4ba4bf46aa5d \
     --hash=sha256:aef77c9fb94a3ac588e87841208bdec464471d9871bd5050a287cc9a475cd0ba

bigframes/core/blocks.py

@@ -38,6 +38,7 @@
 import bigframes.core as core
 import bigframes.core.guid as guid
 import bigframes.core.indexes as indexes
+import bigframes.core.joins as joins
 import bigframes.core.ordering as ordering
 import bigframes.core.utils
 import bigframes.core.utils as utils
@@ -1403,6 +1404,78 @@ def concat(
             result_block = result_block.reset_index()
         return result_block
 
+    def merge(
+        self,
+        other: Block,
+        how: typing.Literal[
+            "inner",
+            "left",
+            "outer",
+            "right",
+        ],
+        left_col_ids: typing.Sequence[str],
+        right_col_ids: typing.Sequence[str],
+        sort: bool,
+        suffixes: tuple[str, str] = ("_x", "_y"),
+    ) -> Block:
+        (
+            joined_expr,
+            coalesced_join_cols,
+            (get_column_left, get_column_right),
+        ) = joins.join_by_column(
+            self.expr,
+            left_col_ids,
+            other.expr,
+            right_col_ids,
+            how=how,
+            sort=sort,
+        )
+
+        # which join key parts should be coalesced
+        merge_join_key_mask = [
+            str(self.col_id_to_label[left_id]) == str(other.col_id_to_label[right_id])
+            for left_id, right_id in zip(left_col_ids, right_col_ids)
+        ]
+        labels_to_coalesce = [
+            self.col_id_to_label[col_id]
+            for i, col_id in enumerate(left_col_ids)
+            if merge_join_key_mask[i]
+        ]
+
+        def left_col_mapping(col_id: str) -> str:
+            if col_id in left_col_ids:
+                join_key_part = left_col_ids.index(col_id)
+                if merge_join_key_mask[join_key_part]:
+                    return coalesced_join_cols[join_key_part]
+            return get_column_left(col_id)
+
+        def right_col_mapping(col_id: str) -> typing.Optional[str]:
+            if col_id in right_col_ids:
+                join_key_part = right_col_ids.index(col_id)
+                if merge_join_key_mask[join_key_part]:
+                    return None
+            return get_column_right(col_id)
+
+        left_columns = [left_col_mapping(col_id) for col_id in self.value_columns]
+
+        right_columns = [
+            typing.cast(str, right_col_mapping(col_id))
+            for col_id in other.value_columns
+            if right_col_mapping(col_id)
+        ]
+
+        expr = joined_expr.select_columns([*left_columns, *right_columns])
+        labels = utils.merge_column_labels(
+            self.column_labels,
+            other.column_labels,
+            coalesce_labels=labels_to_coalesce,
+            suffixes=suffixes,
+        )
+
+        # Constructs default index
+        expr, offset_index_id = expr.promote_offsets()
+        return Block(expr, index_columns=[offset_index_id], column_labels=labels)
+
     def _force_reproject(self) -> Block:
         """Forces a reprojection of the underlying tables expression. Used to force predicate/order application before subsequent operations."""
         return Block(

bigframes/core/indexers.py

@@ -332,8 +332,6 @@ def _iloc_getitem_series_or_dataframe(
     elif isinstance(key, slice):
         return series_or_dataframe._slice(key.start, key.stop, key.step)
     elif pd.api.types.is_list_like(key):
-        # TODO(henryjsolberg): support MultiIndex
-
         if len(key) == 0:
             return typing.cast(
                 typing.Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
@@ -346,15 +344,18 @@ def _iloc_getitem_series_or_dataframe(
                 original_series_name if original_series_name is not None else "0"
             )
             df = series_or_dataframe.to_frame()
-        original_index_name = df.index.name
-        temporary_index_name = guid.generate_guid(prefix="temp_iloc_index_")
-        df = df.rename_axis(temporary_index_name)
+        original_index_names = df.index.names
+        temporary_index_names = [
+            guid.generate_guid(prefix="temp_iloc_index_")
+            for _ in range(len(df.index.names))
+        ]
+        df = df.rename_axis(temporary_index_names)
 
         # set to offset index and use regular loc, then restore index
         df = df.reset_index(drop=False)
         result = df.loc[key]
-        result = result.set_index(temporary_index_name)
-        result = result.rename_axis(original_index_name)
+        result = result.set_index(temporary_index_names)
+        result = result.rename_axis(original_index_names)
 
         if isinstance(series_or_dataframe, bigframes.series.Series):
             result = result[series_name]

bigframes/core/joins/single_column.py

@@ -44,7 +44,6 @@ def join_by_column(
         "right",
     ],
     sort: bool = False,
-    coalesce_join_keys: bool = True,
     allow_row_identity_join: bool = True,
 ) -> Tuple[
     core.ArrayValue,
@@ -59,8 +58,6 @@ def join_by_column(
         right: Expression for right table to join.
         right_column_ids: Column IDs (not label) to join by.
         how: The type of join to perform.
-        coalesce_join_keys: if set to False, returned column ids will contain
-            both left and right join key columns.
         allow_row_identity_join (bool):
             If True, allow matching by row identity. Set to False to always
             perform a true JOIN in generated SQL.
@@ -71,8 +68,6 @@ def join_by_column(
         * Sequence[str]: Column IDs of the coalesced join columns. Sometimes either the
           left/right table will have missing rows. This column pulls the
           non-NULL value from either left/right.
-          If coalesce_join_keys is False, will return uncombined left and
-          right key columns.
         * Tuple[Callable, Callable]: For a given column ID from left or right,
           respectively, return the new column id from the combined expression.
     """
@@ -100,9 +95,7 @@ def join_by_column(
         right_join_keys = [
             combined_expr.get_column(get_column_right(col)) for col in right_column_ids
         ]
-        join_key_cols = get_join_cols(
-            left_join_keys, right_join_keys, how, coalesce_join_keys
-        )
+        join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how)
         join_key_ids = [col.get_name() for col in join_key_cols]
         combined_expr = combined_expr.projection(
             [*join_key_cols, *combined_expr.columns]
@@ -182,9 +175,7 @@ def get_column_right(col_id):
         right_join_keys = [
             combined_table[get_column_right(col)] for col in right_column_ids
         ]
-        join_key_cols = get_join_cols(
-            left_join_keys, right_join_keys, how, coalesce_join_keys
-        )
+        join_key_cols = get_coalesced_join_cols(left_join_keys, right_join_keys, how)
         # We could filter out the original join columns, but predicates/ordering
         # might still reference them in implicit joins.
         columns = (
@@ -226,46 +217,35 @@ def get_column_right(col_id):
         )
 
 
-def get_join_cols(
+def get_coalesced_join_cols(
     left_join_cols: typing.Iterable[ibis_types.Value],
     right_join_cols: typing.Iterable[ibis_types.Value],
     how: str,
-    coalesce_join_keys: bool = True,
 ) -> typing.List[ibis_types.Value]:
     join_key_cols: list[ibis_types.Value] = []
     for left_col, right_col in zip(left_join_cols, right_join_cols):
-        if not coalesce_join_keys:
+        if how == "left" or how == "inner":
             join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_")))
+        elif how == "right":
             join_key_cols.append(right_col.name(guid.generate_guid(prefix="index_")))
-        else:
-            if how == "left" or how == "inner":
+        elif how == "outer":
+            # The left index and the right index might contain null values, for
+            # example due to an outer join with different numbers of rows. Coalesce
+            # these to take the index value from either column.
+            # Use a random name in case the left index and the right index have the
+            # same name. In such a case, _x and _y suffixes will already be used.
+            # Don't need to coalesce if they are exactly the same column.
+            if left_col.name("index").equals(right_col.name("index")):
                 join_key_cols.append(left_col.name(guid.generate_guid(prefix="index_")))
-            elif how == "right":
-                join_key_cols.append(
-                    right_col.name(guid.generate_guid(prefix="index_"))
-                )
-            elif how == "outer":
-                # The left index and the right index might contain null values, for
-                # example due to an outer join with different numbers of rows. Coalesce
-                # these to take the index value from either column.
-                # Use a random name in case the left index and the right index have the
-                # same name. In such a case, _x and _y suffixes will already be used.
-                # Don't need to coalesce if they are exactly the same column.
-                if left_col.name("index").equals(right_col.name("index")):
-                    join_key_cols.append(
-                        left_col.name(guid.generate_guid(prefix="index_"))
-                    )
-                else:
-                    join_key_cols.append(
-                        ibis.coalesce(
-                            left_col,
-                            right_col,
-                        ).name(guid.generate_guid(prefix="index_"))
-                    )
             else:
-                raise ValueError(
-                    f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}"
+                join_key_cols.append(
+                    ibis.coalesce(
+                        left_col,
+                        right_col,
+                    ).name(guid.generate_guid(prefix="index_"))
                 )
+        else:
+            raise ValueError(f"Unexpected join type: {how}. {constants.FEEDBACK_LINK}")
     return join_key_cols
 
 

bigframes/core/utils.py

@@ -104,3 +104,36 @@ def get_standardized_ids(
     idx_ids, col_ids = ids[: len(idx_ids)], ids[len(idx_ids) :]
 
     return col_ids, idx_ids
+
+
+def merge_column_labels(
+    left_labels: pd.Index,
+    right_labels: pd.Index,
+    coalesce_labels: typing.Sequence,
+    suffixes: tuple[str, str] = ("_x", "_y"),
+) -> pd.Index:
+    result_labels = []
+
+    for col_label in left_labels:
+        if col_label in right_labels:
+            if col_label in coalesce_labels:
+                # Merging on the same column only returns 1 key column from coalesce both.
+                # Take the left key column.
+                result_labels.append(col_label)
+            else:
+                result_labels.append(str(col_label) + suffixes[0])
+        else:
+            result_labels.append(col_label)
+
+    for col_label in right_labels:
+        if col_label in left_labels:
+            if col_label in coalesce_labels:
+                # Merging on the same column only returns 1 key column from coalesce both.
+                # Pass the right key column.
+                pass
+            else:
+                result_labels.append(str(col_label) + suffixes[1])
+        else:
+            result_labels.append(col_label)
+
+    return pd.Index(result_labels)