path1 = tmp_path / "dir1"

path2 = tmp_path / "dir2"

path1 = str(path1) if path_to_str else path1

path2 = str(path2) if path_to_str else path2

runtime = RuntimeEnvBuilder().with_disk_manager_specified(path1, path2)

config = SessionConfig().with_default_catalog_and_schema("foo", "bar")

ctx = SessionContext(config, runtime)

assert ctx is not None

db = ctx.catalog("foo").schema("bar")

path = str(tmp_path) if path_to_str else tmp_path

runtime = RuntimeEnvBuilder().with_temp_file_path(path)

config = SessionConfig().with_default_catalog_and_schema("foo", "bar")

ctx = SessionContext(config, runtime)

assert ctx is not None

db = ctx.catalog("foo").schema("bar")

runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)

config = (

SessionConfig()

.with_create_default_catalog_and_schema(enabled=True)

.with_default_catalog_and_schema("foo", "bar")

.with_target_partitions(1)

.with_information_schema(enabled=True)

.with_repartition_joins(enabled=False)

.with_repartition_aggregations(enabled=False)

.with_repartition_windows(enabled=False)

.with_parquet_pruning(enabled=False)

)

ctx = SessionContext(config, runtime)

# verify that at least some of the arguments worked

ctx.catalog("foo").schema("bar")

with pytest.raises(KeyError):

# create a RecordBatch and register it as memtable

batch = pa.RecordBatch.from_arrays(

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

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

)

ctx.register_record_batches("t", [[batch]])

assert ctx.catalog().schema().names() == {"t"}

result = ctx.sql("SELECT a+b, a-b FROM t").collect()

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

# A partition list with no record batches carries no schema, so this used to

# panic on unchecked `[0][0]` indexing. It should now raise a clear error.

with pytest.raises(ValueError, match="no record batches"):

ctx.register_record_batches("t", [[]])

# An empty outer partition list carries no schema either, and raises the same error.

with pytest.raises(ValueError, match="no record batches"):

ctx.register_record_batches("t", [])

# The schema is still recovered from a later non-empty partition.

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

ctx.register_record_batches("t2", [[], [batch]])

# create a RecordBatch and register it as memtable

batch = pa.RecordBatch.from_arrays(

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

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

)

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

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert len(tables[0]) == 33

assert tables[0].startswith("c")

# ensure that the rest of the table name contains

# only hexadecimal numbers

for c in tables[0][1:]:

# create a RecordBatch and register it as memtable

batch = pa.RecordBatch.from_arrays(

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

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

)

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

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

# create a PyArrow table

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

table = pa.Table.from_pydict(data)

# convert to DataFrame

df = ctx.from_arrow(table)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

schema = pa.schema([("a", pa.int64()), ("b", pa.int64())])

for _i in range(num_batches):

