# create a RecordBatch and a new DataFrame from it

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])],

names=["a", "b", "c"],

)

ctx = SessionContext()

rows = 100000

data = {

"a": list(range(rows)),

"b": [f"s-{i}" for i in range(rows)],

"c": [float(i + 0.1) for i in range(rows)],

}

batch = pa.record_batch(data)

"""Create a DataFrame with multiple record batches for testing stream behavior.

ctx = SessionContext()

# Create multiple batches, each with 10,000 rows

batches = []

rows_per_batch = 10000

num_batches = 10

for batch_idx in range(num_batches):

start_row = batch_idx * rows_per_batch

end_row = start_row + rows_per_batch

data = {

"a": list(range(start_row, end_row)),

"b": [f"s-{i}" for i in range(start_row, end_row)],

"c": [float(i + 0.1) for i in range(start_row, end_row)],

}

batch = pa.record_batch(data)

batches.append(batch)

# Register as record batches to maintain multi-batch structure

# Using [batches] wraps list in another list as required by register_record_batches

ctx.register_record_batches("large_multi_batch_data", [batches])

ctx = SessionContext()

# create a RecordBatch and a new DataFrame from it

batch = pa.RecordBatch.from_arrays(

[pa.array([{"c": 1}, {"c": 2}, {"c": 3}]), pa.array([4, 5, 6])],

names=["a", "b"],

)

ctx = SessionContext()

# create a RecordBatch and a new DataFrame from it

# Intentionally make each array of different length

batch = pa.RecordBatch.from_arrays(

[pa.array([[1], [2, 3], [4, 5, 6], None]), pa.array([7, 8, 9, 10])],

names=["a", "b"],

)

ctx = SessionContext()

# create a RecordBatch and a new DataFrame from it

batch = pa.RecordBatch.from_arrays(

[

pa.array([0, 1, 2, 3, 4, 5, 6]),

pa.array([7, None, 7, 8, 9, None, 9]),

pa.array(["A", "A", "A", "A", "B", "B", "B"]),

],

names=["a", "b", "c"],

)

"""Create a DataFrame with null values of different types."""

ctx = SessionContext()

# Create a RecordBatch with nulls across different types

batch = pa.RecordBatch.from_arrays(

[

pa.array([1, None, 3, None], type=pa.int64()),

pa.array([4.5, 6.7, None, None], type=pa.float64()),

pa.array(["a", None, "c", None], type=pa.string()),

pa.array([True, None, False, None], type=pa.bool_()),

pa.array(

[10957, None, 18993, None], type=pa.date32()

), # 2000-01-01, null, 2022-01-01, null

pa.array(

[946684800000, None, 1640995200000, None], type=pa.date64()

), # 2000-01-01, null, 2022-01-01, null

],

names=[

"int_col",

"float_col",

"str_col",

"bool_col",

"date32_col",

"date64_col",

],

)

def get_cell_style(self) -> str:

return (

"background-color: #f5f5f5; color: #333; padding: 8px; border: "

"1px solid #ddd;"

)

def get_header_style(self) -> str:

return (

"background-color: #4285f4; color: white; font-weight: bold; "

"padding: 10px; border: 1px solid #3367d6;"

"""Count the number of table rows in HTML content.

df_1 = df.select(

column("a") + column("b"),

column("a") - column("b"),

)

# execute and collect the first (and only) batch

result = df_1.collect()[0]

assert result.column(0) == pa.array([5, 7, 9])

assert result.column(1) == pa.array([-3, -3, -3])

df_2 = df.select("b", "a")

# execute and collect the first (and only) batch

result = df_2.collect()[0]

assert result.column(0) == pa.array([4, 5, 6])

df_1 = df.select_exprs(

"a + b",

"a - b",

)

# execute and collect the first (and only) batch

result = df_1.collect()[0]

assert result.column(0) == pa.array([5, 7, 9])

assert result.column(1) == pa.array([-3, -3, -3])

df_2 = df.select_exprs("b", "a")

# execute and collect the first (and only) batch

result = df_2.collect()[0]

assert result.column(0) == pa.array([4, 5, 6])

