df = test_ctx.sql("select c1, 123, c1 < 123 from test")

plan = df.logical_plan()

plan = plan.to_variant()

assert isinstance(plan, Projection)

expr = plan.projections()

col1 = expr[0].to_variant()

assert isinstance(col1, Column)

assert col1.name() == "c1"

assert col1.qualified_name() == "test.c1"

col2 = expr[1].to_variant()

assert isinstance(col2, Literal)

assert col2.data_type() == "Int64"

assert col2.value_i64() == 123

col3 = expr[2].to_variant()

assert isinstance(col3, BinaryExpr)

assert isinstance(col3.left().to_variant(), Column)

assert col3.op() == "<"

assert isinstance(col3.right().to_variant(), Literal)

plan = plan.input()[0].to_variant()

df = test_ctx.sql("select c1 from test WHERE c1 > 5")

plan = df.logical_plan()

plan = plan.to_variant()

assert isinstance(plan, Projection)

plan = plan.input()[0].to_variant()

df = test_ctx.sql("select c1 from test LIMIT 10")

plan = df.logical_plan()

plan = plan.to_variant()

assert isinstance(plan, Limit)

assert "Skip: None" in str(plan)

df = test_ctx.sql("select c1 from test LIMIT 10 OFFSET 5")

plan = df.logical_plan()

plan = plan.to_variant()

assert isinstance(plan, Limit)

df = test_ctx.sql("select c1, count(*) from test group by c1")

plan = df.logical_plan()

projection = plan.to_variant()

assert isinstance(projection, Projection)

aggregate = projection.input()[0].to_variant()

assert isinstance(aggregate, Aggregate)

col1 = aggregate.group_by_exprs()[0].to_variant()

assert isinstance(col1, Column)

assert col1.name() == "c1"

assert col1.qualified_name() == "test.c1"

col2 = aggregate.aggregate_exprs()[0].to_variant()

df = test_ctx.sql("select c1 from test order by c1")

plan = df.logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

batch = pa.RecordBatch.from_arrays(

[

pa.array([1, 2, 3]),

pa.array(["alpha", "beta", "gamma"], type=pa.string_view()),

],

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

)

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

assert df.filter(col("a") == 1).count() == 1

assert df.filter(col("a") != 1).count() == 2

assert df.filter(col("a") >= 1).count() == 3

assert df.filter(col("a") > 1).count() == 2

assert df.filter(col("a") <= 3).count() == 3

assert df.filter(col("a") < 3).count() == 2

assert df.filter(col("b") == "beta").count() == 1

assert df.filter(col("b") != "beta").count() == 2

# Taken from https://github.com/apache/datafusion-python/issues/781

from datafusion import SessionContext

from datafusion.expr import Filter

def traverse_logical_plan(plan):

cur_node = plan.to_variant()

if isinstance(cur_node, Filter):

return cur_node.predicate().to_variant()

if hasattr(plan, "inputs"):

for input_plan in plan.inputs():

res = traverse_logical_plan(input_plan)

if res is not None:

return res

return None

ctx = SessionContext()

data = {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}

ctx.from_pydict(data, name="table1")

query = "SELECT * FROM table1 t1 WHERE t1.name IN ('dfa', 'ad', 'dfre', 'vsa')"

logical_plan = ctx.sql(query).optimized_logical_plan()

variant = traverse_logical_plan(logical_plan)

assert variant is not None

assert variant.expr().to_variant().qualified_name() == "table1.name"

assert (

str(variant.list())

== '[Expr(Utf8("dfa")), Expr(Utf8("ad")), Expr(Utf8("dfre")), Expr(Utf8("vsa"))]' # noqa: E501

)

case_builder = functions.when(lit(True), lit(1))

with pytest.raises(Exception) as exc_info:

_ = case_builder.otherwise(lit("bad"))

err_msg = str(exc_info.value)

assert "multiple data types" in err_msg

assert "CaseBuilder has already been consumed" not in err_msg

_ = case_builder.end()

err_msg = str(exc_info.value)

assert "multiple data types" in err_msg

ctx = SessionContext()

df = ctx.from_pydict({"flag": [False]}, name="tbl")

case_builder = functions.when(col("flag"), lit("true"))

expr_default_one = case_builder.otherwise(lit("default-1")).alias("result")

result_one = df.select(expr_default_one).collect()

assert result_one[0].column(0).to_pylist() == ["default-1"]

expr_default_two = case_builder.otherwise(lit("default-2")).alias("result")

result_two = df.select(expr_default_two).collect()

assert result_two[0].column(0).to_pylist() == ["default-2"]

expr_end_one = case_builder.end().alias("result")

end_one = df.select(expr_end_one).collect()

ctx = SessionContext()

df = ctx.from_pydict(

{

"flag": [True, False, False, False],

"value": [1, 15, 25, 5],

},

name="tbl",

)

base_builder = functions.when(col("flag"), lit("flag-true"))

first_builder = base_builder.when(col("value") > lit(10), lit("gt10"))

second_builder = base_builder.when(col("value") > lit(20), lit("gt20"))

first_builder = first_builder.when(lit(True), lit("final-one"))

expr_first = first_builder.otherwise(lit("fallback-one")).alias("first")

expr_second = second_builder.otherwise(lit("fallback-two")).alias("second")

result = df.select(expr_first, expr_second).collect()[0]

assert result.column(0).to_pylist() == [

"flag-true",

"gt10",

"gt10",

"final-one",

]