yield pa.RecordBatch.from_arrays(

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

df = ctx.from_arrow(source)

assert df

assert isinstance(df, DataFrame)

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

# create a PyArrow table

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

table = pa.Table.from_pydict(data)

# convert to DataFrame with optional name

df = ctx.from_arrow(table, name="tbl")

tables = list(ctx.catalog().schema().names())

assert df

data = {"a": [], "b": []}

schema = pa.schema([("a", pa.int32()), ("b", pa.string())])

table = pa.Table.from_pydict(data, schema=schema)

# convert to DataFrame

df = ctx.from_arrow(table)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

data = {"a": [], "b": []}

table = pa.Table.from_pydict(data)

# convert to DataFrame

df = ctx.from_arrow(table)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

# create a dataframe from Python list

data = [

{"a": 1, "b": 4},

{"a": 2, "b": 5},

{"a": 3, "b": 6},

]

df = ctx.from_pylist(data)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

# create a dataframe from Python dictionary

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

df = ctx.from_pydict(data)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

# create a dataframe from pandas dataframe

pd = pytest.importorskip("pandas")

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

pandas_df = pd.DataFrame(data)

df = ctx.from_pandas(pandas_df)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

# create a dataframe from Polars dataframe

pd = pytest.importorskip("polars")

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

polars_df = pd.DataFrame(data)

df = ctx.from_polars(polars_df)

tables = list(ctx.catalog().schema().names())

assert df

assert len(tables) == 1

assert isinstance(df, DataFrame)

assert set(df.schema().names) == {"a", "b"}

default = ctx.catalog()

public = default.schema("public")

assert public.names() == {"csv", "csv1", "csv2"}

table = public.table("csv")

ctx.register_table("csv3", table)

default = ctx.catalog()

public = default.schema("public")

assert public.names() == {"csv", "csv1", "csv2"}

table = public.table("csv")

table_df = ctx.read_table(table)

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

result = ctx.read_table(dataset).collect()

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

default = ctx.catalog()

public = default.schema("public")

assert public.names() == {"csv", "csv1", "csv2"}

ctx.deregister_table("csv")

ctx = SessionContext()

is_null = udf(

lambda x: x.is_null(),

[pa.float64()],

pa.bool_(),

volatility="immutable",

name="my_is_null",

)

ctx.register_udf(is_null)

# Verify it works

df = ctx.from_pydict({"a": [1.0, None]})

ctx.register_table("t", df.into_view())

result = ctx.sql("SELECT my_is_null(a) FROM t").collect()

assert result[0].column(0) == pa.array([False, True])

# Deregister and verify it's gone

ctx.deregister_udf("my_is_null")

with pytest.raises(ValueError):

import pyarrow.compute as pc

ctx = SessionContext()

from datafusion import Accumulator, udaf

class MySum(Accumulator):

def __init__(self):

self._sum = 0.0

def update(self, values: pa.Array) -> None:

self._sum += pc.sum(values).as_py()

def merge(self, states: list[pa.Array]) -> None:

self._sum += pc.sum(states[0]).as_py()

def state(self) -> list:

return [self._sum]

def evaluate(self) -> pa.Scalar:

return self._sum

my_sum = udaf(

MySum,

[pa.float64()],

pa.float64(),

[pa.float64()],

volatility="immutable",

name="my_sum",

)

ctx.register_udaf(my_sum)

df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})

ctx.register_table("t", df.into_view())

result = ctx.sql("SELECT my_sum(a) FROM t").collect()

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

ctx.deregister_udaf("my_sum")

with pytest.raises(ValueError):

ctx = SessionContext()

from datafusion import udwf

from datafusion.user_defined import WindowEvaluator

class MyRowNumber(WindowEvaluator):

def __init__(self):

self._row = 0

def evaluate_all(self, values, num_rows):

return pa.array(list(range(1, num_rows + 1)), type=pa.uint64())

my_row_number = udwf(

MyRowNumber,

[pa.float64()],

pa.uint64(),

volatility="immutable",

name="my_row_number",

)

ctx.register_udwf(my_row_number)

df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})

ctx.register_table("t", df.into_view())

result = ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect()

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

ctx.deregister_udwf("my_row_number")

with pytest.raises(ValueError):

import pyarrow.dataset as ds

ctx = SessionContext()

from datafusion import Table, udtf

class MyTable:

def __call__(self):

batch = pa.RecordBatch.from_pydict({"x": [1, 2, 3]})

return Table(ds.dataset([batch]))

my_table = udtf(MyTable(), "my_table")

ctx.register_udtf(my_table)

result = ctx.sql("SELECT * FROM my_table()").collect()

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

ctx.deregister_udtf("my_table")

with pytest.raises(ValueError):

df = ctx.from_pydict({"a": [1, 2]})

ctx.register_table("df_tbl", df)

result = ctx.sql("SELECT * FROM df_tbl").collect()

df = ctx.from_pydict({"a": [1, 2]})

table = df.into_view(temporary=temporary)

assert isinstance(table, Table)

if temporary:

assert table.kind == "temporary"

else:

assert table.kind == "view"

ctx.register_table("view_tbl", table)

result = ctx.sql("SELECT * FROM view_tbl").collect()

df = ctx.from_pydict({"a": [1, 2]})

table = Table(df)

assert isinstance(table, Table)

ctx.register_table("from_dataframe_tbl", table)

result = ctx.sql("SELECT * FROM from_dataframe_tbl").collect()

df = ctx.from_pydict({"a": [1, 2]})

table = Table(df.df)

assert isinstance(table, Table)

ctx.register_table("from_internal_dataframe_tbl", table)

result = ctx.sql("SELECT * FROM from_internal_dataframe_tbl").collect()

# create a RecordBatch and register it as a pyarrow.dataset.Dataset

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

assert ctx.catalog().schema().names() == {"t"}

result = ctx.sql("SELECT a+b, a-b FROM t").collect()

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

# create a RecordBatch and register it as a pyarrow.dataset.Dataset

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

