# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

"""Tests for predicate pushdown"""

import random

from datetime import date, datetime, timedelta

from decimal import Decimal

from pathlib import Path

import lance

import numpy as np

import pandas as pd

import pandas.testing as tm

import pyarrow as pa

import pyarrow.compute as pc

import pytest

from lance.vector import vec_to_table

def create_table(nrows=100):

intcol = pa.array(range(nrows))

floatcol = pa.array(np.arange(nrows) * 2 / 3, type=pa.float32())

arr = np.arange(nrows) < nrows / 2

structcol = pa.StructArray.from_arrays(

[

pa.array(arr, type=pa.bool_()),

pa.array([date(2021, 1, 1) + timedelta(days=i) for i in range(nrows)]),

pa.array([datetime(2021, 1, 1) + timedelta(hours=i) for i in range(nrows)]),

],

names=["bool", "date", "dt"],

)

random.seed(42)

def gen_str(n):

return "".join(random.choices("abc", k=n))

string_col = pa.array([gen_str(2) for _ in range(nrows)])

decimal_col = pa.array([Decimal(f"{str(i)}.000") for i in range(nrows)])

tbl = pa.Table.from_arrays(

[intcol, floatcol, structcol, string_col, decimal_col],

names=["int", "float", "rec", "str", "decimal"],

)

return tbl

@pytest.fixture()

def dataset(tmp_path: Path):

tbl = create_table()

yield lance.write_dataset(tbl, tmp_path)

def test_simple_predicates(dataset):

predicates = [

pc.field("int") >= 50,

pc.field("int") == 50,

pc.field("int") != 50,

pc.field("float") < 30.0,

pc.field("float") > 30.0,

pc.field("float") <= 30.0,

pc.field("float") >= 30.0,

pc.field("str") != "aa",

pc.field("str") == "aa",

(pc.field("int") >= 50) & (pc.field("int") < 200),

pc.invert(pc.field("int") >= 50),

pc.is_null(pc.field("int")),

pc.field("int") + 3 >= 50,

pc.is_valid(pc.field("int")),

]

# test simple

for expr in predicates:

assert dataset.to_table(filter=expr) == dataset.to_table().filter(expr)

def test_sql_predicates(dataset):

# Predicate and expected number of rows

predicates_nrows = [

("int >= 50", 50),

("int = 50", 1),

("int != 50", 99),

("int BETWEEN 50 AND 60", 11),

("float < 30.0", 45),

("str = 'aa'", 16),

("str in ('aa', 'bb')", 26),

("rec.bool", 50),

("rec.bool is true", 50),

("rec.bool is not true", 50),

("rec.bool is false", 50),

("rec.bool is not false", 50),

("rec.date = cast('2021-01-01' as date)", 1),

("rec.dt = cast('2021-01-01 00:00:00' as timestamp(6))", 1),

("rec.dt = cast('2021-01-01 00:00:00' as timestamp)", 1),

("rec.dt = cast('2021-01-01 00:00:00' as datetime(6))", 1),

("rec.dt = cast('2021-01-01 00:00:00' as datetime)", 1),

("rec.dt = TIMESTAMP '2021-01-01 00:00:00'", 1),

("rec.dt = TIMESTAMP(6) '2021-01-01 00:00:00'", 1),

("rec.date = DATE '2021-01-01'", 1),

("rec.date >= cast('2021-01-31' as date)", 70),

("cast(rec.date as string) = '2021-01-01'", 1),

("decimal = DECIMAL(5,3) '12.000'", 1),

("decimal >= DECIMAL(5,3) '50.000'", 50),

]

for expr, expected_num_rows in predicates_nrows:

assert dataset.to_table(filter=expr).num_rows == expected_num_rows

def test_sql_current_date(tmp_path: Path):

table = pa.table(

{"date": pa.array([date(2020, 1, 1), date(2020, 1, 2)], type=pa.date32())}

)

dataset = lance.write_dataset(table, tmp_path / "current_date")

filtered = dataset.to_table(filter="date <= current_date()")

assert filtered.equals(dataset.to_table())

def test_illegal_predicates(dataset):

bad_parse = [

"str BETWEEN 10 AND 20",

"str > 10",

"str AN",

"🥞",

]

