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