assert ctx.catalog().schema().names() == {"t"}

df = ctx.sql("SELECT a+b, a-b FROM t WHERE a BETWEEN 2 and 3 AND b > 5")

# Make sure the filter was pushed down in Physical Plan

df.explain()

captured = capfd.readouterr()

assert "filter_expr=(((a >= 2) and (a <= 3)) and (b > 5))" in captured.out

result = df.collect()

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

# `datafusion-python` issue: https://github.com/apache/datafusion-python/issues/800

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

# Testing the dataframe API

df = ctx.table("t")

assert df.count() == 3

# Testing the SQL API

count = ctx.sql("SELECT COUNT(*) FROM t")

count = count.collect()

"""Ensure that pyarrow filter gets pushed down for `IsNull`"""

# create a RecordBatch and register it as a pyarrow.dataset.Dataset

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

# Make sure the filter was pushed down in Physical Plan

df = ctx.sql("SELECT a FROM t WHERE c is NULL")

df.explain()

captured = capfd.readouterr()

assert "filter_expr=is_null(c, {nan_is_null=false})" in captured.out

result = df.collect()

"""Ensure that pyarrow filter gets pushed down for timestamp"""

# Ref: https://github.com/apache/datafusion-python/issues/703

# create pyarrow dataset with no actual files

col_type = pa.timestamp("ns", "+00:00")

nyd_2000 = pa.scalar(dt.datetime(2000, 1, 1, tzinfo=dt.timezone.utc), col_type)

pa_dataset_fs = pa.fs.SubTreeFileSystem(str(tmpdir), pa.fs.LocalFileSystem())

pa_dataset_format = pa.dataset.ParquetFileFormat()

pa_dataset_partition = pa.dataset.field("a") <= nyd_2000

fragments = [

# NOTE: we never actually make this file.

# Working predicate pushdown means it never gets accessed

pa_dataset_format.make_fragment(

"1.parquet",

filesystem=pa_dataset_fs,

partition_expression=pa_dataset_partition,

)

]

pa_dataset = pa.dataset.FileSystemDataset(

fragments,

pa.schema([pa.field("a", col_type)]),

pa_dataset_format,

pa_dataset_fs,

)

ctx.register_dataset("t", pa_dataset)

# the partition for our only fragment is for a < 2000-01-01.

# so querying for a > 2024-01-01 should not touch any files

df = ctx.sql("SELECT * FROM t WHERE a > '2024-01-01T00:00:00+00:00'")

# create Arrow StructArrays to test nested data types

data = pa.StructArray.from_arrays(

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

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

)

