ctx = SessionContext()

# 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()

# 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 = 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])

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 = 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]}

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", 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.cast({"a": pa.float16(), "b": pa.list_(pa.uint32())})

expected = pa.schema(

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

)

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.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

df2 = df.join(df1, left_on="a", right_on="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

# 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.show()

df2 = df2.sort(column("l.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("l.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, 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],

),

("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],

),

(

"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],

),

(

"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],

),

(

"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],

),

(

"cume_dist_w_params",

f.round(

f.cume_dist(

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

),

literal(3),

),

[0.5, 0.25, 0.75, 1.0, 0.667, 0.333, 1.0],

),

(

"ntile",

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

[1, 1, 1, 2, 2, 1, 2],

),

(

"ntile_w_params",

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

[1, 1, 2, 2, 1, 1, 2],

),

("lead", f.lead(column("b"), order_by=[column("b")]), [7, None, 8, 9, 9, 7, None]),

(

"lead_w_params",

f.lead(

column("b"),

shift_offset=2,

default_value=-1,

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

partition_by=[column("c")],

),

[8, 7, -1, -1, -1, 9, -1],

),

("lag", f.lag(column("b"), order_by=[column("b")]), [None, None, 7, 7, 8, None, 9]),

(

"lag_w_params",

f.lag(

column("b"),

shift_offset=2,

default_value=-1,

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

partition_by=[column("c")],

),

[-1, -1, None, 7, -1, -1, None],

),

(

"first_value",

f.first_value(column("a")).over(

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

),

[1, 1, 1, 1, 5, 5, 5],

),

(

"last_value",

f.last_value(column("a")).over(

Window(

partition_by=[column("c")],

order_by=[column("b")],

window_frame=WindowFrame("rows", None, None),

)

),

[3, 3, 3, 3, 6, 6, 6],

),

(

"3rd_value",

f.nth_value(column("b"), 3).over(Window(order_by=[column("a")])),

[None, None, 7, 7, 7, 7, 7],

),

(

"avg",

f.round(f.avg(column("b")).over(Window(order_by=[column("a")])), literal(3)),

[7.0, 7.0, 7.0, 7.333, 7.75, 7.75, 8.0],

),

df = partitioned_df.select(

column("a"), column("b"), column("c"), f.alias(expr, name)

)

df.sort(column("a")).show()

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

expected = {

"a": [0, 1, 2, 3, 4, 5, 6],

"b": [7, None, 7, 8, 9, None, 9],

"c": ["A", "A", "A", "A", "B", "B", "B"],

name: result,

}

# ref: https://github.com/apache/datafusion-python/issues/688

window_frame = WindowFrame("rows", None, None)

col_a = column("a")

# Using `f.window` with or without an unbounded window_frame produces the same

# results. These tests are included as a regression check but can be removed when

# f.window() is deprecated in favor of using the .over() approach.

no_frame = f.window("avg", [col_a]).alias("no_frame")

with_frame = f.window("avg", [col_a], window_frame=window_frame).alias("with_frame")

df_1 = partitioned_df.select(col_a, no_frame, with_frame)

expected = {

"a": [0, 1, 2, 3, 4, 5, 6],

"no_frame": [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0],

"with_frame": [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0],

}

assert df_1.sort(col_a).to_pydict() == expected

# When order is not set, the default frame should be unounded preceeding to

# unbounded following. When order is set, the default frame is unbounded preceeding

# to current row.

no_order = f.avg(col_a).over(Window()).alias("over_no_order")

with_order = f.avg(col_a).over(Window(order_by=[col_a])).alias("over_with_order")

df_2 = partitioned_df.select(col_a, no_order, with_order)

expected = {

"a": [0, 1, 2, 3, 4, 5, 6],

"over_no_order": [3.0, 3.0, 3.0, 3.0, 3.0, 3.0, 3.0],

"over_with_order": [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0],

}

"""Test configuring the HTML formatter with different options."""

# Configure with custom settings

configure_formatter(

max_width=500,

max_height=200,

enable_cell_expansion=False,

)

html_output = df._repr_html_()

# Verify our configuration was applied

assert "max-height: 200px" in html_output

assert "max-width: 500px" in html_output

# With cell expansion disabled, we shouldn't see expandable-container elements

"""Test using custom style providers with the HTML formatter."""

# Configure with custom style provider

configure_formatter(style_provider=CustomStyleProvider())

html_output = df._repr_html_()

# Verify our custom styles were applied

assert "background-color: #4285f4" in html_output

assert "color: white" in html_output

"""Test registering custom type formatters for specific data types."""

# Get current formatter and register custom formatters

formatter = get_formatter()

# Format integers with color based on value

# Using int as the type for the formatter will work since we convert

# Arrow scalar values to Python native types in _get_cell_value

def format_int(value):

return f'<span style="color: {"red" if value > 2 else "blue"}">{value}</span>'

formatter.register_formatter(int, format_int)

html_output = df._repr_html_()

# Our test dataframe has values 1,2,3 so we should see:

"""Test using a custom cell builder function."""

# Create a custom cell builder with distinct styling for different value ranges

def custom_cell_builder(value, row, col, table_id):

try:

num_value = int(value)

if num_value > 5: # Values > 5 get green background with indicator

return (

'<td style="background-color: #d9f0d3" '

f'data-test="high">{value}-high</td>'

)

if num_value < 3: # Values < 3 get blue background with indicator

return (

'<td style="background-color: #d3e9f0" '

f'data-test="low">{value}-low</td>'

)

except (ValueError, TypeError):

pass

# Default styling for other cells (3, 4, 5)

return f'<td style="border: 1px solid #ddd" data-test="mid">{value}-mid</td>'

# Set our custom cell builder

formatter = get_formatter()

formatter.set_custom_cell_builder(custom_cell_builder)

html_output = df._repr_html_()

# Extract cells with specific styling using regex

low_cells = re.findall(

r'<td style="background-color: #d3e9f0"[^>]*>(\d+)-low</td>', html_output

)

mid_cells = re.findall(

r'<td style="border: 1px solid #ddd"[^>]*>(\d+)-mid</td>', html_output

)

high_cells = re.findall(

r'<td style="background-color: #d9f0d3"[^>]*>(\d+)-high</td>', html_output

)

# Sort the extracted values for consistent comparison

low_cells = sorted(map(int, low_cells))

mid_cells = sorted(map(int, mid_cells))

high_cells = sorted(map(int, high_cells))

# Verify specific values have the correct styling applied

assert low_cells == [1, 2] # Values < 3

assert mid_cells == [3, 4, 5, 5] # Values 3-5

assert high_cells == [6, 8, 8] # Values > 5

# Verify the exact content with styling appears in the output

assert (

'<td style="background-color: #d3e9f0" data-test="low">1-low</td>'

in html_output

)

assert (

'<td style="background-color: #d3e9f0" data-test="low">2-low</td>'

in html_output

)

assert (

'<td style="border: 1px solid #ddd" data-test="mid">3-mid</td>' in html_output

)

assert (

'<td style="border: 1px solid #ddd" data-test="mid">4-mid</td>' in html_output

)

assert (

'<td style="background-color: #d9f0d3" data-test="high">6-high</td>'

in html_output

)

assert (

'<td style="background-color: #d9f0d3" data-test="high">8-high</td>'

in html_output

)

# Count occurrences to ensure all cells are properly styled

assert html_output.count("-low</td>") == 2 # Two low values (1, 2)

assert html_output.count("-mid</td>") == 4 # Four mid values (3, 4, 5, 5)

assert html_output.count("-high</td>") == 3 # Three high values (6, 8, 8)

# Create a custom cell builder that changes background color based on value

def custom_cell_builder(value, row, col, table_id):

# Handle numeric values regardless of their exact type

try:

num_value = int(value)

if num_value > 5: # Values > 5 get green background

return f'<td style="background-color: #d9f0d3">{value}</td>'

if num_value < 3: # Values < 3 get light blue background

return f'<td style="background-color: #d3e9f0">{value}</td>'

except (ValueError, TypeError):

pass

# Default styling for other cells

return f'<td style="border: 1px solid #ddd">{value}</td>'

# Set our custom cell builder

formatter = get_formatter()

formatter.set_custom_cell_builder(custom_cell_builder)

html_output = df._repr_html_()

# Verify our custom cell styling was applied

"""Test using a custom header builder function."""

# Create a custom header builder with tooltips

def custom_header_builder(field):

tooltips = {

"a": "Primary key column",

"b": "Secondary values",

"c": "Additional data",

}

tooltip = tooltips.get(field.name, "")

return (

f'<th style="background-color: #333; color: white" '

f'title="{tooltip}">{field.name}</th>'

)

# Set our custom header builder

formatter = get_formatter()

formatter.set_custom_header_builder(custom_header_builder)

html_output = df._repr_html_()

# Verify our custom headers were applied

assert 'title="Primary key column"' in html_output

assert 'title="Secondary values"' in html_output

"""Test combining multiple customization options together."""

# Create a dark mode style provider

class DarkModeStyleProvider:

def get_cell_style(self) -> str:

return (

"background-color: #222; color: #eee; "

"padding: 8px; border: 1px solid #444;"

)

def get_header_style(self) -> str:

return (

"background-color: #111; color: #fff; padding: 10px; "

"border: 1px solid #333;"

)

# Configure with dark mode style

configure_formatter(

max_cell_length=10,

style_provider=DarkModeStyleProvider(),

custom_css="""

""",

)

# Add type formatters for special formatting - now working with native int values

formatter = get_formatter()

formatter.register_formatter(

int,

lambda n: f'<span style="color: {"#5af" if n % 2 == 0 else "#f5a"}">{n}</span>',

)

html_output = df._repr_html_()

# Verify our customizations were applied

assert "background-color: #222" in html_output

assert "background-color: #111" in html_output

assert ".datafusion-table" in html_output

"""Test the memory and row control parameters in DataFrameHtmlFormatter."""

configure_formatter(max_memory_bytes=10, min_rows_display=1)

html_output = df._repr_html_()

# Count the number of table rows in the output

tr_count = count_table_rows(html_output)

# With a tiny memory limit of 10 bytes, the formatter should display

# the minimum number of rows (1) plus a message about truncation

assert tr_count == 2 # 1 for header row, 1 for data row

assert "data truncated" in html_output.lower()

configure_formatter(max_memory_bytes=10 * MB, min_rows_display=1)

html_output = df._repr_html_()

# With larger memory limit and min_rows=2, should display all rows

tr_count = count_table_rows(html_output)

# Table should have header row (1) + 3 data rows = 4 rows

assert tr_count == 4

# No truncation message should appear

configure_formatter(min_rows_display=2, repr_rows=2)

html_output = df._repr_html_()

tr_count = count_table_rows(html_output)

# Tabe should have header row (1) + 2 data rows = 3 rows

assert tr_count == 3

configure_formatter(min_rows_display=2, repr_rows=3)

html_output = df._repr_html_()