diff --git a/evadb/models/storage/batch.py b/evadb/models/storage/batch.py index 86ec28727..f178add21 100644 --- a/evadb/models/storage/batch.py +++ b/evadb/models/storage/batch.py @@ -173,16 +173,18 @@ def apply_function_expression(self, expr: Callable) -> Batch: """ Execute function expression on frames. """ - - if ( - hasattr(expr, "forward") - and hasattr(expr.forward, "tags") - and (len(expr.forward.tags["input"]) != 0) - ): - input_tags = expr.forward.tags["input"][0] - output_tags = expr.forward.tags["output"][0] - self.drop_column_alias(metadata=(input_tags, output_tags)) - else: + try: + if ( + hasattr(expr, "forward") + and hasattr(expr.forward, "tags") + and (len(expr.forward.tags["input"]) != 0) + ): + input_tags = expr.forward.tags["input"][0] + output_tags = expr.forward.tags["output"][0] + self.drop_column_alias(metadata=(input_tags, output_tags)) + else: + self.drop_column_alias() + except (TypeError, KeyError): self.drop_column_alias() return Batch(expr(self._frames))