# 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 pathlib import Path

from typing import Any

import pyarrow as pa

import pyarrow.parquet as pq

import pytest

from datafusion import (

DataFrame,

ExplainFormat,

InsertOp,

ParquetColumnOptions,

ParquetWriterOptions,

RecordBatch,

SessionContext,

WindowFrame,

column,

literal,

udf,

)

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

pa_cffi = pytest.importorskip("pyarrow.cffi")

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

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

This fixture creates 10 batches of 10,000 rows each (100,000 rows total),

ensuring the DataFrame spans multiple batches. This is essential for testing

that memory limits actually cause early stream termination rather than

truncating all collected data.

"""

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

return ctx.table("large_multi_batch_data")

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

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

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

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_filter_string_predicates(df):

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

assert result_strings.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_equivalent(df):

df1 = df.filter("a > 1").to_pydict()

df2 = df.filter(column("a") > literal(1)).to_pydict()

assert df1 == df2

def test_filter_string_invalid(df):

with pytest.raises(Exception) as excinfo:

df.filter("this is not valid sql").collect()

assert "Expected Expr" not in str(excinfo.value)

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_sql_expression(df):

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

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

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_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_str(df):

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

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

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_iter_batches(df):

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]

assert pa_batch.column(2).to_pylist() == [8, 5, 8]

def test_iter_returns_datafusion_recordbatch(df):

for batch in df:

assert isinstance(batch, RecordBatch)

def test_execute_stream_basic(df):

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]

assert pa_batch.column(2).to_pylist() == [8, 5, 8]

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

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

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("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_full_with_drop_duplicate_keys():

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

assert result.column(0).to_pylist() == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

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