Improve INCREMENTAL_BY_UNIQUE_KEY merge for MSSQL

[codykonior]

For MSSQL/T-SQL the MERGE statement generated for a INCREMENTAL_BY_UNIQUE_KEY model could use improvement.

It typically does:

MERGE INTO ... [__MERGE_TARGET__] 
USING (...)  AS [__MERGE_SOURCE__] 
ON Target Unique Key Columns = Source Unique Key Columns
WHEN MATCHED 
THEN UPDATE SET Each Target Column = Each Source Column

This is bad for MSSQL, updating a row to all of its existing values isn't cancelled out as a no-operation. It overwrites every single column and is all written to the transaction log. So we'd like to prevent row updates where it's not necessary, such as when the row has not changed.

So for MSSQL it's much better to do:

MERGE INTO ... [__MERGE_TARGET__] 
USING (...)  AS [__MERGE_SOURCE__] 
ON Target Unique Key Columns = Source Unique Key Columns
WHEN MATCHED AND EXISTS (SELECT Each Target Column EXCEPT SELECT Each Source Column) (except for the unique columns)
THEN UPDATE SET Each Target Column = Each Source Column (except for the unique columns)

The addition of EXISTS prevents an update from happening on a row where all of the other columns are the same. The syntax for this with SELECT/EXCEPT/SELECT is valid for MSSQL because it honours NULL comparisons, and in this way it's much better than trying to do WHEN MATCHED AND (ISNULL(a1, '') <> ISNULL(a2, '') OR b1 <> b2 OR c1 <> c2) and all of that kind of stuff. Note that this may not work the same way with other database engines, I don't know.

The second part of that change is being a bit more careful with the UPDATE SET columns to not update the unique columns, after all we already checked them as part of the match. But there's a specific reason for this too, because in MSSQL on a normal clustered table, trying to set a key column to itself will cause a deferred update where the engine internally deletes and inserts the entire row. By not setting them, it can update the row in place instead (assuming all the data types fit, etc). I'm not 100% sure this occurs on the heap tables sqlmesh is using, but it should be avoided anyway.

Together these will result in much reduced CPU and IO.

A third thing but slightly less important and that could be done separately 😓 is another keyword so that when you materialize the table it creates the table with a clustered primary key using the unique_key field. I know that we can create indexes using a post-statement, I don't think we can create them on the model table as part of the pre-statement), but for MSSQL converting the fully populated heap at the end to a clustered table is a huge amount of CPU/IO versus creating it that way up front. MSSQL doesn't perform well with heaps, or MERGE on heaps, and so for me every single model will end up as a clustered table if that was an option.

For the user-presented overall implementation I'd like to suggest allowing a keyword on a model like:

MODEL (
    name [xxx].[yyy],
    kind INCREMENTAL_BY_UNIQUE_KEY (
        unique_key [My Id],
        use_merge_exists True,
        use_merge_exists_except [Optionally, Some Column I Want To Exclude],
        type 'clustered'
    ),
    cron '*/5 * * * *'
);

Having use_merge_exists would generate the syntax as above and following the above rules, and use_merge_exists_except would allow you to exclude columns from the EXISTS comparison section.

This is to allow a pattern where you include a [Load Date] = SYSDATETIMEOFFSET() or similar column in your model. This is something that is useful because it would get inserted into the table when a row suddenly appears for the first time, or when the row was actually changed, but shouldn't trigger a change just in itself because it would otherwise force an update every single time, which we don't want.

The type could be used to specify clustered instead of heap tables. When sqlglot materializes the table, it can chuck on a create table ... ( ... constraint PK_modelname primary key clustered (unique_key)) with (data_compression = page). That would be a huge help.

I'm not great with Python but did jump in to try to give this all a shot, but got nowhere fast, as it requires a lot of internal and SQLglot usage knowledge that I just don't have.

[codykonior]

I've worked out how to do a little bit of it in source. It's in base.py in the merge function, ~line 1869.

