test_col.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import operator
import pathlib
from typing import Generator

import numpy as np
import pandas as pd
import pytest

import bigframes
import bigframes.pandas as bpd
from bigframes.testing.utils import assert_frame_equal, convert_pandas_dtypes

pytest.importorskip("polars")
pytest.importorskip("pandas", minversion="3.0.0")


CURRENT_DIR = pathlib.Path(__file__).parent
DATA_DIR = CURRENT_DIR.parent / "data"


@pytest.fixture(scope="module", autouse=True)
def session() -> Generator[bigframes.Session, None, None]:
    import bigframes.core.global_session
    from bigframes.testing import polars_session

    session = polars_session.TestSession()
    with bigframes.core.global_session._GlobalSessionContext(session):
        yield session


@pytest.fixture(scope="module")
def scalars_pandas_df_index() -> pd.DataFrame:
    """pd.DataFrame pointing at test data."""

    df = pd.read_json(
        DATA_DIR / "scalars.jsonl",
        lines=True,
    )
    convert_pandas_dtypes(df, bytes_col=True)

    df = df.set_index("rowindex", drop=False)
    df.index.name = None
    return df.set_index("rowindex").sort_index()


@pytest.fixture(scope="module")
def scalars_df_index(
    session: bigframes.Session, scalars_pandas_df_index
) -> bpd.DataFrame:
    return session.read_pandas(scalars_pandas_df_index)


@pytest.fixture(scope="module")
def scalars_df_2_index(
    session: bigframes.Session, scalars_pandas_df_index
) -> bpd.DataFrame:
    return session.read_pandas(scalars_pandas_df_index)


@pytest.fixture(scope="module")
def scalars_dfs(
    scalars_df_index,
    scalars_pandas_df_index,
):
    return scalars_df_index, scalars_pandas_df_index


@pytest.mark.parametrize(
    ("op",),
    [
        (operator.invert,),
    ],
)
def test_pd_col_unary_operators(scalars_dfs, op):
    scalars_df, scalars_pandas_df = scalars_dfs
    bf_kwargs = {
        "result": op(bpd.col("float64_col")),
    }
    pd_kwargs = {
        "result": op(pd.col("float64_col")),  # type: ignore
    }
    df = scalars_df.assign(**bf_kwargs)

    bf_result = df.to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)

    assert_frame_equal(bf_result, pd_result)


@pytest.mark.parametrize(
    ("op"),
    [
        (lambda x: x.sum()),
        (lambda x: x.mean()),
        (lambda x: x.min()),
        (lambda x: x.max()),
        (lambda x: x.std()),
        (lambda x: x.var()),
    ],
    ids=[
        "sum",
        "mean",
        "min",
        "max",
        "std",
        "var",
    ],
)
def test_pd_col_aggregate_op(scalars_dfs, op):
    scalars_df, scalars_pandas_df = scalars_dfs
    bf_kwargs = {
        "result": op(bpd.col("float64_col")),
    }
    pd_kwargs = {
        "result": op(pd.col("float64_col")),  # type: ignore
    }
    df = scalars_df.assign(**bf_kwargs)

    bf_result = df.to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)

    assert_frame_equal(bf_result, pd_result)


def test_pd_col_aggregate_of_aggregate(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs
    bf_kwargs = {
        "result": (bpd.col("int64_col") - bpd.col("int64_col").mean()).mean(),
    }
    pd_kwargs = {
        "result": (pd.col("int64_col") - pd.col("int64_col").mean()).mean(),  # type: ignore
    }
    df = scalars_df.assign(**bf_kwargs)

    bf_result = df.to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)

    assert_frame_equal(bf_result, pd_result)