for expr in bad_parse:

with pytest.raises(ValueError, match="Invalid user input: *"):

dataset.to_table(filter=expr)

with pytest.raises(ValueError, match="No field named foo"):

dataset.to_table(filter="foo = 7")

with pytest.raises(ValueError, match="does not return a boolean"):

dataset.to_table(filter="int")

def test_compound(dataset):

predicates = [

pc.field("int") >= 50,

pc.field("float") < 90.0,

pc.field("str") == "aa",

]

# test compound

for expr in predicates:

for other_expr in predicates:

compound = expr & other_expr

assert dataset.to_table(filter=compound) == dataset.to_table().filter(

compound

)

compound = expr | other_expr

assert dataset.to_table(filter=compound) == dataset.to_table().filter(

compound

)

def test_match(tmp_path: Path, provide_pandas: bool):

array = pa.array(["aaa", "bbb", "abc", "bca", "cab", "cba"])

table = pa.Table.from_arrays([array], names=["str"])

dataset = lance.write_dataset(table, tmp_path / "test_match")

result = dataset.to_table(filter="str LIKE 'a%'").to_pandas()

pd.testing.assert_frame_equal(result, pd.DataFrame({"str": ["aaa", "abc"]}))

result = dataset.to_table(filter="str NOT LIKE 'a%'").to_pandas()

pd.testing.assert_frame_equal(

result, pd.DataFrame({"str": ["bbb", "bca", "cab", "cba"]})

)

result = dataset.to_table(filter="regexp_match(str, 'c.+')").to_pandas()

pd.testing.assert_frame_equal(result, pd.DataFrame({"str": ["bca", "cab", "cba"]}))

def test_escaped_name(tmp_path: Path, provide_pandas: bool):

table = pa.table({"silly :name": pa.array([0, 1, 2])})

dataset = lance.write_dataset(table, tmp_path / "test_escaped_name")

dataset = lance.dataset(tmp_path / "test_escaped_name")

result = dataset.to_table(filter="`silly :name` > 1").to_pandas()

pd.testing.assert_frame_equal(result, pd.DataFrame({"silly :name": [2]}))

# nested case

table = pa.table({"outer field": pa.array([{"inner field": i} for i in range(3)])})

dataset = lance.write_dataset(table, tmp_path / "test_escaped_name_nested")

dataset = lance.dataset(tmp_path / "test_escaped_name_nested")

result = dataset.to_table(filter="`outer field`.`inner field` > 1").to_pandas()

pd.testing.assert_frame_equal(

result, pd.DataFrame({"outer field": [{"inner field": 2}]})

)

# test uppercase name

table = pa.table({"ALLCAPSNAME": pa.array([0, 1]), "other": pa.array([2, 3])})

_ = lance.write_dataset(table, tmp_path / "test_uppercase_name")

dataset = lance.dataset(tmp_path / "test_uppercase_name")

result = dataset.to_table(filter="`ALLCAPSNAME` > 0").to_pandas()

pd.testing.assert_frame_equal(

result, pd.DataFrame([{"ALLCAPSNAME": 1, "other": 3}])

)

table = pa.table(

{"Nested with Space": pa.array([{"Inner With Caps": i} for i in range(3)])}

)

_ = lance.write_dataset(table, tmp_path / "test_escaped_name_nested_and_capped")

dataset = lance.dataset(tmp_path / "test_escaped_name_nested_and_capped")

result = dataset.to_table(

filter="`Nested with Space`.`Inner With Caps` > 1"

).to_pandas()

pd.testing.assert_frame_equal(

result, pd.DataFrame({"Nested with Space": [{"Inner With Caps": 2}]})

)

def test_functions(tmp_path: Path):

# Ensure that we can use complex functions

table = pa.table(

{"genres": [["action", "comedy"], ["anime", "drama"], ["adventure"]]}

)

expected = table.slice(1, 2)

dataset = lance.write_dataset(table, tmp_path / "test_neg_expr")

assert (

dataset.to_table(filter="array_has_any(genres, Array['anime', 'adventure'])")

== expected

)

expected = table.slice(0, 1)

assert dataset.to_table(filter="array_contains(genres, 'comedy')") == expected

def test_negative_expressions(tmp_path: Path):