assert result.column(1) == pa.array([1, 2, 3])

df_3 = df.select_exprs(

"abs(a + b)",

"abs(a - b)",

)

# execute and collect the first (and only) batch

result = df_3.collect()[0]

assert result.column(0) == pa.array([5, 7, 9])

left = df.alias("l")

right = df.alias("r")

joined = left.join(right, left_on="a", right_on="a").select(

"a",

column("l.b").alias("lb"),

column("r.b").alias("rb"),

)

result = joined.sort(column("a")).collect()[0]

assert result.column(0) == pa.array([1, 2, 3])

assert result.column(1) == pa.array([4, 5, 6])

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["ID_For_Students"])

df = ctx.create_dataframe([[batch]])

# here we must quote to match the original column name

assert df.drop('"ID_For_Students"').schema().names == []

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])], names=["a", "b"]

)

df = ctx.create_dataframe([[batch]])

# with a lower case column, both 'a' and '"a"' work

assert df.drop("a").schema().names == ["b"]

df = ctx.create_dataframe([[batch]])

df = df.select(column("b"), "a")

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.column(0) == pa.array([4, 5, 6])

with pytest.raises(

TypeError,

match=f"Expected Expr or column name.*{re.escape(EXPR_TYPE_ERROR)}",

):

df1 = df.filter(column("a") > literal(2)).select(

column("a") + column("b"),

column("a") - column("b"),

)

# execute and collect the first (and only) batch

result = df1.collect()[0]

assert result.column(0) == pa.array([9])

assert result.column(1) == pa.array([-3])

df.show()

# verify that if there is no filter applied, internal dataframe is unchanged

df2 = df.filter()

assert df.df == df2.df

df3 = df.filter(column("a") > literal(1), column("b") != literal(6))

result = df3.collect()[0]

assert result.column(0) == pa.array([2])

assert result.column(1) == pa.array([5])

df_str = df.filter("a > 2")

result = df_str.collect()[0]

assert result.column(0) == pa.array([3])

assert result.column(1) == pa.array([6])

assert result.column(2) == pa.array([8])

df_mixed = df.filter("a > 1", column("b") != literal(6))

result_mixed = df_mixed.collect()[0]

assert result_mixed.column(0) == pa.array([2])

assert result_mixed.column(1) == pa.array([5])

assert result_mixed.column(2) == pa.array([5])

df_strings = df.filter("a > 1", "b < 6")

result_strings = df_strings.collect()[0]

assert result_strings.column(0) == pa.array([2])

assert result_strings.column(1) == pa.array([5])

plan1 = df.filter(df.parse_sql_expr("a > 2")).logical_plan()

plan2 = df.filter(column("a") > literal(2)).logical_plan()

# object equality not implemented but string representation should match

assert str(plan1) == str(plan2)

df1 = df.filter(df.parse_sql_expr("a > 2")).select(

column("a") + column("b"),

column("a") - column("b"),

)

# execute and collect the first (and only) batch

result = df1.collect()[0]

assert result.column(0) == pa.array([9])

assert result.column(1) == pa.array([-3])

df.show()

# verify that if there is no filter applied, internal dataframe is unchanged

df2 = df.filter()

assert df.df == df2.df

df3 = df.filter(df.parse_sql_expr("a > 1"), df.parse_sql_expr("b != 6"))

result = df3.collect()[0]

assert result.column(0) == pa.array([2])

assert result.column(1) == pa.array([5])

df_empty = df.filter(column("a") > literal(3))

df_empty.show()

captured = capsys.readouterr()

ctx.sql("explain select 1").show()

captured = capsys.readouterr()

assert "1 as Int64(1)" in captured.out

ctx.sql("explain analyze select 1").show()

captured = capsys.readouterr()

df = df.sort(column("b").sort(ascending=False))

table = pa.Table.from_batches(df.collect())

expected = {"a": [3, 2, 1], "b": [6, 5, 4], "c": [8, 5, 8]}

from datafusion import col

result_str = df.sort("a").to_pydict()

result_expr = df.sort(col("a")).to_pydict()

