def column_expr(dataset,
column_name,
query: Query,
parsing_context: ParsingContext,
alias=None,
aggregate=None):
"""
Certain special column names expand into more complex expressions. Return
a 2-tuple of:
(expanded column expression, sanitized alias)
Needs the body of the request for some extra data used to expand column expressions.
"""
assert column_name or aggregate
assert not aggregate or (aggregate and (column_name or alias))
column_name = column_name or ''
if is_function(column_name, 0):
return complex_column_expr(dataset, column_name, query,
parsing_context)
elif isinstance(column_name, (list, tuple)) and aggregate:
return complex_column_expr(dataset, [aggregate, column_name, alias],
query, parsing_context)
elif isinstance(column_name, str) and QUOTED_LITERAL_RE.match(column_name):
return escape_literal(column_name[1:-1])
else:
expr = dataset.column_expr(column_name, query, parsing_context)
if aggregate:
expr = function_expr(aggregate, expr)
alias = escape_alias(alias or column_name)
return alias_expr(expr, alias, parsing_context)
def __tags_expr(self,
column_name: str,
query: Query,
parsing_context: ParsingContext,
table_alias: str="",
) -> str:
"""
Return an expression that array-joins on tags to produce an output with one
row per tag.
"""
assert column_name in ['tags_key', 'tags_value']
col, k_or_v = column_name.split('_', 1)
nested_tags_only = state.get_config('nested_tags_only', 1)
qualified_col = qualified_column(col, table_alias)
# Generate parallel lists of keys and values to arrayJoin on
if nested_tags_only:
key_list = '{}.key'.format(qualified_col)
val_list = '{}.value'.format(qualified_col)
else:
promoted = self.__promoted_columns[col]
col_map = self.__column_tag_map[col]
key_list = u'arrayConcat([{}], {}.key)'.format(
u', '.join(u'\'{}\''.format(col_map.get(p, p)) for p in promoted),
qualified_col
)
val_list = u'arrayConcat([{}], {}.value)'.format(
', '.join(self.__string_col(p) for p in promoted),
qualified_col
)
qualified_key = qualified_column("tags_key", table_alias)
qualified_value = qualified_column("tags_value", table_alias)
cols_used = query.get_all_referenced_columns() & set([qualified_key, qualified_value])
if len(cols_used) == 2:
# If we use both tags_key and tags_value in this query, arrayjoin
# on (key, value) tag tuples.
expr = (u'arrayJoin(arrayMap((x,y) -> [x,y], {}, {}))').format(
key_list,
val_list
)
# put the all_tags expression in the alias cache so we can use the alias
# to refer to it next time (eg. 'all_tags[1] AS tags_key'). instead of
# expanding the whole tags expression again.
expr = alias_expr(expr, 'all_tags', parsing_context)
return u'({})[{}]'.format(expr, 1 if k_or_v == 'key' else 2)
else:
# If we are only ever going to use one of tags_key or tags_value, don't
# bother creating the k/v tuples to arrayJoin on, or the all_tags alias
# to re-use as we won't need it.
return 'arrayJoin({})'.format(key_list if k_or_v == 'key' else val_list)
def complex_column_expr(dataset,
expr,
query: Query,
parsing_context: ParsingContext,
depth=0):
function_tuple = is_function(expr, depth)
if function_tuple is None:
raise ValueError(
'complex_column_expr was given an expr %s that is not a function at depth %d.'
% (expr, depth))
name, args, alias = function_tuple
out = []
i = 0
while i < len(args):
next_2 = args[i:i + 2]
if is_function(next_2, depth + 1):
out.append(
complex_column_expr(dataset, next_2, query, parsing_context,
depth + 1))
i += 2
else:
nxt = args[i]
if is_function(nxt, depth + 1): # Embedded function
out.append(
complex_column_expr(dataset, nxt, query, parsing_context,
depth + 1))
elif isinstance(nxt, str):
out.append(column_expr(dataset, nxt, query, parsing_context))
else:
out.append(escape_literal(nxt))
i += 1
ret = function_expr(name, ', '.join(out))
if alias:
ret = alias_expr(ret, alias, parsing_context)
return ret
def __tags_expr(
self,
parsed_col: ParsedNestedColumn,
query: Query,
parsing_context: ParsingContext,
table_alias: str = "",
) -> str:
"""
Return an expression that array-joins on tags to produce an output with one
row per tag.
It can also apply an arrayFilter in the arrayJoin if an equivalent condition
is found in the query in order to reduce the size of the arrayJoin.
"""
col, k_or_v = parsed_col.col_name.split("_", 1)
nested_tags_only = state.get_config("nested_tags_only", 1)
qualified_col = qualified_column(col, table_alias)
# Generate parallel lists of keys and values to arrayJoin on
if nested_tags_only:
key_list = "{}.key".format(qualified_col)
val_list = "{}.value".format(qualified_col)
else:
promoted = self.__promoted_columns[col]
col_map = self.__column_tag_map[col]
key_list = "arrayConcat([{}], {}.key)".format(
", ".join("'{}'".format(col_map.get(p, p)) for p in promoted),
qualified_col,
)
val_list = "arrayConcat([{}], {}.value)".format(
", ".join(self.__string_col(p) for p in promoted),
qualified_col)
qualified_key = qualified_column("tags_key", table_alias)
qualified_value = qualified_column("tags_value", table_alias)
cols_used = query.get_all_referenced_columns() & set(
[qualified_key, qualified_value])
filter_tags = ",".join(
[f"'{tag}'" for tag in self.__get_filter_tags(query)])
if len(cols_used) == 2:
# If we use both tags_key and tags_value in this query, arrayjoin
# on (key, value) tag tuples.
mapping = f"arrayMap((x,y) -> [x,y], {key_list}, {val_list})"
if filter_tags:
filtering = (
f"arrayFilter(pair -> pair[1] IN ({filter_tags}), {mapping})"
)
else:
filtering = mapping
expr = f"arrayJoin({filtering})"
# put the all_tags expression in the alias cache so we can use the alias
# to refer to it next time (eg. 'all_tags[1] AS tags_key'). instead of
# expanding the whole tags expression again.
expr = alias_expr(expr, "all_tags", parsing_context)
return "({})[{}]".format(expr, 1 if k_or_v == "key" else 2)
else:
# If we are only ever going to use one of tags_key or tags_value, don't
# bother creating the k/v tuples to arrayJoin on, or the all_tags alias
# to re-use as we won't need it.
if filter_tags:
return (
f"arrayJoin(arrayFilter(tag -> tag IN ({filter_tags}), {key_list}))"
)
else:
return f"arrayJoin({key_list if k_or_v == 'key' else val_list})"
def output_builder(alias: Optional[str], name: str, params: List[Any]) -> Any:
ret = function_expr(name, ", ".join(params))
if alias:
ret = alias_expr(ret, alias, parsing_context)
return ret