table = pa.table({"x": [-1, 0, 1, 1], "y": [1, 2, 3, 4]})

dataset = lance.write_dataset(table, tmp_path / "test_neg_expr")

filters_expected = [

("x = -1", [-1]),

("x > -1", [0, 1, 1]),

("x = 1 * -1", [-1]),

("x <= 2 + -2 ", [-1, 0]),

("x = y - 2", [-1, 0, 1]),

]

for filter, expected in filters_expected:

assert dataset.scanner(filter=filter).to_table()["x"].to_pylist() == expected

def create_table_for_duckdb(nvec=10000, ndim=768):

mat = np.random.randn(nvec, ndim)

price = (np.random.rand(nvec) + 1) * 100

def gen_str(n):

return "".join(random.choices("abc"))

meta = np.array([gen_str(1) for _ in range(nvec)])

tbl = (

vec_to_table(data=mat)

.append_column("price", pa.array(price))

.append_column("meta", pa.array(meta))

.append_column("id", pa.array(range(nvec)))

)

return tbl

def test_datatypes(tmp_path):

table = pa.table(

{

"binary": pa.array([b"abc", None], type=pa.binary()),

"largebin": pa.array([b"abc", None], type=pa.large_binary()),

}

)

dataset = lance.write_dataset(table, tmp_path)

for filter, expected_matches in [

("binary = X'616263'", 1),

("binary is NULL", 1),

("largebin = X'616263'", 1),

("largebin is NULL", 1),

]:

assert dataset.count_rows(filter=filter) == expected_matches

def test_duckdb(tmp_path):

duckdb = pytest.importorskip("duckdb")

tbl = create_table_for_duckdb()

ds = lance.write_dataset(tbl, str(tmp_path)) # noqa: F841

actual = duckdb.query("SELECT id, meta, price FROM ds WHERE id==1000").to_df()

expected = duckdb.query("SELECT id, meta, price FROM ds").to_df()

expected = expected[expected.id == 1000].reset_index(drop=True)

tm.assert_frame_equal(actual, expected)

actual = duckdb.query("SELECT id, meta, price FROM ds WHERE id=1000").to_df()

expected = duckdb.query("SELECT id, meta, price FROM ds").to_df()

expected = expected[expected.id == 1000].reset_index(drop=True)

tm.assert_frame_equal(actual, expected)

actual = duckdb.query(

"SELECT id, meta, price FROM ds WHERE price>20.0 and price<=90"

).to_df()

expected = duckdb.query("SELECT id, meta, price FROM ds").to_df()

expected = expected[(expected.price > 20.0) & (expected.price <= 90)].reset_index(

drop=True

)

tm.assert_frame_equal(actual, expected, check_dtype=False)

actual = duckdb.query("SELECT id, meta, price FROM ds WHERE meta=='aa'").to_df()

expected = duckdb.query("SELECT id, meta, price FROM ds").to_df()

expected = expected[expected.meta == "aa"].reset_index(drop=True)

tm.assert_frame_equal(actual, expected, check_dtype=False)

def test_struct_field_order(tmp_path):

"""

This test regresses some old behavior where the order of struct fields would get

messed up due to late materialization and we would get {y,x} instead of {x,y}

"""

data = pa.table({"struct": [{"x": i, "y": i} for i in range(10)]})

dataset = lance.write_dataset(data, tmp_path)

for late_materialization in [True, False]:

result = dataset.to_table(

filter="struct.y > 5", late_materialization=late_materialization

)

expected = pa.table({"struct": [{"x": i, "y": i} for i in range(6, 10)]})

assert result == expected

@pytest.mark.skip(

reason="enable this in recurring test https://github.com/lance-format/lance/pull/4190"

" as it requires release mode"

)

def test_filter_depth_limit():

column_name = "a_very_long_column_name"

ds = lance.write_dataset(pa.table({column_name: [1, 2]}), "memory://")

ds.create_scalar_index(column_name, "BTREE")

filter = " AND ".join([f"{column_name} = {i}" for i in range(500)])

ds.to_table(filter=filter)

with pytest.raises(ValueError, match="the filter expression is too long"):

filter = " AND ".join([f"{column_name} = {i}" for i in range(501)])

ds.to_table(filter=filter)