def test_formatting() -> None: """ Validates the formatting of the arrayFilter expressions. """ assert tupleElement( "tags_key", arrayJoin( "snuba_all_tags", zip_columns( Column(None, None, "tags.key"), Column(None, None, "tags.value"), ), ), Literal(None, 1), ).accept(ClickhouseExpressionFormatter()) == ( "(tupleElement((arrayJoin(arrayMap((x, y -> tuple(x, y)), " "tags.key, tags.value)) AS snuba_all_tags), 1) AS tags_key)" ) assert tupleElement( "tags_key", arrayJoin( "snuba_all_tags", filter_key_values( zip_columns( Column(None, None, "tags.key"), Column(None, None, "tags.value"), ), [Literal(None, "t1"), Literal(None, "t2")], ), ), Literal(None, 1), ).accept(ClickhouseExpressionFormatter()) == ( "(tupleElement((arrayJoin(arrayFilter((pair -> in(" "tupleElement(pair, 1), tuple('t1', 't2'))), " "arrayMap((x, y -> tuple(x, y)), tags.key, tags.value))) AS snuba_all_tags), 1) AS tags_key)" )
def _filtered_mapping_pairs( alias: Optional[str], column_name: str, pair_alias: str, filtered_tags: Sequence[LiteralExpr], array_index: LiteralExpr, ) -> Expression: # (arrayJoin(arrayFilter( # pair -> tupleElement(pair, 1) IN (tags), # arrayMap((x,y) -> (x,y), tags.key, tags.value) # )) as all_tags).1 return tupleElement( alias, arrayJoin( pair_alias, filter_key_values( zip_columns( ColumnExpr(None, None, key_column(column_name)), ColumnExpr(None, None, val_column(column_name)), ), filtered_tags, ), ), array_index, )
def filtered_mapping_keys(alias: Optional[str], column_name: str, filtered: Sequence[str]) -> Expression: return arrayJoin( alias, filter_column( ColumnExpr(None, None, column_name), [LiteralExpr(None, f) for f in filtered], ), )
def _filtered_mapping_keys(alias: Optional[str], column_name: str, filtered_tags: Sequence[LiteralExpr]) -> Expression: # arrayJoin(arrayFilter( # tag -> tag IN (tags), # tags.key # )) return arrayJoin( alias, filter_keys(ColumnExpr(None, None, key_column(column_name)), filtered_tags), )
def unfiltered_mapping_tuples( alias: Optional[str], tuple_alias: str, tuple_index: LiteralExpr, column_names: Sequence[str], ) -> Expression: return tupleElement( alias, arrayJoin( tuple_alias, zip_columns( *[ColumnExpr(None, None, column) for column in column_names]), ), tuple_index, )
def _unfiltered_mapping_pairs(alias: Optional[str], column_name: str, pair_alias: str, tuple_index: LiteralExpr) -> Expression: # (arrayJoin( # arrayMap((x,y) -> (x,y), tags.key, tags.value) # as all_tags).1 return tupleElement( alias, arrayJoin( pair_alias, zip_columns( ColumnExpr(None, None, key_column(column_name)), ColumnExpr(None, None, val_column(column_name)), ), ), tuple_index, )
def filtered_mapping_tuples( alias: Optional[str], tuple_alias: str, tuple_index: LiteralExpr, column_names: Sequence[str], single_filtered: Dict[LiteralExpr, Sequence[str]], multiple_filtered: Dict[Tuple[LiteralExpr, ...], Sequence[Tuple[str, ...]]], ) -> Expression: return tupleElement( alias, arrayJoin( tuple_alias, filter_expression( zip_columns(*[ ColumnExpr(None, None, column) for column in column_names ]), single_filtered, multiple_filtered, ), ), tuple_index, )
def array_join_col(ops=None, groups=None, op_groups=None): conditions: List[Expression] = [] argument_name = "arg" argument = Argument(None, argument_name) if ops: conditions.append( binary_condition( ConditionFunctions.IN, tupleElement(None, argument, Literal(None, 1)), FunctionCall(None, "tuple", tuple(Literal(None, op) for op in ops)), )) if groups: conditions.append( binary_condition( ConditionFunctions.IN, tupleElement(None, argument, Literal(None, 2)), FunctionCall(None, "tuple", tuple(Literal(None, group) for group in groups)), )) if op_groups: conditions.append( binary_condition( ConditionFunctions.IN, FunctionCall( None, "tuple", ( tupleElement(None, argument, Literal(None, 1)), tupleElement(None, argument, Literal(None, 2)), ), ), FunctionCall( None, "tuple", tuple( FunctionCall(None, "tuple", (Literal(None, op), Literal(None, group))) for op, group in op_groups), ), )) cols = FunctionCall( None, "arrayMap", ( Lambda( None, ("x", "y", "z"), FunctionCall( None, "tuple", tuple(Argument(None, arg) for arg in ("x", "y", "z"))), ), Column(None, None, "spans.op"), Column(None, None, "spans.group"), Column(None, None, "spans.exclusive_time"), ), ) if conditions: cols = FunctionCall( None, "arrayFilter", ( Lambda(None, (argument_name, ), combine_and_conditions(conditions)), cols, ), ) return arrayJoin("snuba_all_spans", cols)
pytest.param( { "aggregations": [], "groupby": [], "selected_columns": ["col1"], "conditions": [["tags_key", "IN", ["t1", "t2"]]], }, ClickhouseQuery( None, selected_columns=[ SelectedExpression(name="col1", expression=Column("_snuba_col1", None, "col1")) ], condition=in_condition( arrayJoin("_snuba_tags_key", Column(None, None, "tags.key")), [Literal(None, "t1"), Literal(None, "t2")], ), ), id="no tag in select clause", ), # Individual tag, no change pytest.param( { "aggregations": [], "groupby": [], "selected_columns": ["tags_key", "tags_value"], "conditions": [["col", "IN", ["t1", "t2"]]], }, ClickhouseQuery( None, selected_columns=[