def merge(
        self,
        target_table: TableName,
        source_table: QueryOrDF,
        columns_to_types: t.Optional[t.Dict[str, exp.DataType]],
        unique_key: t.Sequence[exp.Expression],
        use_match_exists: t.Optional[bool] = True,
        use_match_exists_except: t.Optional[t.Sequence[exp.Expression]] = [exp.column("DateLoaded")],
        when_matched: t.Optional[exp.Whens] = None,
        merge_filter: t.Optional[exp.Expression] = None,
    ) -> None:
        source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
            source_table, columns_to_types, target_table=target_table
        )
        columns_to_types = columns_to_types or self.columns(target_table)
        on = exp.and_(
            *(
                add_table(part, MERGE_TARGET_ALIAS).eq(add_table(part, MERGE_SOURCE_ALIAS))
                for part in unique_key
            )
        )
        if merge_filter:
            on = exp.and_(merge_filter, on)

        if not when_matched:
            match_condition = None
            if use_match_exists and unique_key_names:
                unique_key_names = [y.name for y in unique_key]
                if use_match_exists_except:
                    unique_key_names += [y.name for y in use_match_exists_except]
                target_columns = [exp.column(c, MERGE_TARGET_ALIAS) for c in columns_to_types if c not in unique_key_names]
                source_columns = [exp.column(c, MERGE_SOURCE_ALIAS) for c in columns_to_types if c not in unique_key_names]
                match_condition = exp.Exists(this = exp.select(*target_columns).except_(exp.select(*source_columns)))

            match_expressions = [
                exp.When(
                    matched=True,
                    source=False,
                    condition=match_condition,
                    then=exp.Update(
                        expressions=[
                            exp.column(col, MERGE_TARGET_ALIAS).eq(
                                exp.column(col, MERGE_SOURCE_ALIAS)
                            )
                            for col in columns_to_types if not use_match_exists or col not in unique_key_names
                        ],
                    ),
                )
            ]
        else:
            match_expressions = when_matched.copy().expressions

This achieves the first and second parts of what I wanted; adding WHEN MATCHED AND EXISTS ... and not updating the unique keys, and allowing a hardcoded column name to exclude.

What I couldn't work out how to do was add a keyword to the model definition and pass it down to this function's input.

[izeigerman]

Hey @codykonior ! Thanks for your submission. Overall it makes a lot of sense to me.

I only have concerns around custom keywords. I'm hesitant to introduce engine-specific properties as top-level attributes of the model kind which is shared across all supported engines.

Having use_merge_exists would generate the syntax as above and following the above rules, and use_merge_exists_except would allow you to exclude columns from the EXISTS comparison section.

I'd rather not make this behavior optional and instead opt into the best default behavior for MSSQL. If WHEN MATCHED AND EXISTS is the most optimal way and it doesn't come with any tradeoffs, I don't see any reason to make it a conditional behavior.

@erindru @treysp any thoughts?

[erindru]

I'd rather not make this behavior optional and instead opt into the best default behavior for MSSQL

Yes, that makes sense. We can improve the MERGE implementation for MSSQL as described here and use the default implementation for the other engines.

I guess if the user specifies a custom when_matched then all bets are off and we just use the expression they supplied?

@codykonior , regarding:

A third thing but slightly less important and that could be done separately is another keyword so that when you materialize the table it creates the table with a clustered primary key using the unique_key field

I think this is a separate issue. As far as I am aware, we dont handle indexes automatically because SQLMesh is designed for data warehouses that tend to be column stores so dont have indexes in the same way as a row store like MSSQL does.

However, wouldn't this be a case for a pre statement? SQLMesh creates empty tables before loading them with data, so something like:

@if(
   @runtime_stage = 'creating',
   create clustered index idx_pk on @{this_model} (col1, col2)
);

I cant remember if MSSQL index names are globally unique like they are in other databases. If they are, then you might need to use @resolve_template macro to generate a unique index name

[codykonior]

If this becomes the default behaviour for SQL Server/Azure SQL Database then I don't know of any issues. I just didn't want to assume too much.

Regarding use_match_exists_except vs when_matched I will try to explain my use-case a little further.

If you have a hundred models, each with a hundred columns, I think the source of truth about "what columns are output" should be the query itself.

If you attempted to write your own when_matched statement that covers almost all of them, that's a few dozen dense rows of text (or few screens full of text one column at a time), and it becomes impossible to compare that versus the query to confirm you are actually checking all of the columns you intended to. Over time, as the models change and columns are added and removed, it becomes extremely difficult to trust that the when_matched columns have also been similarly maintained.

I know because I see the code of manual merge statements a lot, thousands of lines long with columns and comparisons, and it drives me crazy! 😸

So that's why I wanted use_match_exists_except. Auto-all of the columns, except something. This way I don't need to spend energy adding, maintaining, or worrying about, my column lists in the merge comparison. And this is specifically for that [Date Loaded] field which helps track downstream when specific rows were inserted or updated to the materialized table.

Alternatives:

• If you had some way to control scripting out the columns I want into the when_exists clause of the model definition? I know you have these macros and stuff but I wasn't sure if they could be used for this, or how.
• People have to choose between [Date Loaded] or efficient merges 😿 That would really suck.
• I've thought about NOT using [Date Loaded] but instead using the pre-statement to turn the output table into a system versioned temporal table, and I can then use its date columns to track the date loaded and deletes. But I was hesitant to go that far, in case it has side effects with other sqlmesh functionality. For example it won't be able to drop or destroy those tables easily... I would have to provide some code to disconnect the temporal table and drop it also. Still, if that's your preference, I can try it...

