def test_query_env_changing():
df = DataFrame()
df["a"] = aa = np.arange(100)
expr = "a < @c"
# first attempt
c = 10
got = df.query(expr)
np.testing.assert_array_equal(aa[aa < c], got["a"].to_array())
# change env
c = 50
got = df.query(expr)
np.testing.assert_array_equal(aa[aa < c], got["a"].to_array())
def test_query_local_dict():
df = DataFrame()
df["a"] = aa = np.arange(100)
expr = "a < @val"
got = df.query(expr, local_dict={"val": 10})
np.testing.assert_array_equal(aa[aa < 10], got["a"].to_array())
# test for datetime
df = DataFrame()
data = np.array(["2018-10-07", "2018-10-08"], dtype="datetime64")
df["datetimes"] = data
search_date = datetime.datetime.strptime("2018-10-08", "%Y-%m-%d")
expr = "datetimes==@search_date"
got = df.query(expr, local_dict={"search_date": search_date})
np.testing.assert_array_equal(data[1], got["datetimes"].to_array())
def test_query_ref_env(data, fn):
# prepare
nelem, seed = data
expect_fn, query_expr = fn
np.random.seed(seed)
df = DataFrame()
df["a"] = aa = np.arange(nelem)
df["b"] = bb = np.random.random(nelem) * nelem
c = 2.3
d = 1.2
# udt
expect_mask = expect_fn(aa, bb, c, d)
print(expect_mask)
df2 = df.query(query_expr)
# check
assert len(df2) == np.count_nonzero(expect_mask)
np.testing.assert_array_almost_equal(df2["a"].to_array(), aa[expect_mask])
np.testing.assert_array_almost_equal(df2["b"].to_array(), bb[expect_mask])
