def test_blending_2_2(self):
q1 = Select(
columns=[
column('ad_id'),
column('impressions'),
column(HUSKY_QUERY_DATA_SOURCE_COLUMN_NAME)
],
from_obj=table('table1'),
)
df1 = Dataframe(
q1, get_mocked_dataframe_columns_map(['ad_id', 'impressions']),
set(), {'SF'})
q2 = Select(
columns=[
column('ad_id'),
column('campaign_id'),
column('impressions'),
column(HUSKY_QUERY_DATA_SOURCE_COLUMN_NAME),
],
from_obj=table('table2'),
)
df2 = Dataframe(
q2,
get_mocked_dataframe_columns_map(
['ad_id', 'impressions', 'campaign_id']), set(), {'SF'})
blended_df = blend_dataframes(SNOWFLAKE_HUSKY_CONTEXT, [df1, df2])
self.write_test_expectations('query.sql',
compile_query(blended_df.query))
expected_query = self.read_test_expectations('query.sql')
self.assertEqual(expected_query, compile_query(blended_df.query))
self.assertEqual({'ad_id', 'impressions', 'campaign_id'},
set(blended_df.slug_to_column.keys()))
def _build_data_blend_query(
cls,
ctx: HuskyQueryContext,
taxon_manager: BlendingTaxonManager,
config_arg: BlendingDataRequest,
query_info: BlendingQueryInfo,
) -> Dataframe:
"""
Builds subquery for each subrequest and then blends them all into one dataframe.
:param ctx: Husky query context
"""
dataframes = []
request = BlendingDataRequest(config_arg.to_native(
)) # Clone, coz we will be modifying subqueries
for subrequest in request.data_subrequests:
# add comparison taxons to data subrequest taxons
subrequest.taxons = taxon_manager.get_subrequest_taxons(subrequest)
sub_query_info = QueryInfo({
'used_raw_taxons': subrequest.taxons,
})
query_info.subrequests_info.append(sub_query_info)
# Build query for subrequest and add it to the list
data_source = subrequest.properties.data_source
dimension_templates = taxon_manager.plan.data_source_formula_templates[
data_source]
filter_templates = taxon_manager.plan.data_source_filter_templates[
data_source]
df = MainQueryBuilder.build_query(
ctx,
subrequest.to_internal_model(),
sub_query_info,
taxon_manager.used_taxons,
dimension_templates,
filter_templates=filter_templates,
allowed_physical_data_sources=set(
request.physical_data_sources)
if request.physical_data_sources else None,
)
dataframes.append(df)
return blend_dataframes(
ctx, dataframes, taxon_manager.plan.data_source_formula_templates)
def _build_comparison_blend_query(
cls,
ctx: HuskyQueryContext,
config_arg: BlendingDataRequest,
taxon_manager: BlendingTaxonManager,
query_info: BlendingQueryInfo,
allowed_physical_data_sources: Optional[Set[str]] = None,
) -> Optional[Dataframe]:
"""
Builds comparison query for each subrequest and then blends them all into one comparison dataframe.
"""
dataframes = []
config = BlendingDataRequest(config_arg.to_native(
)) # Clone, coz we will be modifying subqueries
assert config.comparison, 'Comparison must be defined when trying to build comparison query..'
comparison: ComparisonConfig = config.comparison
for _subrequest in config.data_subrequests:
subrequest = cls._build_comparison_subrequest(
_subrequest, comparison, taxon_manager)
data_source = subrequest.properties.data_source
# if no comparison taxons were found for this subrequest, skip creating comparison query for it as well
if len(subrequest.taxons) == 0:
continue
bm_sub_query_info = QueryInfo.create(subrequest)
query_info.comparison_subrequests_info.append(bm_sub_query_info)
# Build comparison dataframe and add it to a list.
# TODO pass down TelPlan for comparisons
# ComparisonRequestBuilder might have added filters (typically for company id project id)
# Me create new filter templates for this comparison subrequest.
filter_templates = TelPlanner.get_preaggregation_filter_templates(
ctx,
[
subrequest.preaggregation_filters,
subrequest.scope.preaggregation_filters
],
taxon_manager.taxon_map,
data_source,
)
dataframes.append(
QueryBuilder.build_query(
ctx,
subrequest,
bm_sub_query_info,
taxon_manager.used_taxons,
dimension_templates=taxon_manager.plan.
comparison_data_source_formula_templates[data_source],
filter_templates=filter_templates,
allowed_physical_data_sources=allowed_physical_data_sources,
))
# if no comparison subrequests were created, there is no need to blend data frames
if len(dataframes) == 0:
return None
# Blend all comparison dataframes into one
# TODO pass down TelPlan for comparisons
data_source_formula_templates = taxon_manager.plan.comparison_data_source_formula_templates
dataframe = blend_dataframes(ctx, dataframes,
data_source_formula_templates)
# Prefix all comparison metric columns with 'comparison@' and create comparison taxon for it.
query = dataframe.query
final_columns = []
aliased_taxon_by_slug: Dict[TaxonExpressionStr,
DataframeColumn] = dict()
for slug, df_column in dataframe.slug_to_column.items():
# Alias metrics with comparison@ prefix, and select dimensions..
if df_column.taxon.is_dimension:
new_taxon = df_column.taxon.copy(deep=True)
new_slug = TaxonExpressionStr(f'{slug}')
else:
new_slug, new_taxon = BlendingTaxonManager.create_comparison_taxon(
df_column.taxon)
final_columns.append(query.c[safe_identifier(slug)].label(
new_taxon.slug_safe_sql_identifier))
aliased_taxon_by_slug[new_slug] = DataframeColumn(
new_slug, new_taxon, df_column.quantity_type)
for pre_formulas in data_source_formula_templates.values():
# and also select the dim columns from dim templates.
for pre_formula in pre_formulas:
final_columns.append(
literal_column(
quote_identifier(pre_formula.label, ctx.dialect)))
renamed_cols_query = select(sort_columns(final_columns)).select_from(
dataframe.query)
return Dataframe(renamed_cols_query, aliased_taxon_by_slug,
dataframe.used_model_names,
dataframe.used_physical_data_sources)