Ok I will try out the pre-statement for the indexing, thanks.

[codykonior]

I was not able to get the idea you provided to work because @{this_model} ends up double-quoted and so the command fails. e.g. ... [[Database]].[Schema]].[Table]]. Those mismatched parenthesis aren't a typo, that's exactly what it tried to use!

@if(
   @runtime_stage = 'creating',
   create clustered index idx_pk on @{this_model} (col1)
);
-- CREATE CLUSTERED INDEX [idx_pk] ON [[my_db_name]].[my_schema_name]].[xxx__xxxx__558102588]] /* my_db_name.my_schema_name.xxx */]([col1]);

I saw you have @resolve_template to extract just portions of the model name, and I thought that might work, but in this macro it gets sent through straight to the database unexecuted as plain text and fails. I ALSO tried this with the mode := 'table' parameter and it made no difference.

@if(
   @runtime_stage = 'creating',
   CREATE CLUSTERED INDEX @resolve_template('PK_{table_name}') ON @resolve_template('{schema_name}.{table_name}') ([xxx]) WITH (DATA_COMPRESSION = PAGE)
);
-- CREATE CLUSTERED INDEX [PK_{table_name}] ON [{schema_name}.{table_name}]([xxx]) WITH (DATA_COMPRESSION=PAGE);

I can hardcode like:

@if(
   @runtime_stage = 'creating',
   create clustered index idx_pk on model_schema.model_name (col1, col2)
);

But ideally I'd like to not have to hardcode the index name, model name, and columns for every model when sqlmesh already knows it all internally and I just want to follow a template. Can't I do something like:

@if(
   @runtime_stage = 'creating',
   CREATE CLUSTERED INDEX @resolve_template('PK_{table_name}') ON @resolve_template('{schema_name}.{table_name}') ({model.unique_key}) WITH (DATA_COMPRESSION = PAGE)
);

?

[tobymao]

can you try

@if(@runtime_stage = 'creating' create clustered index @resolve_template('PK_@table_name') ON my_model_schema.model_name (col1, col2)

[codykonior]

I wanted to have the my_model_schema.model_name part dynamically populated. Here's how I got it working on Slack with the help of Toby.

Create a file macros/get_unique_key.py. It may look like I'm checking runtime_stage twice, but, both are necessary:

from sqlmesh.core.macros import macro

@macro()
def get_unique_key(evaluator):
  if evaluator.runtime_stage == 'creating':
    return evaluator.get_snapshot(evaluator._model_fqn).model.kind.unique_key

At the bottom of the SQL model file (after the SELECT statement, that's critical):

@if(
   @runtime_stage = 'creating',
   CREATE CLUSTERED INDEX @resolve_template('PK_@{table_name}') ON @SQL('@this_model') (@get_unique_key()) WITH (DATA_COMPRESSION = PAGE)
);

This way if you just want to create a primary key on your unique columns on any model, you can throw it at the end without re-entering the model name and unique keys you already entered at the top.

Now we can go back to discussing the other pieces.

[georgesittas]

Hey @codykonior, please take a look at this PR: #4599. It aims to explain the difference between @x and @{x}, which applies to this discussion as well. In your case, you need to use @this_model.

If SQLGlot is unable to parse your statement and falls back into parsing a Command, then you can do something like this to still render the physical table name of your model:

JINJA_STATEMENT_BEGIN;

CREATE CLUSTERED INDEX idx_pk on {{ this_model }} (col1)

JINJA_END;

[codykonior]

@treysp Thanks it looks great, but I think it has a bug if you have a model that is composed entirely of unique keys.

MODEL   (
        name [restricted].[Test],
        kind INCREMENTAL_BY_UNIQUE_KEY (
                unique_key ([ClientPrefix], [PolicyName])
        ),
        cron '*/5 * * * *',
        enabled True
);

SELECT  r.ClientPrefix,
        r.PolicyName
FROM
(
        SELECT  N'ABC' AS ClientPrefix,
                CONVERT(NVARCHAR(100), GETDATE()) AS PolicyName
) r
;

• The WHEN MATCHED AND clause will become EXISTS (SELECT EXCEPT SELECT), which is invalid
• The SET clause becomes empty

I'd like to suggest wrapping the match_expressions in this, which causes it to pass:

            if target_columns_no_keys:
                match_expressions = [
                    exp.When(
                        matched=True,
                        source=False,
                        condition=match_condition,
                        then=exp.Update(
                            expressions=[
                                exp.column(col, MERGE_TARGET_ALIAS).eq(
                                    exp.column(col, MERGE_SOURCE_ALIAS)
                                )
                                for col in columns_to_types_no_keys
                            ],
                        ),
                    )
                ]
            else:
                match_expressions = []

Because in that case if it matched there is no need to check for exists and no set statements required either.