assert result.column(1).to_pylist() == [

"flag-true",

"fallback-two",

"gt20",

"fallback-two",

case_builder = functions.when(lit(True), lit(1))

def build_expr(value: int) -> bool:

builder = case_builder.when(lit(True), lit(value))

builder.otherwise(lit(value))

return True

with ThreadPoolExecutor(max_workers=8) as executor:

futures = [executor.submit(build_expr, idx) for idx in range(16)]

results = [future.result() for future in futures]

assert all(results)

# Ensure the shared builder remains usable after concurrent `when` calls.

follow_up_builder = case_builder.when(lit(True), lit(42))

assert isinstance(follow_up_builder, type(case_builder))

ctx = SessionContext()

data = {

"array_values": [[1, 2, 3], [4, 5], [6], []],

"struct_values": [

{"name": "Alice", "age": 15},

{"name": "Bob", "age": 14},

{"name": "Charlie", "age": 13},

{"name": None, "age": 12},

],

}

df = ctx.from_pydict(data, name="table1")

names = df.select(col("struct_values")["name"].alias("name")).collect()

names = [r.as_py() for rs in names for r in rs["name"]]

array_values = df.select(col("array_values")[1].alias("value")).collect()

array_values = [r.as_py() for rs in array_values for r in rs["value"]]

assert names == ["Alice", "Bob", "Charlie", None]

import warnings

expr = col("foo")

with warnings.catch_warnings(record=True) as w:

# Cause all warnings to always be triggered

warnings.simplefilter("always")

# should trigger warning

name = expr.display_name()

# Verify some things

assert len(w) == 1

assert issubclass(w[-1].category, DeprecationWarning)

assert "deprecated" in str(w[-1].message)

# returns appropriate result

assert name == expr.schema_name()

ctx = SessionContext()

# create a RecordBatch and a new DataFrame from it

batch = pa.RecordBatch.from_arrays(

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

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

)

df = df.select(

col("a").fill_null(100).alias("a"),

col("b").fill_null(25).alias("b"),

col("c").fill_null(1234).alias("c"),

)

df.show()

result = df.collect()[0]

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

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

ctx = SessionContext()

ctx.sql("CREATE TABLE foo (a int, b int)").collect()

df = ctx.sql("COPY foo TO bar STORED AS CSV")

plan = df.logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

ctx.sql("CREATE TABLE foo (a int, b int)").collect()

plan = ctx.sql("create index idx on foo (a)").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

ctx.sql("CREATE TABLE foo (a int, b int)").collect()

plan = ctx.sql("describe foo").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

ctx.sql("CREATE TABLE foo (a int, b int)").collect()

plan = ctx.sql("insert into foo values (1, 2)").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

plan = ctx.sql("drop schema cat").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

plan = ctx.sql(

"""

).logical_plan()

plan = plan.inputs()[0].inputs()[0].to_variant()

ctx = SessionContext()

plan = ctx.sql("values (1, 'foo'), (2, 'bar')").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

plan = ctx.sql("START TRANSACTION").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

plan = ctx.sql("COMMIT").logical_plan()

plan = plan.to_variant()

ctx = SessionContext()

data = {

"array_values": [[1, 2, 3], [4, 5], [6], []],

"struct_values": [

{"name": "Alice", "age": 15},

{"name": "Bob", "age": 14},

{"name": "Charlie", "age": 13},

{"name": None, "age": 12},

],

}

df = ctx.from_pydict(data, name="table1")

names = df.select(col.struct_values["name"].alias("name")).collect()

names = [r.as_py() for rs in names for r in rs["name"]]

array_values = df.select(col.array_values[1].alias("value")).collect()

array_values = [r.as_py() for rs in array_values for r in rs["value"]]

assert names == ["Alice", "Bob", "Charlie", None]

batch = pa.RecordBatch.from_arrays(

[

pa.array([-0.75, 0.5, 0.0, None], type=pa.float64()),

pa.array([-30, 42, 0, None], type=pa.int64()),

pa.array(["Hello", " world ", "!", None], type=pa.string_view()),

pa.array(

[[-1, 1, 0], [5, 10, 15, 20], [], None], type=pa.list_(pa.int64())

),

pa.array([-0.75, float("nan"), 0.0, None], type=pa.float64()),

pa.array(

[[[-1, 1, 0], [4, 4]], [[5, 10, 15, 20], [3]], [[]], [None]],

type=pa.list_(pa.list_(pa.int64())),

),

],

names=["a", "b", "c", "d", "e", "f"],

)

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

result = df.collect()

assert len(result) == 1

result = (

ctx.from_pydict({"a": [1]})

.select(

lit(1).alias("no_metadata"),

lit_with_metadata(2, {"key1": "value1"}).alias("lit_with_metadata_fn"),

literal_with_metadata(3, {"key2": "value2"}).alias(

"literal_with_metadata_fn"

),

)

.collect()

)

expected_schema = pa.schema(

[

pa.field("no_metadata", pa.int64(), nullable=False),

pa.field(

"lit_with_metadata_fn",

pa.int64(),

nullable=False,

metadata={"key1": "value1"},

),

pa.field(

"literal_with_metadata_fn",

pa.int64(),

nullable=False,

metadata={"key2": "value2"},

),

]

)

expected = pa.RecordBatch.from_pydict(

{

"no_metadata": pa.array([1]),

"lit_with_metadata_fn": pa.array([2]),

"literal_with_metadata_fn": pa.array([3]),

},

schema=expected_schema,

)

assert result[0] == expected

# Testing result[0].schema == expected_schema does not check each key/value pair

# so we want to explicitly test these

for expected_field in expected_schema:

actual_field = result[0].schema.field(expected_field.name)

e = col("a")

assert ensure_expr(e) is e.expr

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