# Copyright 2025 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 datetime as dt

import json

import math

import operator

import pathlib

import re

import tempfile

from typing import Generator

import db_dtypes # type: ignore

import geopandas as gpd # type: ignore

import google.api_core.exceptions

import numpy

from packaging.version import Version

import pandas as pd

import pyarrow as pa # type: ignore

import pytest

import shapely.geometry # type: ignore

import bigframes

import bigframes.dtypes as dtypes

import bigframes.features

import bigframes.pandas

import bigframes.pandas as bpd

import bigframes.series as series

from bigframes.testing.utils import (

assert_frame_equal,

assert_series_equal,

convert_pandas_dtypes,

get_first_file_from_wildcard,

pandas_major_version,

)

pytest.importorskip("polars")

pytest.importorskip("pandas", minversion="2.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_default_index(

session: bigframes.Session, scalars_pandas_df_index

) -> bpd.DataFrame:

return session.read_pandas(scalars_pandas_df_index).reset_index(drop=False)

@pytest.fixture(scope="module")

def scalars_df_2_default_index(

session: bigframes.Session, scalars_pandas_df_index

) -> bpd.DataFrame:

return session.read_pandas(scalars_pandas_df_index).reset_index(drop=False)

@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

def test_series_construct_copy(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = series.Series(

scalars_df["int64_col"], name="test_series", dtype="Float64"

).to_pandas()

pd_result = pd.Series(

scalars_pandas_df["int64_col"], name="test_series", dtype="Float64"

)

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_nullable_ints():

bf_result = series.Series(

[1, 3, bigframes.pandas.NA], index=[0, 4, bigframes.pandas.NA]

).to_pandas()

# TODO(b/340885567): fix type error

expected_index = pd.Index( # type: ignore

[0, 4, None],

dtype=pd.Int64Dtype(),

)

expected = pd.Series([1, 3, pd.NA], dtype=pd.Int64Dtype(), index=expected_index)

pd.testing.assert_series_equal(bf_result, expected)

def test_series_construct_timestamps():

datetimes = [

dt.datetime(2020, 1, 20, 20, 20, 20, 20),

dt.datetime(2019, 1, 20, 20, 20, 20, 20),

None,

]

bf_result = series.Series(datetimes).to_pandas()

pd_result = pd.Series(datetimes, dtype=pd.ArrowDtype(pa.timestamp("us")))

assert_series_equal(bf_result, pd_result, check_index_type=False)

def test_series_construct_copy_with_index(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = series.Series(

scalars_df["int64_col"],

name="test_series",

dtype="Float64",

index=scalars_df["int64_too"],

).to_pandas()

pd_result = pd.Series(

scalars_pandas_df["int64_col"],

name="test_series",

dtype="Float64",

index=scalars_pandas_df["int64_too"],

)

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_copy_index(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = series.Series(

scalars_df.index,

name="test_series",

dtype="Float64",

index=scalars_df["int64_too"],

).to_pandas()

pd_result = pd.Series(

scalars_pandas_df.index,

name="test_series",

dtype="Float64",

index=scalars_pandas_df["int64_too"],

)

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_pandas(scalars_dfs):

_, scalars_pandas_df = scalars_dfs

bf_result = series.Series(

scalars_pandas_df["int64_col"], name="test_series", dtype="Float64"

)

pd_result = pd.Series(

scalars_pandas_df["int64_col"], name="test_series", dtype="Float64"

)

assert bf_result.shape == pd_result.shape

pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)

def test_series_construct_from_list():

bf_result = series.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64").to_pandas()

pd_result = pd.Series([1, 1, 2, 3, 5, 8, 13], dtype="Int64")

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_reindex():

bf_result = series.Series(

series.Series({1: 10, 2: 30, 3: 30}), index=[3, 2], dtype="Int64"

).to_pandas()

pd_result = pd.Series(pd.Series({1: 10, 2: 30, 3: 30}), index=[3, 2], dtype="Int64")

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_from_list_w_index():

bf_result = series.Series(

[1, 1, 2, 3, 5, 8, 13], index=[10, 20, 30, 40, 50, 60, 70], dtype="Int64"

).to_pandas()

pd_result = pd.Series(

[1, 1, 2, 3, 5, 8, 13], index=[10, 20, 30, 40, 50, 60, 70], dtype="Int64"

)

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_empty(session: bigframes.Session):

bf_series: series.Series = series.Series(session=session)

pd_series: pd.Series = pd.Series()

bf_result = bf_series.empty

pd_result = pd_series.empty

assert pd_result

assert bf_result == pd_result

def test_series_construct_scalar_no_index():

bf_result = series.Series("hello world", dtype="string[pyarrow]").to_pandas()

pd_result = pd.Series("hello world", dtype="string[pyarrow]")

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_scalar_w_index():

bf_result = series.Series(

"hello world", dtype="string[pyarrow]", index=[0, 2, 1]

).to_pandas()

pd_result = pd.Series("hello world", dtype="string[pyarrow]", index=[0, 2, 1])

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_nan():

bf_result = series.Series(numpy.nan).to_pandas()

pd_result = pd.Series(numpy.nan)

pd_result.index = pd_result.index.astype("Int64")

pd_result = pd_result.astype("Float64")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_scalar_w_bf_index():

bf_result = series.Series(

"hello", index=bigframes.pandas.Index([1, 2, 3])

).to_pandas()

pd_result = pd.Series("hello", index=pd.Index([1, 2, 3], dtype="Int64"))

pd_result = pd_result.astype("string[pyarrow]")

pd.testing.assert_series_equal(bf_result, pd_result)

def test_series_construct_from_list_escaped_strings():

"""Check that special characters are supported."""

strings = [

"string\nwith\nnewline",

"string\twith\ttabs",

"string\\with\\backslashes",

]

bf_result = series.Series(strings, name="test_series", dtype="string[pyarrow]")

pd_result = pd.Series(strings, name="test_series", dtype="string[pyarrow]")

# BigQuery DataFrame default indices use nullable Int64 always

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)

def test_series_construct_geodata():

pd_series = pd.Series(

[

shapely.geometry.Point(1, 1),

shapely.geometry.Point(2, 2),

shapely.geometry.Point(3, 3),

],

dtype=gpd.array.GeometryDtype(),

)

series = bigframes.pandas.Series(pd_series)

assert_series_equal(pd_series, series.to_pandas(), check_index_type=False)

@pytest.mark.parametrize(

("dtype"),

[

pytest.param(pd.Int64Dtype(), id="int"),

pytest.param(pd.Float64Dtype(), id="float"),

pytest.param(pd.StringDtype(storage="pyarrow"), id="string"),

],

)

def test_series_construct_w_dtype(dtype):

data = [1, 2, 3]

expected = pd.Series(data, dtype=dtype)

expected.index = expected.index.astype("Int64")

series = bigframes.pandas.Series(data, dtype=dtype)

pd.testing.assert_series_equal(series.to_pandas(), expected)

def test_series_construct_w_dtype_for_struct():

# The data shows the struct fields are disordered and correctly handled during

# construction.

data = [

{"a": 1, "c": "pandas", "b": dt.datetime(2020, 1, 20, 20, 20, 20, 20)},

{"a": 2, "c": "pandas", "b": dt.datetime(2019, 1, 20, 20, 20, 20, 20)},

{"a": 1, "c": "numpy", "b": None},

]

dtype = pd.ArrowDtype(

pa.struct([("a", pa.int64()), ("c", pa.string()), ("b", pa.timestamp("us"))])

)

series = bigframes.pandas.Series(data, dtype=dtype)

expected = pd.Series(data, dtype=dtype)

expected.index = expected.index.astype("Int64")

pd.testing.assert_series_equal(series.to_pandas(), expected)

def test_series_construct_w_dtype_for_array_string():

data = [["1", "2", "3"], [], ["4", "5"]]

dtype = pd.ArrowDtype(pa.list_(pa.string()))

series = bigframes.pandas.Series(data, dtype=dtype)

expected = pd.Series(data, dtype=dtype)

expected.index = expected.index.astype("Int64")

# Skip dtype check due to internal issue b/321013333. This issue causes array types

# to be converted to the `object` dtype when calling `to_pandas()`, resulting in

# a mismatch with the expected Pandas type.

if bigframes.features.PANDAS_VERSIONS.is_arrow_list_dtype_usable:

check_dtype = True

else:

check_dtype = False

pd.testing.assert_series_equal(

series.to_pandas(), expected, check_dtype=check_dtype

)

def test_series_construct_w_dtype_for_array_struct():

data = [[{"a": 1, "c": "aa"}, {"a": 2, "c": "bb"}], [], [{"a": 3, "c": "cc"}]]

dtype = pd.ArrowDtype(pa.list_(pa.struct([("a", pa.int64()), ("c", pa.string())])))

series = bigframes.pandas.Series(data, dtype=dtype)

expected = pd.Series(data, dtype=dtype)

expected.index = expected.index.astype("Int64")

# Skip dtype check due to internal issue b/321013333. This issue causes array types

# to be converted to the `object` dtype when calling `to_pandas()`, resulting in

# a mismatch with the expected Pandas type.

if bigframes.features.PANDAS_VERSIONS.is_arrow_list_dtype_usable:

check_dtype = True

else:

check_dtype = False

pd.testing.assert_series_equal(

series.to_pandas(), expected, check_dtype=check_dtype

)

def test_series_construct_local_unordered_has_sequential_index(session):

series = bigframes.pandas.Series(

["Sun", "Mon", "Tues", "Wed", "Thurs", "Fri", "Sat"], session=session

)

expected: pd.Index = pd.Index([0, 1, 2, 3, 4, 5, 6], dtype=pd.Int64Dtype())

pd.testing.assert_index_equal(series.index.to_pandas(), expected)

@pytest.mark.parametrize(

("json_type"),

[

pytest.param(dtypes.JSON_DTYPE),

pytest.param("json"),

],

)

def test_series_construct_w_json_dtype(json_type):

data = [

"1",

'"str"',

"false",

'["a", {"b": 1}, null]',

None,

'{"a": {"b": [1, 2, 3], "c": true}}',

]

s = bigframes.pandas.Series(data, dtype=json_type)

assert s.dtype == dtypes.JSON_DTYPE

assert s[0] == "1"

assert s[1] == '"str"'

assert s[2] == "false"

assert s[3] == '["a",{"b":1},null]'

assert pd.isna(s[4])

assert s[5] == '{"a":{"b":[1,2,3],"c":true}}'

def test_series_keys(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df["int64_col"].keys().to_pandas()

pd_result = scalars_pandas_df["int64_col"].keys()

pd.testing.assert_index_equal(bf_result, pd_result)

@pytest.mark.parametrize(

["data", "index"],

[

(["a", "b", "c"], None),

([1, 2, 3], ["a", "b", "c"]),

([1, 2, None], ["a", "b", "c"]),

([1, 2, 3], [pd.NA, "b", "c"]),

([numpy.nan, 2, 3], ["a", "b", "c"]),

],

)

def test_series_items(data, index):

bf_series = series.Series(data, index=index)

pd_series = pd.Series(data, index=index)

for (bf_index, bf_value), (pd_index, pd_value) in zip(

bf_series.items(), pd_series.items()

):

# TODO(jialuo): Remove the if conditions after b/373699458 is addressed.

if not pd.isna(bf_index) or not pd.isna(pd_index):

assert bf_index == pd_index

if not pd.isna(bf_value) or not pd.isna(pd_value):

assert bf_value == pd_value

@pytest.mark.parametrize(

["col_name", "expected_dtype"],

[

("bool_col", pd.BooleanDtype()),

# TODO(swast): Use a more efficient type.

("bytes_col", pd.ArrowDtype(pa.binary())),

("date_col", pd.ArrowDtype(pa.date32())),

("datetime_col", pd.ArrowDtype(pa.timestamp("us"))),

("float64_col", pd.Float64Dtype()),

("geography_col", gpd.array.GeometryDtype()),

("int64_col", pd.Int64Dtype()),

# TODO(swast): Use a more efficient type.

("numeric_col", pd.ArrowDtype(pa.decimal128(38, 9))),

("int64_too", pd.Int64Dtype()),

("string_col", pd.StringDtype(storage="pyarrow")),

("time_col", pd.ArrowDtype(pa.time64("us"))),

("timestamp_col", pd.ArrowDtype(pa.timestamp("us", tz="UTC"))),

],

)

def test_get_column(scalars_dfs, col_name, expected_dtype):

scalars_df, scalars_pandas_df = scalars_dfs

series = scalars_df[col_name]

series_pandas = series.to_pandas()

assert series_pandas.dtype == expected_dtype

assert series_pandas.shape[0] == scalars_pandas_df.shape[0]

def test_series_get_column_default(scalars_dfs):

scalars_df, _ = scalars_dfs

result = scalars_df.get(123123123123123, "default_val")

assert result == "default_val"

@pytest.mark.parametrize(

("key",),

[

("hello",),

(2,),

("int64_col",),

(None,),

],

)

def test_series_contains(scalars_df_index, scalars_pandas_df_index, key):

bf_result = key in scalars_df_index["int64_col"]

pd_result = key in scalars_pandas_df_index["int64_col"]

assert bf_result == pd_result

def test_series_equals_identical(scalars_df_index, scalars_pandas_df_index):

bf_result = scalars_df_index.int64_col.equals(scalars_df_index.int64_col)

pd_result = scalars_pandas_df_index.int64_col.equals(

scalars_pandas_df_index.int64_col

)

assert pd_result == bf_result

def test_series_equals_df(scalars_df_index, scalars_pandas_df_index):

bf_result = scalars_df_index["int64_col"].equals(scalars_df_index[["int64_col"]])

pd_result = scalars_pandas_df_index["int64_col"].equals(

scalars_pandas_df_index[["int64_col"]]

)

assert pd_result == bf_result

def test_series_equals_different_dtype(scalars_df_index, scalars_pandas_df_index):

bf_series = scalars_df_index["int64_col"]

pd_series = scalars_pandas_df_index["int64_col"]

bf_result = bf_series.equals(bf_series.astype("Float64"))

pd_result = pd_series.equals(pd_series.astype("Float64"))

assert pd_result == bf_result

def test_series_equals_different_values(scalars_df_index, scalars_pandas_df_index):

bf_series = scalars_df_index["int64_col"]

pd_series = scalars_pandas_df_index["int64_col"]

bf_result = bf_series.equals(bf_series + 1)

pd_result = pd_series.equals(pd_series + 1)

assert pd_result == bf_result

def test_series_get_with_default_index(scalars_dfs):

col_name = "float64_col"

key = 2

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].get(key)

pd_result = scalars_pandas_df[col_name].get(key)

assert bf_result == pd_result

@pytest.mark.parametrize(

("index_col", "key"),

(

("int64_too", 2),

("string_col", "Hello, World!"),

("int64_too", slice(2, 6)),

),

)

def test_series___getitem__(scalars_dfs, index_col, key):

col_name = "float64_col"

scalars_df, scalars_pandas_df = scalars_dfs

scalars_df = scalars_df.set_index(index_col, drop=False)

scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)