with pytest.raises(

TypeError,

match=f"Expected Expr or column name.*{re.escape(EXPR_TYPE_ERROR)}",

):

from datafusion import col

result_str = df.aggregate("a", [f.count()]).sort("a").to_pydict()

result_expr = df.aggregate(col("a"), [f.count()]).sort("a").to_pydict()

df = df.drop("c")

# execute and collect the first (and only) batch

result = df.collect()[0]

assert df.schema().names == ["a", "b"]

assert result.column(0) == pa.array([1, 2, 3])

df = df.limit(1)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert len(result.column(0)) == 1

# only 3 rows, but limit past the end to ensure that offset is working

df = df.limit(5, offset=2)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert len(result.column(0)) == 1

df = df.head(1)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.column(0) == pa.array([1])

assert result.column(1) == pa.array([4])

df = df.tail(1)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.column(0) == pa.array([3])

assert result.column(1) == pa.array([6])

df = df.with_column("c", "a + b")

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.schema.field(0).name == "a"

assert result.schema.field(1).name == "b"

assert result.schema.field(2).name == "c"

assert result.column(0) == pa.array([1, 2, 3])

assert result.column(1) == pa.array([4, 5, 6])

df = df.with_column("c", column("a") + column("b"))

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.schema.field(0).name == "a"

assert result.schema.field(1).name == "b"

assert result.schema.field(2).name == "c"

assert result.column(0) == pa.array([1, 2, 3])

assert result.column(1) == pa.array([4, 5, 6])

df = df.with_columns(

(column("a") + column("b")).alias("c"),

(column("a") + column("b")).alias("d"),

[

(column("a") + column("b")).alias("e"),

(column("a") + column("b")).alias("f"),

],

g=(column("a") + column("b")),

)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.schema.field(0).name == "a"

assert result.schema.field(1).name == "b"

assert result.schema.field(2).name == "c"

assert result.schema.field(3).name == "d"

assert result.schema.field(4).name == "e"

assert result.schema.field(5).name == "f"

assert result.schema.field(6).name == "g"

assert result.column(0) == pa.array([1, 2, 3])

assert result.column(1) == pa.array([4, 5, 6])

assert result.column(2) == pa.array([5, 7, 9])

assert result.column(3) == pa.array([5, 7, 9])

assert result.column(4) == pa.array([5, 7, 9])

assert result.column(5) == pa.array([5, 7, 9])

df = df.with_columns(

"a + b as c",

"a + b as d",

[

"a + b as e",

"a + b as f",

],

g="a + b",

)

# execute and collect the first (and only) batch

result = df.collect()[0]

assert result.schema.field(0).name == "a"

assert result.schema.field(1).name == "b"

assert result.schema.field(2).name == "c"

assert result.schema.field(3).name == "d"

assert result.schema.field(4).name == "e"

assert result.schema.field(5).name == "f"

assert result.schema.field(6).name == "g"

assert result.column(0) == pa.array([1, 2, 3])

assert result.column(1) == pa.array([4, 5, 6])

assert result.column(2) == pa.array([5, 7, 9])

assert result.column(3) == pa.array([5, 7, 9])

assert result.column(4) == pa.array([5, 7, 9])

assert result.column(5) == pa.array([5, 7, 9])

df = df.cast({"a": pa.float16(), "b": pa.list_(pa.uint32())})

expected = pa.schema(

[("a", pa.float16()), ("b", pa.list_(pa.uint32())), ("c", pa.int64())]

)

batches = []

for batch in df:

batches.append(batch) # noqa: PERF402

# Delete DataFrame to ensure RecordBatches remain valid

del df

assert len(batches) == 1

batch = batches[0]

assert isinstance(batch, RecordBatch)

pa_batch = batch.to_pyarrow()

assert pa_batch.column(0).to_pylist() == [1, 2, 3]

assert pa_batch.column(1).to_pylist() == [4, 5, 6]

stream = df.execute_stream()

batches = list(stream)

assert len(batches) == 1

assert isinstance(batches[0], RecordBatch)

pa_batch = batches[0].to_pyarrow()

