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 build(self, operator: Operator):
self.__init__()
self._visit(operator)
create_udf_plan = CreateUDFPlan(self._name, self._if_not_exists,
self._inputs, self._outputs,
self._impl_path, self._udf_type)
return create_udf_plan
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_create_udf_plan(self):
udf_name = 'udf'
if_not_exists = True
udfIO = 'udfIO'
inputs = [udfIO, udfIO]
outputs = [udfIO]
impl_path = 'test'
ty = 'classification'
node = CreateUDFPlan(udf_name, if_not_exists, inputs, outputs,
impl_path, ty)
self.assertEqual(node.node_type, PlanNodeType.CREATE_UDF)
self.assertEqual(node.if_not_exists, True)
self.assertEqual(node.inputs, [udfIO, udfIO])
self.assertEqual(node.outputs, [udfIO])
self.assertEqual(node.impl_path, impl_path)
self.assertEqual(node.udf_type, ty)
def apply(self, before: LogicalCreateUDF, context: OptimizerContext):
after = CreateUDFPlan(before.name, before.if_not_exists, before.inputs,
before.outputs, before.impl_path,
before.udf_type)
return after