feat: support upcasting numeric columns in concat (#294)

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -1506,8 +1506,10 @@ def concat(
         blocks: typing.List[Block] = [self, *other]
         if ignore_index:
             blocks = [block.reset_index() for block in blocks]
-
-        result_labels = _align_indices(blocks)
+            level_names = None
+        else:
+            level_names, level_types = _align_indices(blocks)
+            blocks = [_cast_index(block, level_types) for block in blocks]
 
         index_nlevels = blocks[0].index.nlevels
 
@@ -1522,7 +1524,7 @@ def concat(
             result_expr,
             index_columns=list(result_expr.column_ids)[:index_nlevels],
             column_labels=aligned_blocks[0].column_labels,
-            index_labels=result_labels,
+            index_labels=level_names,
         )
         if ignore_index:
             result_block = result_block.reset_index()
@@ -1783,16 +1785,40 @@ def block_from_local(data) -> Block:
     )
 
 
+def _cast_index(block: Block, dtypes: typing.Sequence[bigframes.dtypes.Dtype]):
+    original_block = block
+    result_ids = []
+    for idx_id, idx_dtype, target_dtype in zip(
+        block.index_columns, block.index_dtypes, dtypes
+    ):
+        if idx_dtype != target_dtype:
+            block, result_id = block.apply_unary_op(idx_id, ops.AsTypeOp(target_dtype))
+            result_ids.append(result_id)
+        else:
+            result_ids.append(idx_id)
+
+    expr = block.expr.select_columns((*result_ids, *original_block.value_columns))
+    return Block(
+        expr,
+        index_columns=result_ids,
+        column_labels=original_block.column_labels,
+        index_labels=original_block.index_labels,
+    )
+
+
 def _align_block_to_schema(
     block: Block, schema: dict[Label, bigframes.dtypes.Dtype]
 ) -> Block:
-    """For a given schema, remap block to schema by reordering columns and inserting nulls."""
+    """For a given schema, remap block to schema by reordering columns,  and inserting nulls."""
     col_ids: typing.Tuple[str, ...] = ()
     for label, dtype in schema.items():
-        # TODO: Support casting to lcd type - requires mixed type support
         matching_ids: typing.Sequence[str] = block.label_to_col_id.get(label, ())
         if len(matching_ids) > 0:
             col_id = matching_ids[-1]
+            col_dtype = block.expr.get_column_type(col_id)
+            if dtype != col_dtype:
+                # If _align_schema worked properly, this should always be an upcast
+                block, col_id = block.apply_unary_op(col_id, ops.AsTypeOp(dtype))
             col_ids = (*col_ids, col_id)
         else:
             block, null_column = block.create_constant(None, dtype=dtype)
@@ -1810,38 +1836,44 @@ def _align_schema(
     return functools.reduce(reduction, schemas)
 
 
-def _align_indices(blocks: typing.Sequence[Block]) -> typing.Sequence[Label]:
-    """Validates that the blocks have compatible indices and returns the resulting label names."""
+def _align_indices(
+    blocks: typing.Sequence[Block],
+) -> typing.Tuple[typing.Sequence[Label], typing.Sequence[bigframes.dtypes.Dtype]]:
+    """Validates that the blocks have compatible indices and returns the resulting label names and dtypes."""
     names = blocks[0].index.names
     types = blocks[0].index.dtypes
+
     for block in blocks[1:]:
         if len(names) != block.index.nlevels:
             raise NotImplementedError(
                 f"Cannot combine indices with different number of levels. Use 'ignore_index'=True. {constants.FEEDBACK_LINK}"
             )
-        if block.index.dtypes != types:
-            raise NotImplementedError(
-                f"Cannot combine different index dtypes. Use 'ignore_index'=True. {constants.FEEDBACK_LINK}"
-            )
         names = [
             lname if lname == rname else None
             for lname, rname in zip(names, block.index.names)
         ]
-    return names
+        types = [
+            bigframes.dtypes.lcd_type_or_throw(ltype, rtype)
+            for ltype, rtype in zip(types, block.index.dtypes)
+        ]
+    types = typing.cast(typing.Sequence[bigframes.dtypes.Dtype], types)
+    return names, types
 
 
 def _combine_schema_inner(
     left: typing.Dict[Label, bigframes.dtypes.Dtype],
     right: typing.Dict[Label, bigframes.dtypes.Dtype],
 ) -> typing.Dict[Label, bigframes.dtypes.Dtype]:
     result = dict()
-    for label, type in left.items():
+    for label, left_type in left.items():
         if label in right:
-            if type != right[label]:
+            right_type = right[label]
+            output_type = bigframes.dtypes.lcd_type(left_type, right_type)
+            if output_type is None:
                 raise ValueError(
                     f"Cannot concat rows with label {label} due to mismatched types. {constants.FEEDBACK_LINK}"
                 )
-            result[label] = type
+            result[label] = output_type
     return result
 
 
@@ -1850,15 +1882,20 @@ def _combine_schema_outer(
     right: typing.Dict[Label, bigframes.dtypes.Dtype],
 ) -> typing.Dict[Label, bigframes.dtypes.Dtype]:
     result = dict()
-    for label, type in left.items():
-        if (label in right) and (type != right[label]):
-            raise ValueError(
-                f"Cannot concat rows with label {label} due to mismatched types. {constants.FEEDBACK_LINK}"
-            )
-        result[label] = type
-    for label, type in right.items():
+    for label, left_type in left.items():
+        if label not in right:
+            result[label] = left_type
+        else:
+            right_type = right[label]
+            output_type = bigframes.dtypes.lcd_type(left_type, right_type)
+            if output_type is None:
+                raise NotImplementedError(
+                    f"Cannot concat rows with label {label} due to mismatched types. {constants.FEEDBACK_LINK}"
+                )
+            result[label] = output_type
+    for label, right_type in right.items():
         if label not in left:
-            result[label] = type
+            result[label] = right_type
     return result
 
 

bigframes/dtypes.py

@@ -40,6 +40,7 @@
     pd.Int64Dtype,
     pd.StringDtype,
     pd.ArrowDtype,
+    gpd.array.GeometryDtype,
 ]
 
 # On BQ side, ARRAY, STRUCT, GEOGRAPHY, JSON are not orderable
@@ -139,7 +140,7 @@
 
 ARROW_TO_IBIS = {arrow: ibis for ibis, arrow in IBIS_TO_ARROW.items()}
 
-IBIS_TO_BIGFRAMES: Dict[ibis_dtypes.DataType, Union[Dtype, np.dtype[Any]]] = {
+IBIS_TO_BIGFRAMES: Dict[ibis_dtypes.DataType, Dtype] = {
     ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS
 }
 # Allow REQUIRED fields to map correctly.
@@ -179,7 +180,7 @@
 
 def ibis_dtype_to_bigframes_dtype(
     ibis_dtype: ibis_dtypes.DataType,
-) -> Union[Dtype, np.dtype[Any]]:
+) -> Dtype:
     """Converts an Ibis dtype to a BigQuery DataFrames dtype
 
     Args:
@@ -340,6 +341,11 @@ def literal_to_ibis_scalar(
         ValueError: if passed literal cannot be coerced to a
         BigQuery DataFrames compatible scalar
     """
+    # Special case: Can create nulls for non-bidirectional types
+    if (force_dtype == gpd.array.GeometryDtype()) and pd.isna(literal):
+        # Ibis has bug for casting nulltype to geospatial, so we perform intermediate cast first
+        geotype = ibis_dtypes.GeoSpatial(geotype="geography", srid=4326, nullable=True)
+        return ibis.literal(None, geotype)
     ibis_dtype = BIGFRAMES_TO_IBIS[force_dtype] if force_dtype else None
 
     if pd.api.types.is_list_like(literal):
@@ -538,6 +544,8 @@ def is_compatible(scalar: typing.Any, dtype: Dtype) -> typing.Optional[Dtype]:
 
 
 def lcd_type(dtype1: Dtype, dtype2: Dtype) -> typing.Optional[Dtype]:
+    if dtype1 == dtype2:
+        return dtype1
     # Implicit conversion currently only supported for numeric types
     hierarchy: list[Dtype] = [
         pd.BooleanDtype(),
@@ -550,3 +558,12 @@ def lcd_type(dtype1: Dtype, dtype2: Dtype) -> typing.Optional[Dtype]:
         return None
     lcd_index = max(hierarchy.index(dtype1), hierarchy.index(dtype2))
     return hierarchy[lcd_index]
+
+
+def lcd_type_or_throw(dtype1: Dtype, dtype2: Dtype) -> Dtype:
+    result = lcd_type(dtype1, dtype2)
+    if result is None:
+        raise NotImplementedError(
+            f"BigFrames cannot upcast {dtype1} and {dtype2} to common type. {constants.FEEDBACK_LINK}"
+        )
+    return result

bigframes/operations/aggregations.py

@@ -396,6 +396,11 @@ def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value:
         )
 
 
+class LastOp(WindowOp):
+    def _as_ibis(self, column: ibis_types.Column, window=None) -> ibis_types.Value:
+        return _apply_window_if_present(column.last(), window)
+
+
 class LastNonNullOp(WindowOp):
     @property
     def skips_nulls(self):

tests/system/small/test_dataframe.py

@@ -3141,9 +3141,9 @@ def test_df___array__(scalars_df_index, scalars_pandas_df_index):
 
 
 def test_getattr_attribute_error_when_pandas_has(scalars_df_index):
-    # asof is implemented in pandas but not in bigframes
+    # swapaxes is implemented in pandas but not in bigframes
     with pytest.raises(AttributeError):
-        scalars_df_index.asof()
+        scalars_df_index.swapaxes()
 
 
 def test_getattr_attribute_error(scalars_df_index):

tests/system/small/test_pandas.py

@@ -185,6 +185,38 @@ def test_concat_dataframe_mismatched_columns(scalars_dfs, how):
     pd.testing.assert_frame_equal(bf_result, pd_result)
 
 
+def test_concat_dataframe_upcasting(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    bf_input1 = scalars_df[["int64_col", "float64_col", "int64_too"]].set_index(
+        "int64_col", drop=True
+    )
+    bf_input1.columns = ["a", "b"]
+    bf_input2 = scalars_df[["int64_too", "int64_col", "float64_col"]].set_index(
+        "float64_col", drop=True
+    )
+    bf_input2.columns = ["a", "b"]
+    bf_result = bpd.concat([bf_input1, bf_input2], join="outer")
+    bf_result = bf_result.to_pandas()
+
+    bf_input1 = (
+        scalars_pandas_df[["int64_col", "float64_col", "int64_too"]]
+        .set_index("int64_col", drop=True)
+        .set_axis(["a", "b"], axis=1)
+    )
+    bf_input2 = (
+        scalars_pandas_df[["int64_too", "int64_col", "float64_col"]]
+        .set_index("float64_col", drop=True)
+        .set_axis(["a", "b"], axis=1)
+    )
+    pd_result = pd.concat(
+        [bf_input1, bf_input2],
+        join="outer",
+    )
+
+    pd.testing.assert_frame_equal(bf_result, pd_result)
+
+
 @pytest.mark.parametrize(
     ("how",),
     [