bf_result = scalars_df[col_name][key]

pd_result = scalars_pandas_df[col_name][key]

pd.testing.assert_series_equal(bf_result.to_pandas(), pd_result)

@pytest.mark.parametrize(

("key",),

(

(-2,),

(-1,),

(0,),

(1,),

),

)

def test_series___getitem___with_int_key(scalars_dfs, key):

if pd.__version__.startswith("3."):

pytest.skip("pandas 3.0 dropped getitem with int key")

col_name = "int64_too"

index_col = "string_col"

scalars_df, scalars_pandas_df = scalars_dfs

scalars_df = scalars_df.set_index(index_col, drop=False)

scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)

bf_result = scalars_df[col_name][key]

pd_result = scalars_pandas_df[col_name][key]

assert bf_result == pd_result

def test_series___getitem___with_default_index(scalars_dfs):

col_name = "float64_col"

key = 2

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name][key]

pd_result = scalars_pandas_df[col_name][key]

assert bf_result == pd_result

@pytest.mark.parametrize(

("index_col", "key", "value"),

(

("int64_too", 2, "new_string_value"),

("string_col", "Hello, World!", "updated_value"),

("int64_too", 0, None),

),

)

def test_series___setitem__(scalars_dfs, index_col, key, value):

col_name = "string_col"

