def test_visit_should_call_drop(self):
stmt = MagicMock(spec=DropTableStatement)
convertor = StatementToPlanConvertor()
mock = MagicMock()
convertor.visit_drop = mock
convertor.visit(stmt)
mock.assert_called_once()
mock.assert_called_with(stmt)
def test_visit_projection_should_add_logical_predicate(
self, mock_lproject):
converter = StatementToPlanConvertor()
projects = MagicMock()
converter._visit_projection(projects)
mock_lproject.assert_called_with(projects)
mock_lproject.return_value.append_child.assert_called()
self.assertEqual(mock_lproject.return_value, converter._plan)
def test_visit_select_predicate_should_add_logical_filter(
self, mock_lfilter):
converter = StatementToPlanConvertor()
select_predicate = MagicMock()
converter._visit_select_predicate(select_predicate)
mock_lfilter.assert_called_with(select_predicate)
mock_lfilter.return_value.append_child.assert_called()
self.assertEqual(mock_lfilter.return_value, converter._plan)
def test_visit_table_ref_should_create_logical_get_opr(self, mock_lget):
converter = StatementToPlanConvertor()
table_ref = MagicMock(spec=TableRef, alias="alias")
table_ref.is_select.return_value = False
table_ref.sample_freq = None
converter.visit_table_ref(table_ref)
mock_lget.assert_called_with(table_ref, table_ref.table.table_obj,
"alias")
self.assertEqual(mock_lget.return_value, converter._plan)
def test_visit_should_call_create_udf(self):
stmt = MagicMock(spec=CreateUDFStatement)
convertor = StatementToPlanConvertor()
mock = MagicMock()
convertor.visit_create_udf = mock
convertor.visit(stmt)
mock.assert_called_once()
mock.assert_called_with(stmt)
def execute_query(query) -> Iterator[Batch]:
"""
Execute the query and return a result generator.
"""
stmt = Parser().parse(query)[0]
try:
StatementBinder(StatementBinderContext()).bind(stmt)
except Exception as error:
raise RuntimeError(f'Binder failed: {error}')
l_plan = StatementToPlanConvertor().visit(stmt)
p_plan = PlanGenerator().build(l_plan)
return PlanExecutor(p_plan).execute_plan()
def test_visit_select_should_not_call_visits_for_null_values(self):
converter = StatementToPlanConvertor()
converter.visit_table_ref = MagicMock()
converter._visit_projection = MagicMock()
converter._visit_select_predicate = MagicMock()
converter._visit_union = MagicMock()
statement = SelectStatement()
converter.visit_select(statement)
converter.visit_table_ref.assert_not_called()
converter._visit_projection.assert_not_called()
converter._visit_select_predicate.assert_not_called()
def test_visit_create_udf(self, mock, l_create_udf_mock):
convertor = StatementToPlanConvertor()
stmt = MagicMock()
stmt.name = "name"
stmt.if_not_exists = True
stmt.inputs = ["inp"]
stmt.outputs = ["out"]
stmt.impl_path = "tmp.py"
stmt.udf_type = "classification"
mock.side_effect = ["inp", "out"]
convertor.visit_create_udf(stmt)
mock.assert_any_call(stmt.inputs, True)
mock.assert_any_call(stmt.outputs, False)
l_create_udf_mock.assert_called_once()
l_create_udf_mock.assert_called_with(
stmt.name,
stmt.if_not_exists,
"inp",
"out",
stmt.impl_path,
stmt.udf_type,
)
def test_visit_select_should_call_appropriate_visit_methods(self):
converter = StatementToPlanConvertor()
converter.visit_table_ref = MagicMock()
converter._visit_projection = MagicMock()
converter._visit_select_predicate = MagicMock()
converter._visit_union = MagicMock()
statement = MagicMock()
statement.from_table = MagicMock(spec=TableRef)
converter.visit_select(statement)
converter.visit_table_ref.assert_called_with(statement.from_table)
converter._visit_projection.assert_called_with(statement.target_list)
converter._visit_select_predicate.assert_called_with(
statement.where_clause)