def test_build_execution_tree_should_create_correct_exec_node(self):
# SequentialScanExecutor
plan = SeqScanPlan(MagicMock(), [])
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, SequentialScanExecutor)
# PPExecutor
plan = PPScanPlan(MagicMock())
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, PPExecutor)
# CreateExecutor
plan = CreatePlan(MagicMock(), [], False)
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, CreateExecutor)
# InsertExecutor
plan = InsertPlan(0, [], [])
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, InsertExecutor)
# CreateUDFExecutor
plan = CreateUDFPlan('test', False, [], [], MagicMock(), None)
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, CreateUDFExecutor)
# LoadDataExecutor
plan = LoadDataPlan(MagicMock(), MagicMock())
executor = PlanExecutor(plan)._build_execution_tree(plan)
self.assertIsInstance(executor, LoadDataExecutor)
def test_execute_plan_for_create_insert_load_plans(self, mock_batch,
mock_clean, mock_build):
mock_batch.return_value = []
# CreateExecutor
tree = MagicMock(node=CreatePlan(None, [], False))
mock_build.return_value = tree
actual = PlanExecutor(None).execute_plan()
tree.exec.assert_called_once()
mock_batch.assert_called_once()
assert_frame_equal(mock_batch.call_args[0][0], pd.DataFrame())
mock_build.assert_called_once_with(None)
mock_clean.assert_called_once()
self.assertEqual(actual, [])
# InsertExecutor
mock_batch.reset_mock()
mock_build.reset_mock()
mock_clean.reset_mock()
tree = MagicMock(node=InsertPlan(0, [], []))
mock_build.return_value = tree
actual = PlanExecutor(None).execute_plan()
tree.exec.assert_called_once()
mock_batch.assert_called_once()
assert_frame_equal(mock_batch.call_args[0][0], pd.DataFrame())
mock_build.assert_called_once_with(None)
mock_clean.assert_called_once()
self.assertEqual(actual, [])
# CreateUDFExecutor
mock_batch.reset_mock()
mock_build.reset_mock()
mock_clean.reset_mock()
tree = MagicMock(node=CreateUDFPlan(None, False, [], [], None))
mock_build.return_value = tree
actual = PlanExecutor(None).execute_plan()
tree.exec.assert_called_once()
mock_batch.assert_called_once()
assert_frame_equal(mock_batch.call_args[0][0], pd.DataFrame())
mock_build.assert_called_once_with(None)
mock_clean.assert_called_once()
self.assertEqual(actual, [])
# LoadDataExecutor
mock_batch.reset_mock()
mock_build.reset_mock()
mock_clean.reset_mock()
tree = MagicMock(node=LoadDataPlan(None, None))
mock_build.return_value = tree
actual = PlanExecutor(None).execute_plan()
tree.exec.assert_called_once()
mock_batch.assert_called_once()
assert_frame_equal(mock_batch.call_args[0][0], pd.DataFrame())
mock_build.assert_called_once_with(None)
mock_clean.assert_called_once()
self.assertEqual(actual, [])
def test_load_data_plan(self):
table_metainfo = 'meta_info'
file_path = 'test.mp4'
plan_str = 'LoadDataPlan(table_id={},file_path={})'.format(
table_metainfo, file_path)
plan = LoadDataPlan(table_metainfo, file_path)
self.assertEqual(plan.node_type, PlanNodeType.LOAD_DATA)
self.assertEqual(plan.table_metainfo, table_metainfo)
self.assertEqual(plan.file_path, file_path)
self.assertEqual(str(plan), plan_str)
def build(self, operator: Operator):
self.__init__()
self._visit(operator)
load_plan = LoadDataPlan(self._table_metainfo, self._path)
return load_plan
def apply(self, before: LogicalLoadData, context: OptimizerContext):
after = LoadDataPlan(before.table_metainfo, before.path)
return after