scalars_df, scalars_pandas_df = scalars_dfs

scalars_df = scalars_df.set_index(index_col, drop=False)

scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)

bf_series = scalars_df[col_name]

pd_series = scalars_pandas_df[col_name].copy()

bf_series[key] = value

pd_series[key] = value

pd.testing.assert_series_equal(bf_series.to_pandas(), pd_series)

@pytest.mark.parametrize(

("key", "value"),

(

(0, 999),

(1, 888),

(0, None),

(-2345, 777),

),

)

def test_series___setitem___with_int_key_numeric(scalars_dfs, key, value):

col_name = "int64_col"

index_col = "int64_too"

scalars_df, scalars_pandas_df = scalars_dfs

scalars_df = scalars_df.set_index(index_col, drop=False)

scalars_pandas_df = scalars_pandas_df.set_index(index_col, drop=False)

bf_series = scalars_df[col_name]

pd_series = scalars_pandas_df[col_name].copy()

bf_series[key] = value

pd_series[key] = value

pd.testing.assert_series_equal(bf_series.to_pandas(), pd_series)

def test_series___setitem___with_default_index(scalars_dfs):

col_name = "float64_col"

key = 2

value = 123.456

scalars_df, scalars_pandas_df = scalars_dfs

bf_series = scalars_df[col_name]

