Repository: apache/superset
File: tests/unit_tests/dataframe_test.py (Python)

Code before:
                {"a": Timestamp("2262-04-11 23:47:16"), "b": 2},
            ],
            id="timestamp conversion success",
        ),
    ],
)
def test_max_pandas_timestamp(input_, expected) -> None:
    from superset.db_engine_specs import BaseEngineSpec
    from superset.result_set import SupersetResultSet

    cursor_descr: DbapiDescription = [
        ("a", "datetime", None, None, None, None, False),
        ("b", "int", None, None, None, None, False),
    ]
    results = SupersetResultSet(input_, cursor_descr, BaseEngineSpec)
    df = results.to_pandas_df()

    assert df_to_records(df) == expected


def test_df_to_records_with_nan_from_division_by_zero() -> None:
    """Test that NaN values from division by zero are converted to None."""
    import numpy as np

    from superset.db_engine_specs import BaseEngineSpec
    from superset.result_set import SupersetResultSet

    # Simulate Athena query: select 0.00 / 0.00 as test
    data = [(np.nan,), (5.0,), (np.nan,)]
    cursor_descr: DbapiDescription = [("test", "double", None, None, None, None, False)]
    results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
    df = results.to_pandas_df()

    assert df_to_records(df) == [
        {"test": None},
        {"test": 5.0},
        {"test": None},
    ]


def test_df_to_records_with_mixed_nan_and_valid_values() -> None:
    """Test that NaN values are properly handled alongside valid numeric data."""
    import numpy as np

    from superset.db_engine_specs import BaseEngineSpec
    from superset.result_set import SupersetResultSet

    # Simulate a query with multiple columns containing NaN values
    data = [
        ("row1", 10.5, np.nan, 100),
        ("row2", np.nan, 20.3, 200),

Review comment (refactor):
Is it possible to move this to the top of the file? that way we could avoid import duplication?

If this is pattern we are following from previous tests, would be good to know if there is any reason why we are doing this, just for context. 

Code after:
    results = SupersetResultSet(input_, cursor_descr, BaseEngineSpec)
    df = results.to_pandas_df()

    assert df_to_records(df) == expected


def test_df_to_records_with_nan_from_division_by_zero() -> None:
    """Test that NaN values from division by zero are converted to None."""
    # Simulate Athena query: select 0.00 / 0.00 as test
    data = [(np.nan,), (5.0,), (np.nan,)]
    cursor_descr: DbapiDescription = [("test", "double", None, None, None, None, False)]
    results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
    df = results.to_pandas_df()

    assert df_to_records(df) == [
        {"test": None},
        {"test": 5.0},
        {"test": None},
    ]


def test_df_to_records_with_mixed_nan_and_valid_values() -> None:
    """Test that NaN values are properly handled alongside valid numeric data."""

    # Simulate a query with multiple columns containing NaN values
    data = [
        ("row1", 10.5, np.nan, 100),
        ("row2", np.nan, 20.3, 200),
        ("row3", 30.7, 40.2, np.nan),
        ("row4", np.nan, np.nan, np.nan),
    ]
    cursor_descr: DbapiDescription = [
        ("name", "varchar", None, None, None, None, False),
        ("value1", "double", None, None, None, None, False),
        ("value2", "double", None, None, None, None, False),
        ("value3", "int", None, None, None, None, False),
    ]
    results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
    df = results.to_pandas_df()

    assert df_to_records(df) == [
        {"name": "row1", "value1": 10.5, "value2": None, "value3": 100},
        {"name": "row2", "value1": None, "value2": 20.3, "value3": 200},
        {"name": "row3", "value1": 30.7, "value2": 40.2, "value3": None},
        {"name": "row4", "value1": None, "value2": None, "value3": None},
    ]


def test_df_to_records_with_inf_and_nan() -> None:
    """Test that both NaN and infinity values are handled correctly."""
    # Test various edge cases: NaN, positive infinity, negative infinity