def test_nested_logical_conditions(sample_df):
step = FilterStep(
name='filter',
condition={
'and': [
{
'or': [
{
'column': 'colA',
'operator': 'eq',
'value': 'toto',
},
{
'column': 'colC',
'operator': 'lt',
'value': 33,
},
]
},
{'column': 'colB', 'operator': 'gt', 'value': 2},
]
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, DataFrame({'colA': ['tata'], 'colB': [3], 'colC': [25]}))
def test_translate_filter(mocker):
step = FilterStep(name='filter',
condition=ComparisonCondition(column='amount',
operator='eq',
value=10))
query = SQLQuery(
query_name='SELECT_STEP_0',
transformed_query='WITH SELECT_STEP_0 AS (SELECT * FROM products)',
selection_query='SELECT * FROM SELECT_STEP_0',
metadata_manager=SqlQueryMetadataManager(
tables_metadata={'table1': {
'toto': 'str',
'tata': 'int'
}}),
)
mocker.patch(
'weaverbird.backends.sql_translator.steps.utils.query_transformation.apply_condition',
return_value='SELECT * FROM SELECT_STEP_0 WHERE amount = 10',
)
res = translate_filter(step, query, index=1)
assert (
res.transformed_query ==
'WITH SELECT_STEP_0 AS (SELECT * FROM products), FILTER_STEP_1 AS (SELECT * FROM SELECT_STEP_0 WHERE amount = 10)'
)
assert res.selection_query == 'SELECT toto, tata FROM FILTER_STEP_1'
def test_simple_notnull_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colA',
'operator': 'notnull',
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, sample_df)
def test_simple_null_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colA',
'operator': 'isnull',
},
)
df_result = execute_filter(step, sample_df)
assert df_result.empty
def test_simple_lt_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colB',
'operator': 'lt',
'value': 2,
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, DataFrame({'colA': ['toto'], 'colB': [1], 'colC': [100]}))
def test_simple_eq_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colA',
'operator': 'eq',
'value': 'tutu',
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, DataFrame({'colA': ['tutu'], 'colB': [2], 'colC': [50]}))
def test_simple_matches_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colA',
'operator': 'matches',
'value': 'a.a',
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, DataFrame({'colA': ['tata'], 'colB': [3], 'colC': [25]}))
def test_simple_ne_filter(sample_df):
step = FilterStep(
name='filter',
condition={
'column': 'colA',
'operator': 'ne',
'value': 'tutu',
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(
df_result, DataFrame({'colA': ['toto', 'tata'], 'colB': [1, 3], 'colC': [100, 25]})
)
def test_translate_filter_error(mocker):
step = FilterStep(name='filter',
condition=ComparisonCondition(column='amount',
operator='eq',
value=10))
query = SQLQuery(
query_name='SELECT_STEP_0',
transformed_query=
'WITH SELECT_STEP_0 AS (SELECT * FROM products), SELECT * FROM SELECT_STEP_0',
selection_query='SELECT * FROM SELECT_STEP_0',
)
mocker.patch(
'weaverbird.backends.sql_translator.steps.filter.apply_condition',
side_effect=NotImplementedError,
)
with pytest.raises(NotImplementedError):
translate_filter(step, query, index=1)
def test_and_logical_conditions(sample_df):
step = FilterStep(
name='filter',
condition={
'and': [
{
'column': 'colB',
'operator': 'le',
'value': 2,
},
{
'column': 'colC',
'operator': 'gt',
'value': 75,
},
]
},
)
df_result = execute_filter(step, sample_df)
assert_dataframes_equals(df_result, DataFrame({'colA': ['toto'], 'colB': [1], 'colC': [100]}))
def test_benchmark_filter(benchmark):
big_df = DataFrame({'value': list(range(1000))})
step = FilterStep(name='filter', condition={'column': 'value', 'operator': 'lt', 'value': 20})
result = benchmark(execute_filter, step, big_df)
assert len(result) == 20