assert pa_batch.column(0).to_pylist() == [1, 2, 3]

assert pa_batch.column(1).to_pylist() == [4, 5, 6]

df = df.with_column("c", column("a") + column("b")).with_column_renamed("c", "sum")

result = df.collect()[0]

assert result.schema.field(0).name == "a"

assert result.schema.field(1).name == "b"

nested_df = nested_df.unnest_columns("a")

# execute and collect the first (and only) batch

result = nested_df.collect()[0]

assert result.column(0) == pa.array([1, 2, 3, 4, 5, 6, None])

nested_df = nested_df.unnest_columns("a", preserve_nulls=False)

# execute and collect the first (and only) batch

result = nested_df.collect()[0]

assert result.column(0) == pa.array([1, 2, 3, 4, 5, 6])

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])],

names=["a", "b"],

)

df = ctx.create_dataframe([[batch]], "l")

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2]), pa.array([8, 10])],

names=["a", "c"],

)

df1 = ctx.create_dataframe([[batch]], "r")

df2 = df.join(df1, on="a", how="inner")

df2 = df2.sort(column("a"))

table = pa.Table.from_batches(df2.collect())

expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]}

assert table.to_pydict() == expected

# Test the default behavior for dropping duplicate keys

# Since we may have a duplicate column name and pa.Table()

# hides the fact, instead we need to explicitly check the

# resultant arrays.

df2 = df.join(

df1, left_on="a", right_on="a", how="inner", coalesce_duplicate_keys=True

)

df2 = df2.sort(column("a"))

result = df2.collect()[0]

assert result.num_columns == 3

assert result.column(0) == pa.array([1, 2], pa.int64())

assert result.column(1) == pa.array([4, 5], pa.int64())

assert result.column(2) == pa.array([8, 10], pa.int64())

df2 = df.join(

df1, left_on="a", right_on="a", how="inner", coalesce_duplicate_keys=False

)

df2 = df2.sort(column("l.a"))

result = df2.collect()[0]

assert result.num_columns == 4

assert result.column(0) == pa.array([1, 2], pa.int64())

assert result.column(1) == pa.array([4, 5], pa.int64())

assert result.column(2) == pa.array([1, 2], pa.int64())

assert result.column(3) == pa.array([8, 10], pa.int64())

# Verify we don't make a breaking change to pre-43.0.0

# where users would pass join_keys as a positional argument

df2 = df.join(df1, (["a"], ["a"]), how="inner")

df2 = df2.sort(column("a"))

table = pa.Table.from_batches(df2.collect())

expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]}

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])],

names=["a", "b"],

)

df = ctx.create_dataframe([[batch]], "l")

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2]), pa.array([8, 10])],

names=["a", "c"],

)

df1 = ctx.create_dataframe([[batch]], "r")

with pytest.deprecated_call():

df2 = df.join(df1, join_keys=(["a"], ["a"]), how="inner")

df2.show()

df2 = df2.sort(column("a"))

table = pa.Table.from_batches(df2.collect())

expected = {"a": [1, 2], "c": [8, 10], "b": [4, 5]}

assert table.to_pydict() == expected

with pytest.raises(

ValueError, match=r"`left_on` or `right_on` should not provided with `on`"

):

df2 = df.join(df1, on="a", how="inner", right_on="test")

with pytest.raises(

ValueError, match=r"`left_on` and `right_on` should both be provided."

):

df2 = df.join(df1, left_on="a", how="inner")

with pytest.raises(

ValueError, match=r"either `on` or `left_on` and `right_on` should be provided."

):

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])],

names=["a", "b"],

)

df = ctx.create_dataframe([[batch]], "l")

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2]), pa.array([-8, 10])],

names=["a", "c"],

)

df1 = ctx.create_dataframe([[batch]], "r")

df2 = df.join_on(df1, column("l.a").__eq__(column("r.a")), how="inner")

df2.show()

df2 = df2.sort(column("l.a"))

table = pa.Table.from_batches(df2.collect())

expected = {"a": [1, 2], "c": [-8, 10], "b": [4, 5]}

assert table.to_pydict() == expected

