def test_accessors() -> None:
func = FunctionCall(
String("f_name"),
(
FunctionCall(String("f"), (Column(None, String("my_col")), )),
Param(
"second_function",
FunctionCall(Param("second_function_name", Any(str)), None),
),
),
)
result = func.match(
FunctionCallExpr(
"irrelevant",
"f_name",
(
FunctionCallExpr(None, "f",
(ColumnExpr(None, None, "my_col"), )),
FunctionCallExpr(None, "second_name", tuple()),
),
))
assert result is not None
assert result.expression("second_function") == FunctionCallExpr(
None, "second_name", tuple())
assert result.scalar("second_function_name") == "second_name"
def filter_expression(
columns: Expression,
single_filtered: Dict[LiteralExpr, Sequence[str]],
multiple_filtered: Dict[Tuple[LiteralExpr, ...], Sequence[Tuple[str,
...]]],
) -> Expression:
argument_name = "arg"
argument = Argument(None, argument_name)
conditions: List[Expression] = []
for index in single_filtered:
conditions.append(
binary_condition(
ConditionFunctions.IN,
tupleElement(None, argument, index),
FunctionCallExpr(
None,
"tuple",
tuple(
LiteralExpr(None, f) for f in single_filtered[index]),
),
))
for indices in multiple_filtered:
conditions.append(
binary_condition(
ConditionFunctions.IN,
FunctionCallExpr(
None,
"tuple",
tuple(
tupleElement(None, argument, index)
for index in indices),
),
FunctionCallExpr(
None,
"tuple",
tuple(
FunctionCallExpr(
None,
"tuple",
tuple(LiteralExpr(None, t) for t in tuples),
) for tuples in multiple_filtered[indices]),
),
))
return FunctionCallExpr(
None,
"arrayFilter",
(Lambda(None, (argument_name, ),
combine_and_conditions(conditions)), columns),
)
def filter_key_values(key_values: Expression,
keys: Sequence[LiteralExpr]) -> Expression:
"""
Filter an array of key value pairs based on a sequence of keys
(tag keys in this case).
"""
return FunctionCallExpr(
None,
"arrayFilter",
(
Lambda(
None,
("pair", ),
in_condition(
# A pair here is a tuple with two elements (key
# and value) and the index of the first element in
# Clickhouse is 1 instead of 0.
tupleElement(
None,
Argument(None, "pair"),
LiteralExpr(None, 1),
),
keys,
),
),
key_values,
),
)
def process_column(exp: Expression) -> Expression:
match = matcher.match(exp)
if match:
inner = replace(exp, alias=None)
return FunctionCallExpr(
exp.alias,
"if",
(
binary_condition(
ConditionFunctions.IN,
inner,
literals_tuple(
None,
[
LiteralExpr(None, "1"),
LiteralExpr(None, "True")
],
),
),
LiteralExpr(None, "True"),
LiteralExpr(None, "False"),
),
)
return exp
def filter_column(column: Expression,
keys: Sequence[LiteralExpr]) -> Expression:
return FunctionCallExpr(
None,
"arrayFilter",
(Lambda(None,
("x", ), in_condition(Argument(None, "x"), keys)), column),
)
def zip_columns(column1: ColumnExpr, column2: ColumnExpr) -> Expression:
"""
Turns two array columns into an array of pairs
"""
return FunctionCallExpr(
None,
"arrayMap",
(
Lambda(
None,
("x", "y"),
FunctionCallExpr(
None, "tuple", (Argument(None, "x"), Argument(None, "y"),),
),
),
column1,
column2,
),
)
def build_nullable_mapping_expr(
alias: Optional[str],
table_name: Optional[str],
col_name: str,
mapping_key: Expression,
) -> FunctionCallExpr:
# TODO: Add a pattern for this expression if we need it.
return FunctionCallExpr(
alias,
"if",
(
FunctionCallExpr(
None,
"has",
(ColumnExpr(None, table_name, f"{col_name}.key"), mapping_key),
),
build_mapping_expr(None, table_name, col_name, mapping_key),
LiteralExpr(None, None),
),
)
def zip_columns(*columns: ColumnExpr) -> Expression:
if len(columns) not in {2, 3}:
raise NotImplementedError("Can only zip between 2 and 3 columns.")
arguments = ("x", "y", "z")[:len(columns)]
return FunctionCallExpr(
None,
"arrayMap",
(
Lambda(
None,
arguments,
FunctionCallExpr(
None, "tuple",
tuple(Argument(None, arg) for arg in arguments)),
),
*columns,
),
)
def filter_keys(column: Expression, keys: Sequence[LiteralExpr]) -> Expression:
"""
Filter a Column array based on a sequence of keys.
"""
return FunctionCallExpr(
None,
"arrayFilter",
(
Lambda(None, ("tag",), in_condition(None, Argument(None, "tag"), keys),),
column,
),
)
def build_mapping_expr(
alias: Optional[str],
table_name: Optional[str],
col_name: str,
mapping_key: Expression,
) -> FunctionCallExpr:
return arrayElement(
alias,
ColumnExpr(None, table_name, f"{col_name}.value"),
FunctionCallExpr(
None,
"indexOf",
(ColumnExpr(None, table_name, f"{col_name}.key"), mapping_key),
),
)
def attempt_map(
self,
expression: FunctionCallExpr,
children_translator: SnubaClickhouseStrictTranslator,
) -> Optional[FunctionCallExpr]:
if expression.function_name != self.from_name:
return None
return FunctionCallExpr(
alias=expression.alias,
function_name=self.to_name,
parameters=tuple(
exp.accept(children_translator)
for exp in expression.parameters),
)
def replace_exp(exp: Expression) -> Expression:
if matcher.match(exp) is not None:
inner = replace(exp, alias=None)
return FunctionCallExpr(
exp.alias,
"multiIf",
(
binary_condition(ConditionFunctions.EQ, inner,
Literal(None, "")),
Literal(None, ""),
binary_condition(
ConditionFunctions.IN,
inner,
literals_tuple(
None,
[Literal(None, "1"),
Literal(None, "True")]),
),
Literal(None, "True"),
Literal(None, "False"),
),
)
return exp
def generate_bloom_filter_condition(
column_name: str,
single_filtered: Dict[str, Sequence[str]],
multiple_filtered: Dict[Tuple[str, ...], Sequence[Tuple[str, ...]]],
) -> Optional[Expression]:
"""
Generate the filters on the array columns to use the bloom filter index on
the spans.op and spans.group columns in order to filter the transactions
prior to the array join.
The bloom filter index is requires the use of the has function, therefore
the final condition is built up from a series of has conditions.
"""
per_key_vals: Dict[str, Set[str]] = defaultdict(set)
for key, single_filter in single_filtered.items():
for val in single_filter:
per_key_vals[key].add(val)
for keys, multiple_filter in multiple_filtered.items():
for val_tuple in multiple_filter:
for key, val in zip(keys, val_tuple):
per_key_vals[key].add(val)
conditions = [
combine_or_conditions([
FunctionCallExpr(
None,
"has",
(ColumnExpr(None, None, key), LiteralExpr(None, val)),
) for val in sorted(vals)
]) for key, vals in per_key_vals.items()
]
return combine_and_conditions(conditions) if conditions else None
def _produce_output(self, expression: ColumnExpr) -> FunctionCallExpr:
return FunctionCallExpr(
alias=expression.alias,
function_name=self.to_function_name,
parameters=self.to_function_params,
)
Column(Param("p_table_name", Any(type(None))), None),
ColumnExpr("irrelevant", "not_none", "irrelevant"),
None,
),
(
"Match any expression of Column type",
Any(ColumnExpr),
ColumnExpr("irrelevant", "irrelevant", "irrelevant"),
MatchResult(),
),
(
"Match any expression of Column type within function",
FunctionCall(None, (Param("p1", Any(ColumnExpr)), )),
FunctionCallExpr(
"irrelevant",
"irrelevant",
(ColumnExpr("relevant", "relevant", "relevant"), ),
),
MatchResult({"p1": ColumnExpr("relevant", "relevant", "relevant")}),
),
(
"Wrong number of parameters, does not match",
FunctionCall(None, (Param("p1", Any(ColumnExpr)), )),
FunctionCallExpr(
"irrelevant",
"irrelevant",
(
ColumnExpr("relevant", "relevant", "relevant"),
ColumnExpr("relevant", "relevant", "relevant"),
),
),