@pytest.mark.parametrize(
    ("op",),
    [
        (operator.add,),
        (operator.sub,),
        (operator.mul,),
        (operator.truediv,),
        (operator.floordiv,),
        (operator.gt,),
        (operator.lt,),
        (operator.ge,),
        (operator.le,),
        (operator.eq,),
        (operator.mod,),
    ],
)
def test_pd_col_binary_operators(scalars_dfs, op):
    scalars_df, scalars_pandas_df = scalars_dfs
    bf_kwargs = {
        "result": op(bpd.col("float64_col"), 2.4),
        "reverse_result": op(2.4, bpd.col("float64_col")),
    }
    pd_kwargs = {
        "result": op(pd.col("float64_col"), 2.4),  # type: ignore
        "reverse_result": op(2.4, pd.col("float64_col")),  # type: ignore
    }
    df = scalars_df.assign(**bf_kwargs)

    bf_result = df.to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)

    assert_frame_equal(bf_result, pd_result)


@pytest.mark.parametrize(
    ("op",),
    [
        (operator.and_,),
        (operator.or_,),
        (operator.xor,),
    ],
)
def test_pd_col_binary_bool_operators(scalars_dfs, op):
    scalars_df, scalars_pandas_df = scalars_dfs
    bf_kwargs = {
        "result": op(bpd.col("bool_col"), True),
        "reverse_result": op(False, bpd.col("bool_col")),
    }
    pd_kwargs = {
        "result": op(pd.col("bool_col"), True),  # type: ignore
        "reverse_result": op(False, pd.col("bool_col")),  # type: ignore
    }
    df = scalars_df.assign(**bf_kwargs)

    bf_result = df.to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)

    assert_frame_equal(bf_result, pd_result)


def test_loc_with_pd_col(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs

    bf_result = scalars_df.loc[bpd.col("float64_col") > 4].to_pandas()
    pd_result = scalars_pandas_df.loc[pd.col("float64_col") > 4]  # type: ignore

    assert_frame_equal(bf_result, pd_result)


def test_getitem_with_pd_col(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs

    bf_result = scalars_df[bpd.col("float64_col") > 4].to_pandas()
    pd_result = scalars_pandas_df[pd.col("float64_col") > 4]  # type: ignore

    assert_frame_equal(bf_result, pd_result)


def test_col_str_accessor(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs

    bf_result = scalars_df.assign(result=bpd.col("string_col").str.lower()).to_pandas()
    pd_result = scalars_pandas_df.assign(result=pd.col("string_col").str.lower())  # type: ignore

    assert_frame_equal(bf_result, pd_result)


def test_col_dt_accessor(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs

    bf_result = scalars_df.assign(result=bpd.col("date_col").dt.year).to_pandas()
    pd_result = scalars_pandas_df.assign(result=pd.col("date_col").dt.year)  # type: ignore

    # int64[pyarrow] vs Int64
    assert_frame_equal(bf_result, pd_result, check_dtype=False)


def test_col_numpy_ufunc(scalars_dfs):
    scalars_df, scalars_pandas_df = scalars_dfs

    bf_kwargs = {
        "sqrt": np.sqrt(bpd.col("float64_col")),  # type: ignore
        "add_const": np.add(bpd.col("float64_col"), 2.4),  # type: ignore
        "radd_const": np.add(2.4, bpd.col("float64_col")),  # type: ignore
        "add_cols": np.add(bpd.col("float64_col"), bpd.col("int64_col")),  # type: ignore
    }
    pd_kwargs = {
        "sqrt": np.sqrt(pd.col("float64_col")),  # type: ignore
        "add_const": np.add(pd.col("float64_col"), 2.4),  # type: ignore
        "radd_const": np.add(2.4, pd.col("float64_col")),  # type: ignore
        "add_cols": np.add(pd.col("float64_col"), pd.col("int64_col")),  # type: ignore
    }

    bf_result = scalars_df.assign(**bf_kwargs).to_pandas()
    pd_result = scalars_pandas_df.assign(**pd_kwargs)  # type: ignore

    # int64[pyarrow] vs Int64
    assert_frame_equal(bf_result, pd_result, check_dtype=False)