# Licensed to the Apache Software Foundation (ASF) under one

# or more contributor license agreements. See the NOTICE file

# distributed with this work for additional information

# regarding copyright ownership. The ASF licenses this file

# to you 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 ctypes

import datetime

import itertools

import os

import re

import threading

import time

from typing import Any

import pyarrow as pa

import pyarrow.parquet as pq

import pytest

from datafusion import (

DataFrame,

InsertOp,

ParquetColumnOptions,

ParquetWriterOptions,

SessionContext,

WindowFrame,

column,

literal,

)

from datafusion import (

col as df_col,

)

from datafusion import (

functions as f,

)

from datafusion.dataframe import DataFrameWriteOptions

from datafusion.dataframe_formatter import (

DataFrameHtmlFormatter,

configure_formatter,

get_formatter,

reset_formatter,

)

from datafusion.expr import EXPR_TYPE_ERROR, Window

from pyarrow.csv import write_csv

MB = 1024 * 1024

@pytest.fixture

def ctx():

return SessionContext()

@pytest.fixture

def df(ctx):

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

)

return ctx.from_arrow(batch)

@pytest.fixture

def large_df():

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)

return ctx.from_arrow(batch)

@pytest.fixture

def struct_df():

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

)

return ctx.create_dataframe([[batch]])

@pytest.fixture

def nested_df():

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

)

return ctx.create_dataframe([[batch]])

@pytest.fixture

def aggregate_df():

ctx = SessionContext()

ctx.register_csv("test", "testing/data/csv/aggregate_test_100.csv")

return ctx.sql("select c1, sum(c2) from test group by c1")

@pytest.fixture

def partitioned_df():

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

)

return ctx.create_dataframe([[batch]])

@pytest.fixture

def clean_formatter_state():

"""Reset the HTML formatter after each test."""

reset_formatter()

@pytest.fixture

def null_df():

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

],

)

return ctx.create_dataframe([[batch]])

# custom style for testing with html formatter

class CustomStyleProvider:

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;"

)

def count_table_rows(html_content: str) -> int:

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

Args:

html_content: HTML string to analyze

Returns:

Number of table rows found (number of <tr> tags)