df3 = df.join_on(

df1,

column("l.a").__eq__(column("r.a")),

column("l.a").__lt__(column("r.c")),

how="inner",

)

df3.show()

df3 = df3.sort(column("l.a"))

table = pa.Table.from_batches(df3.collect())

expected = {"a": [2], "c": [10], "b": [5]}

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 3, 5, 7, 9]), pa.array([True, True, True, True, True])],

names=["log_time", "key_frame"],

)

key_frame = ctx.create_dataframe([[batch]])

batch = pa.RecordBatch.from_arrays(

[pa.array([2, 4, 6, 8, 10])],

names=["log_time"],

)

query_times = ctx.create_dataframe([[batch]])

merged = query_times.join(

key_frame,

left_on="log_time",

right_on="log_time",

how="full",

coalesce_duplicate_keys=True,

)

merged = merged.sort(column("log_time"))

result = merged.collect()[0]

assert result.num_columns == 2

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2]), pa.array([4, 5])],

names=["a", "b"],

)

df = ctx.create_dataframe([[batch]], "l")

df1 = ctx.create_dataframe([[batch]], "r")

with pytest.raises(

TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)"

):

with pytest.raises(

TypeError, match=r"Use col\(\)/column\(\) or lit\(\)/literal\(\)"

):

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3, 1, 2, 3]), pa.array([4, 5, 6, 4, 5, 6])],

names=["a", "b"],

)

df_a = ctx.create_dataframe([[batch]]).distinct().sort(column("a"))

batch = pa.RecordBatch.from_arrays(

[pa.array([1, 2, 3]), pa.array([4, 5, 6])],

names=["a", "b"],

)

df_b = ctx.create_dataframe([[batch]]).sort(column("a"))

(

"row",

f.row_number(order_by=[column("b"), column("a").sort(ascending=False)]),

[4, 2, 3, 5, 7, 1, 6],

),

(

"row_w_params",

f.row_number(

order_by=[column("b"), column("a")],

partition_by=[column("c")],

),

[2, 1, 3, 4, 2, 1, 3],

),

(

"row_w_params_no_lists",

f.row_number(

order_by=column("b"),

partition_by=column("c"),

),

[2, 1, 3, 4, 2, 1, 3],

),

("rank", f.rank(order_by=[column("b")]), [3, 1, 3, 5, 6, 1, 6]),

(

"rank_w_params",

f.rank(order_by=[column("b"), column("a")], partition_by=[column("c")]),

[2, 1, 3, 4, 2, 1, 3],

),

(

"rank_w_params_no_lists",

f.rank(order_by=column("a"), partition_by=column("c")),

[1, 2, 3, 4, 1, 2, 3],

),

(

"dense_rank",

f.dense_rank(order_by=[column("b")]),

[2, 1, 2, 3, 4, 1, 4],

),

(

"dense_rank_w_params",

f.dense_rank(order_by=[column("b"), column("a")], partition_by=[column("c")]),

[2, 1, 3, 4, 2, 1, 3],

),

(

"dense_rank_w_params_no_lists",

f.dense_rank(order_by=column("a"), partition_by=column("c")),

[1, 2, 3, 4, 1, 2, 3],

),

(

"percent_rank",

f.round(f.percent_rank(order_by=[column("b")]), literal(3)),

[0.333, 0.0, 0.333, 0.667, 0.833, 0.0, 0.833],

),

(

"percent_rank_w_params",

f.round(

f.percent_rank(

order_by=[column("b"), column("a")], partition_by=[column("c")]

),

literal(3),

),

[0.333, 0.0, 0.667, 1.0, 0.5, 0.0, 1.0],

),

(

"percent_rank_w_params_no_lists",

f.round(

f.percent_rank(order_by=column("a"), partition_by=column("c")),

literal(3),

),

[0.0, 0.333, 0.667, 1.0, 0.0, 0.5, 1.0],

),

(

"cume_dist",

f.round(f.cume_dist(order_by=[column("b")]), literal(3)),

[0.571, 0.286, 0.571, 0.714, 1.0, 0.286, 1.0],