def test_should_use_the_same_function_if_not_gpu_compatible(self):
mock_function = MagicMock(return_value=pd.DataFrame())
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.assert_called()
def test_simple_function_scan(self):
values = Batch(pd.DataFrame([1, 2, 3], columns=['a']))
expression = FunctionExpression(lambda x: x + 1,
name='test',
alias='test')
expression.output_col_aliases = ['test.a']
plan = type("FunctionScanPlan", (), {"func_expr": expression})
function_scan_executor = FunctionScanExecutor(plan)
actual = list(function_scan_executor.exec(lateral_input=values))[0]
expected = Batch(pd.DataFrame([2, 3, 4], columns=['test.a']))
self.assertEqual(expected, actual)
def test_should_execute_same_function_if_no_gpu(self, context):
context_instance = context.return_value
mock_function = MagicMock(spec=GPUCompatible,
return_value=pd.DataFrame())
context_instance.gpu_device.return_value = NO_GPU
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.assert_called()
def test_function_move_the_device_to_gpu_if_compatible(self, context):
context_instance = context.return_value
mock_function = MagicMock(spec=GPUCompatible)
gpu_mock_function = Mock(return_value=pd.DataFrame())
gpu_device_id = '2'
mock_function.to_device.return_value = gpu_mock_function
context_instance.gpu_device.return_value = gpu_device_id
expression = FunctionExpression(mock_function, name="test")
input_batch = Batch(frames=pd.DataFrame())
expression.evaluate(input_batch)
mock_function.to_device.assert_called_with(gpu_device_id)
gpu_mock_function.assert_called()
def test_lateral_join_with_where(self):
select_query = '''SELECT frame FROM MyVideo JOIN LATERAL
ObjectDet(frame);'''
parser = Parser()
select_stmt = parser.parse(select_query)[0]
tuple_frame = TupleValueExpression('frame')
func_expr = FunctionExpression(func=None,
name='ObjectDet',
children=[tuple_frame])
from_table = TableRef(
JoinNode(TableRef(TableInfo('MyVideo')),
TableRef(func_expr),
join_type=JoinType.LATERAL_JOIN))
expected_stmt = SelectStatement([tuple_frame], from_table)
self.assertEqual(select_stmt, expected_stmt)
def visitUdfFunction(self, ctx: evaql_parser.UdfFunctionContext):
udf_name = None
udf_output = None
if ctx.simpleId():
udf_name = self.visit(ctx.simpleId())
else:
logger.error('UDF function name missing.')
if ctx.dottedId():
udf_output = self.visit(ctx.dottedId())
udf_args = self.visit(ctx.functionArgs())
func_expr = FunctionExpression(None, name=udf_name, output=udf_output)
for arg in udf_args:
func_expr.append_child(arg)
return func_expr
def test_should_return_false_for_unequal_expressions(self):
const_exp1 = ConstantValueExpression(0)
const_exp2 = ConstantValueExpression(1)
func_expr = FunctionExpression(lambda x: x + 1, name='test')
cmpr_exp = ComparisonExpression(ExpressionType.COMPARE_NEQ, const_exp1,
const_exp2)
tuple_expr = TupleValueExpression(col_name='id')
aggr_expr = AggregationExpression(ExpressionType.AGGREGATION_MAX, None,
tuple_expr)
logical_expr = LogicalExpression(ExpressionType.LOGICAL_OR, cmpr_exp,
cmpr_exp)
self.assertNotEqual(const_exp1, const_exp2)
self.assertNotEqual(cmpr_exp, const_exp1)
self.assertNotEqual(func_expr, cmpr_exp)
self.assertNotEqual(tuple_expr, aggr_expr)
self.assertNotEqual(aggr_expr, tuple_expr)
self.assertNotEqual(tuple_expr, cmpr_exp)
self.assertNotEqual(logical_expr, cmpr_exp)