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(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 # udt expect_mask = expect_fn(aa, bb) 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])