pd_series = scalars_pandas_df[col_name].copy()

bf_series[key] = value

pd_series[key] = value

assert bf_series.to_pandas().iloc[key] == pd_series.iloc[key]

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_too",),

),

)

def test_abs(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].abs().to_pandas()

pd_result = scalars_pandas_df[col_name].abs()

assert_series_equal(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_too",),

),

)

def test_series_pos(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = (+scalars_df[col_name]).to_pandas()

pd_result = +scalars_pandas_df[col_name]

assert_series_equal(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_too",),

),

)

def test_series_neg(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = (-scalars_df[col_name]).to_pandas()

pd_result = -scalars_pandas_df[col_name]

assert_series_equal(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("bool_col",),

("int64_col",),

),

)

def test_series_invert(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = (~scalars_df[col_name]).to_pandas()

pd_result = ~scalars_pandas_df[col_name]

assert_series_equal(pd_result, bf_result)

def test_fillna(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

col_name = "string_col"

bf_result = scalars_df[col_name].fillna("Missing").to_pandas()

pd_result = scalars_pandas_df[col_name].fillna("Missing")

assert_series_equal(

pd_result,

bf_result,

)

def test_series_replace_scalar_scalar(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

col_name = "string_col"

bf_result = (

scalars_df[col_name].replace("Hello, World!", "Howdy, Planet!").to_pandas()

)

pd_result = scalars_pandas_df[col_name].replace("Hello, World!", "Howdy, Planet!")

pd.testing.assert_series_equal(

pd_result,

bf_result,

)

def test_series_replace_list_scalar(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

col_name = "string_col"

bf_result = (

scalars_df[col_name]

.replace(["Hello, World!", "T"], "Howdy, Planet!")

.to_pandas()

)

pd_result = scalars_pandas_df[col_name].replace(

["Hello, World!", "T"], "Howdy, Planet!"

)

pd.testing.assert_series_equal(

pd_result,

bf_result,

)

@pytest.mark.parametrize(

("replacement_dict",),

(({},),),

ids=[

"empty",

],

)

def test_series_replace_dict(scalars_dfs, replacement_dict):

scalars_df, scalars_pandas_df = scalars_dfs

col_name = "string_col"

bf_result = scalars_df[col_name].replace(replacement_dict).to_pandas()

pd_result = scalars_pandas_df[col_name].replace(replacement_dict)

pd.testing.assert_series_equal(

pd_result,

bf_result,

)

@pytest.mark.parametrize(

("method",),

(

("linear",),

("values",),

("slinear",),

("nearest",),

("zero",),

("pad",),

),

)

def test_series_interpolate(method):

pytest.importorskip("scipy")

if method == "pad" and pd.__version__.startswith("3."):

pytest.skip("pandas 3.0 dropped method='pad'")

values = [None, 1, 2, None, None, 16, None]

index = [-3.2, 11.4, 3.56, 4, 4.32, 5.55, 76.8]

pd_series = pd.Series(values, index)

bf_series = series.Series(pd_series)

# Pandas can only interpolate on "float64" columns

# https://github.com/pandas-dev/pandas/issues/40252

pd_result = pd_series.astype("float64").interpolate(method=method)

bf_result = bf_series.interpolate(method=method).to_pandas()

# pd uses non-null types, while bf uses nullable types

assert_series_equal(

pd_result,

bf_result,

check_index_type=False,

check_dtype=False,

nulls_are_nan=True,

)

@pytest.mark.parametrize(

("ignore_index",),

(

(True,),

(False,),

),

)

def test_series_dropna(scalars_dfs, ignore_index):

if pd.__version__.startswith("1."):

pytest.skip("ignore_index parameter not supported in pandas 1.x.")

scalars_df, scalars_pandas_df = scalars_dfs

col_name = "string_col"

bf_result = scalars_df[col_name].dropna(ignore_index=ignore_index).to_pandas()

pd_result = scalars_pandas_df[col_name].dropna(ignore_index=ignore_index)

assert_series_equal(pd_result, bf_result, check_index_type=False)

@pytest.mark.parametrize(

("agg",),

(

("sum",),

("size",),

),

)

def test_series_agg_single_string(scalars_dfs, agg):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df["int64_col"].agg(agg)

pd_result = scalars_pandas_df["int64_col"].agg(agg)

assert math.isclose(pd_result, bf_result)

def test_series_agg_multi_string(scalars_dfs):

aggregations = [

"sum",

"mean",

"std",

"var",

"min",

"max",

"nunique",

"count",

"size",

]

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df["int64_col"].agg(aggregations).to_pandas()

pd_result = scalars_pandas_df["int64_col"].agg(aggregations)

# Pandas may produce narrower numeric types, but bigframes always produces Float64

pd_result = pd_result.astype("Float64")

pd.testing.assert_series_equal(pd_result, bf_result, check_index_type=False)

@pytest.mark.parametrize(

("col_name",),

(

("string_col",),

("int64_col",),

),

)

def test_max(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].max()

pd_result = scalars_pandas_df[col_name].max()

assert pd_result == bf_result

@pytest.mark.parametrize(

("col_name",),

(

("string_col",),

("int64_col",),

),

)

def test_min(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].min()

pd_result = scalars_pandas_df[col_name].min()

assert pd_result == bf_result

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_col",),

),

)

def test_std(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].std()

pd_result = scalars_pandas_df[col_name].std()

assert math.isclose(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_col",),

),

)

def test_kurt(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].kurt()

pd_result = scalars_pandas_df[col_name].kurt()

assert math.isclose(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_col",),

),

)