batch = pa.RecordBatch.from_arrays(

[data],

names=["nested_data"],

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

assert ctx.catalog().schema().names() == {"t"}

df = ctx.table("t")

# This filter will not be pushed down to DatasetExec since it

# isn't supported

df = df.filter(column("nested_data")["b"] > literal(5)).select(

column("nested_data")["a"] + column("nested_data")["b"],

column("nested_data")["a"] - column("nested_data")["b"],

)

result = df.collect()

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

batch = pa.RecordBatch.from_arrays(

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

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

)

dataset = ds.dataset([batch])

ctx.register_dataset("t", dataset)

from uuid import uuid4

with pytest.raises(KeyError):

import datetime

import re

st = ctx.session_start_time()

assert isinstance(st, str)

# Truncate nanoseconds to microseconds for Python 3.10 compat

st = re.sub(r"(\.\d{6})\d+", r"\1", st)

dt = datetime.datetime.fromisoformat(st)

from datafusion.common import DFSchema

schema = DFSchema.empty()

expr = ctx.parse_sql_expr("1 + 2", schema)

df = ctx.from_pydict({"a": [1, 2, 3]})

plan = df.logical_plan()

df2 = ctx.execute_logical_plan(plan)

result = df2.collect()

"""The extension options hook names the capsule it was handed.

class WrongCapsule:

def __datafusion_extension_options__(self):

return ctx.__datafusion_task_context_provider__()

with pytest.raises(ValueError, match="datafusion_extension_options"):

"""A getter that refuses the session is named, not left as a bare TypeError.

class PreSessionCodec:

def __datafusion_logical_extension_codec__(self):

msg = "should never be called"

raise AssertionError(msg)

with pytest.raises(ImportError, match="__datafusion_logical_extension_codec__"):

ctx.with_logical_extension_codec(PreSessionCodec())

with pytest.raises(ImportError) as excinfo:

ctx.with_logical_extension_codec(PreSessionCodec())

assert isinstance(excinfo.value.__cause__, TypeError)

"""A correctly-signed getter's own TypeError must survive unchanged.

class RaisesTypeError:

def __datafusion_logical_extension_codec__(self, session):

msg = "bad cast inside the getter"

raise TypeError(msg)

with pytest.raises(TypeError, match="bad cast inside the getter"):

"""Anything that is not a TypeError was never a signature problem."""

class RaisesValueError:

def __datafusion_logical_extension_codec__(self, session):

msg = "something else entirely"

raise ValueError(msg)

with pytest.raises(ValueError, match="something else entirely"):

capsule = ctx.__datafusion_query_planner__()

get_name = ctypes.pythonapi.PyCapsule_GetName

get_name.argtypes = [ctypes.py_object]

get_name.restype = ctypes.c_char_p

assert get_name(capsule) == b"datafusion_query_planner"

ctx.register_record_batches(

"query_planner_test",

[[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],

)

ctx.set_query_planner(capsule)

assert ctx.table_exist("query_planner_test")

batches = ctx.sql("SELECT 1 AS value").collect()

"""The planner is written into the existing session, not a copy of it.

ctx.register_record_batches(

"registered_before",

[[pa.RecordBatch.from_pydict({"value": [1]})]],

)

before = udf(

lambda arr: arr,

[pa.int64()],

pa.int64(),

volatility="immutable",

name="registered_before",

)

ctx.register_udf(before)

ctx.set_query_planner(ctx.__datafusion_query_planner__())

ctx.register_record_batches(

"registered_after",

[[pa.RecordBatch.from_pydict({"value": [2]})]],

)

after = udf(

lambda arr: arr,

[pa.int64()],

pa.int64(),

volatility="immutable",

name="registered_after",

)

ctx.register_udf(after)

assert ctx.table_exist("registered_before")

assert ctx.table_exist("registered_after")

assert ctx.sql("SELECT registered_before(1)").collect()

"""A context derived before the install still plans through the planner.

sibling = ctx.with_python_udf_inlining(enabled=False)

ctx.register_record_batches(

"shared_planner_test",

[[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]],

)

ctx.set_query_planner(ctx.__datafusion_query_planner__())

assert sibling.table_exist("shared_planner_test")

"""Wraps a codec capsule in an object that can name itself.

def __init__(self, capsule, codec_id):

self._capsule = capsule

self.__datafusion_codec_id__ = codec_id

def __datafusion_logical_extension_codec__(self, session=None):

return self._capsule

def __datafusion_physical_extension_codec__(self, session=None):

"""Contributes decline-all codecs exported from an unrelated session.

def __init__(self, prefix="my_library"):

self.exporter = SessionContext()

self.prefix = prefix

self.bound_ctx = None

def __datafusion_session_components__(self, ctx):

self.bound_ctx = ctx

return SessionExtensionComponents(

logical_extension_codecs=(

_NamedCodec(

self.exporter.__datafusion_logical_extension_codec__(),

f"{self.prefix}.logical",

),

),

physical_extension_codecs=(

_NamedCodec(

self.exporter.__datafusion_physical_extension_codec__(),

f"{self.prefix}.physical",

),

),

"""Contributes a planner, recording the fallback it was handed.

def __init__(self, calls=None):

self.fallbacks = []

self.planner_ctx = None

# Shared list the hooks append themselves to, so a test can assert the

# order they ran in rather than only that each ran.

self.calls = [] if calls is None else calls

def __datafusion_session_planner__(self, ctx, fallback):

self.calls.append(self)

self.planner_ctx = ctx

self.fallbacks.append(fallback)

"""No extensions installs nothing and returns a handle on this session.

ctx.register_record_batches(

"empty_extensions_test",

[[pa.RecordBatch.from_pydict({"value": [1]})]],

)

result = ctx.with_extensions()

assert result.session_id() == ctx.session_id()

assert result.table_exist("empty_extensions_test")

"""The empty case must not disturb a planner the session already has.

extension = _CodecOnlyExtension()

installed = ctx.with_extensions(extension, _PlannerExtension())

result = installed.with_extensions()

assert result.logical_extension_codec_ids() == (

installed.logical_extension_codec_ids()

)

batches = result.sql("SELECT 1 AS value").collect()