"""

return len(re.findall(r"<tr", html_content))

def test_select(df):

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

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

def test_select_exprs(df):

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

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

def test_drop_quoted_columns():

ctx = SessionContext()

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

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

# Both should work

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

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

def test_select_mixed_expr_string(df):

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

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

def test_select_unsupported(df):

with pytest.raises(

TypeError,

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

):

df.select(1)

def test_filter(df):

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

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

def test_parse_sql_expr(df):

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

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

def test_show_empty(df, capsys):

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

df_empty.show()

captured = capsys.readouterr()

assert "DataFrame has no rows" in captured.out

def test_sort(df):

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

assert table.to_pydict() == expected

def test_sort_string_and_expression_equivalent(df):

from datafusion import col

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

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

assert result_str == result_expr

def test_sort_unsupported(df):

with pytest.raises(

TypeError,

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

):

df.sort(1)

def test_aggregate_string_and_expression_equivalent(df):

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

assert result_str == result_expr

def test_aggregate_tuple_group_by(df):

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

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

assert result_tuple == result_list

def test_aggregate_tuple_aggs(df):

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

result_tuple = df.aggregate("a", (f.count(),)).sort("a").to_pydict()

assert result_tuple == result_list

def test_filter_string_unsupported(df):

with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):

df.filter("a > 1")

def test_drop(df):

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

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

def test_limit(df):

df = df.limit(1)

# execute and collect the first (and only) batch

result = df.collect()[0]

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

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

def test_limit_with_offset(df):

# 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

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

def test_head(df):

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

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

def test_tail(df):

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

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

def test_with_column(df):

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

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

def test_with_column_invalid_expr(df):

with pytest.raises(

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

):

df.with_column("c", "a")

def test_with_columns(df):

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

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

def test_with_columns_invalid_expr(df):

with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):

df.with_columns("a")

with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):

df.with_columns(c="a")

with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):

df.with_columns(["a"])

with pytest.raises(TypeError, match=re.escape(EXPR_TYPE_ERROR)):

df.with_columns(c=["a"])

def test_cast(df):

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

)

assert df.schema() == expected

def test_with_column_renamed(df):

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"

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

def test_unnest(nested_df):

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

assert result.column(1) == pa.array([7, 8, 8, 9, 9, 9, 10])

def test_unnest_without_nulls(nested_df):

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

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

@pytest.mark.filterwarnings("ignore:`join_keys`:DeprecationWarning")

def test_join():

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

assert table.to_pydict() == expected

def test_join_invalid_params():

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

):

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

def test_join_on():

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

assert table.to_pydict() == expected

def test_join_on_invalid_expr():

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\(\)"

):

df.join_on(df1, "a")

def test_aggregate_invalid_aggs(df):

with pytest.raises(

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

):

df.aggregate([], "a")

def test_distinct():

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"))

assert df_a.collect() == df_b.collect()

data_test_window_functions = [

(

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

),

(

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

),

(

"cume_dist_w_params_no_lists",

f.round(

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

literal(3),

),

[0.25, 0.5, 0.75, 1.0, 0.333, 0.667, 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],

),

(

"ntile_w_params_no_lists",

f.ntile(2, order_by=column("b"), 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],

),

(

"lead_w_params_no_lists",

f.lead(

column("b"),

shift_offset=2,

default_value=-1,

order_by=column("b"),

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

),

(

"lag_w_params_no_lists",

f.lag(

column("b"),

shift_offset=2,

default_value=-1,

order_by=column("b"),

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

),

(

"first_value_without_list_args",

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

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

),

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

),

(

"first_value_order_by_string",

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

Window(partition_by=[column("c")], order_by="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],

),

]

@pytest.mark.parametrize(("name", "expr", "result"), data_test_window_functions)

def test_window_functions(partitioned_df, name, expr, result):

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,

}

assert table.sort_by("a").to_pydict() == expected

@pytest.mark.parametrize("partition", ["c", df_col("c")])

def test_rank_partition_by_accepts_string(partitioned_df, partition):

"""Passing a string to partition_by should match using col()."""

df = partitioned_df.select(

f.rank(order_by=column("a"), partition_by=partition).alias("r")

)

table = pa.Table.from_batches(df.sort(column("a")).collect())

assert table.column("r").to_pylist() == [1, 2, 3, 4, 1, 2, 3]

@pytest.mark.parametrize("partition", ["c", df_col("c")])

def test_window_partition_by_accepts_string(partitioned_df, partition):

"""Window.partition_by accepts string identifiers."""

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

Window(partition_by=partition, order_by=column("b"))

)

df = partitioned_df.select(expr.alias("fv"))

table = pa.Table.from_batches(df.sort(column("a")).collect())

assert table.column("fv").to_pylist() == [1, 1, 1, 1, 5, 5, 5]

@pytest.mark.parametrize(

("units", "start_bound", "end_bound"),

[

(units, start_bound, end_bound)

for units in ("rows", "range")

for start_bound in (None, 0, 1)

for end_bound in (None, 0, 1)

]

+ [

("groups", 0, 0),

],

)

def test_valid_window_frame(units, start_bound, end_bound):

WindowFrame(units, start_bound, end_bound)

@pytest.mark.parametrize(

("units", "start_bound", "end_bound"),

[

("invalid-units", 0, None),

("invalid-units", None, 0),

("invalid-units", None, None),

("groups", None, 0),

("groups", 0, None),

("groups", None, None),

],

)

def test_invalid_window_frame(units, start_bound, end_bound):

with pytest.raises(NotImplementedError, match=f"(?i){units}"):

WindowFrame(units, start_bound, end_bound)

def test_window_frame_defaults_match_postgres(partitioned_df):

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