def test_skew(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].skew()

pd_result = scalars_pandas_df[col_name].skew()

assert math.isclose(pd_result, bf_result)

def test_skew_undefined(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df["int64_col"].iloc[:2].skew()

pd_result = scalars_pandas_df["int64_col"].iloc[:2].skew()

# both should be pd.NA

assert pd_result is bf_result

def test_kurt_undefined(scalars_dfs):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df["int64_col"].iloc[:3].kurt()

pd_result = scalars_pandas_df["int64_col"].iloc[:3].kurt()

# both should be pd.NA

assert pd_result is bf_result

@pytest.mark.parametrize(

("col_name",),

(

("float64_col",),

("int64_col",),

),

)

def test_var(scalars_dfs, col_name):

scalars_df, scalars_pandas_df = scalars_dfs

bf_result = scalars_df[col_name].var()

pd_result = scalars_pandas_df[col_name].var()

assert math.isclose(pd_result, bf_result)

@pytest.mark.parametrize(

("col_name",),

(

("bool_col",),

("int64_col",),

),

)

def test_mode_stat(scalars_df_index, scalars_pandas_df_index, col_name):

bf_result = scalars_df_index[col_name].mode().to_pandas()

pd_result = scalars_pandas_df_index[col_name].mode()

## Mode implicitly resets index, and bigframes default indices use nullable Int64

pd_result.index = pd_result.index.astype("Int64")

pd.testing.assert_series_equal(

bf_result,

pd_result,

)

@pytest.mark.parametrize(

("operator"),

[