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