#!/usr/bin/env python3
import numpy as np
import pandas as pd
import narwhals as nw

rng = np.random.default_rng(0)
df_sparse = pd.DataFrame(
    {
        "A": pd.arrays.SparseArray(rng.permutation([0, 1, 2] * 10)),
        "B": pd.arrays.SparseArray(rng.permutation([0.0, 0.1, 0.2, -0.1, 0.2] * 6)),
        "C": pd.arrays.SparseArray(rng.permutation([True, False] * 15)),
    }
)
df_dense = df_sparse.sparse.to_dense()

sparse_schema = nw.from_native(df_sparse).schema
dense_schema = nw.from_native(df_dense).schema

print("sparse schema:", sparse_schema)
print("dense schema:", dense_schema)

# These operations still execute and produce the same values as the dense frame.
selected_sparse = nw.from_native(df_sparse).select("A", "B").to_native().sparse.to_dense()
selected_dense = nw.from_native(df_dense).select("A", "B").to_native()
print("select(A, B) matches dense output:", selected_sparse.equals(selected_dense))

filtered_sparse = nw.from_native(df_sparse).filter(nw.col("A") > 0).to_native().sparse.to_dense()
filtered_dense = nw.from_native(df_dense).filter(nw.col("A") > 0).to_native()
print("filter(A > 0) matches dense output:", filtered_sparse.equals(filtered_dense))

with_columns_sparse = nw.from_native(df_sparse).with_columns(double_A=nw.col("A") * 2).to_native().sparse.to_dense()
with_columns_dense = nw.from_native(df_dense).with_columns(double_A=nw.col("A") * 2).to_native()
print("with_columns(double_A) matches dense output:", with_columns_sparse.equals(with_columns_dense))