Merge pull request #12 from ydb-platform/incremental_merge

Incremental materializations support

Author: vgvoleg

README.md

@@ -76,6 +76,17 @@ profile_name:
 | `auto_partitioning_partition_size_mb` | Partition size in megabytes for automatic partitioning | `no` | |
 | `ttl` | Time-to-live (TTL) expression for automatic data expiration | `no` | |
 
+#### Incremental
+
+| Option | Description | Required | Default |
+| ------ | ----------- | -------- | ------- |
+| `incremental_strategy` | Strategy of incremental materialization. Current adapter supports only `merge` strategy, which will use `YDB`'s `UPSERT` operation. | `no` | `default` |
+| `primary_key` | Primary key expression to use during table creation | `yes` | |
+| `store_type` | Type of table. Available options are `row` and `column` | `no` | `row` |
+| `auto_partitioning_by_size` | Enable automatic partitioning by size. Available options are `ENABLED` and `DISABLED` | `no` | |
+| `auto_partitioning_partition_size_mb` | Partition size in megabytes for automatic partitioning | `no` | |
+| `ttl` | Time-to-live (TTL) expression for automatic data expiration | `no` | |
+
 ##### Example table configuration
 
 ```sql

dbt/adapters/ydb/column.py

@@ -11,16 +11,13 @@
 @dataclass
 class YDBColumn(Column):
     TYPE_LABELS = {
-        'STRING': 'String',
+        'STRING': 'Text',
         'TIMESTAMP': 'DateTime',
-        'FLOAT': 'Float32',
-        'INTEGER': 'Int32',
+        'FLOAT': 'Float',
+        'INTEGER': 'Int64',
     }
     is_nullable: bool = False
-    is_low_cardinality: bool = False
-    _low_card_regex = re.compile(r'^LowCardinality\((.*)\)$')
     _nullable_regex = re.compile(r'^(.*)\?$')
-    _fix_size_regex = re.compile(r'FixedString\((.*?)\)')
     _decimal_regex = re.compile(r'Decimal\((\d+), (\d+)\)')
 
     def __init__(self, column: str, dtype: str) -> None:
@@ -30,11 +27,6 @@ def __init__(self, column: str, dtype: str) -> None:
 
         dtype = self._inner_dtype(dtype)
 
-        if dtype.lower().startswith('fixedstring'):
-            match_sized = self._fix_size_regex.search(dtype)
-            if match_sized:
-                char_size = int(match_sized.group(1))
-
         if dtype.lower().startswith('decimal'):
             match_dec = self._decimal_regex.search(dtype)
             numeric_precision = 0
@@ -57,26 +49,16 @@ def data_type(self) -> str:
         else:
             data_t = self.dtype
 
-        if self.is_nullable or self.is_low_cardinality:
-            data_t = self.nested_type(data_t, self.is_low_cardinality, self.is_nullable)
+        if self.is_nullable:
+            data_t = self.nested_type(data_t, self.is_nullable)
 
         return data_t
 
     def is_string(self) -> bool:
         return self.dtype.lower() in [
-            'string',
-            'fixedstring',
-            'longblob',
-            'longtext',
-            'tinytext',
             'text',
-            'varchar',
-            'mediumblob',
-            'blob',
-            'tinyblob',
-            'char',
-            'mediumtext',
-        ] or self.dtype.lower().startswith('fixedstring')
+            'utf8',
+        ]
 
     def is_integer(self) -> bool:
         return self.dtype.lower().startswith('int') or self.dtype.lower().startswith('uint')
@@ -98,17 +80,15 @@ def string_size(self) -> int:
 
     @classmethod
     def string_type(cls, size: int) -> str:
-        return 'String'
+        return 'Text'
 
     @classmethod
     def numeric_type(cls, dtype: str, precision: Any, scale: Any) -> str:
         return f'Decimal({precision}, {scale})'
 
     @classmethod
-    def nested_type(cls, dtype: str, is_low_cardinality: bool, is_nullable: bool) -> str:
+    def nested_type(cls, dtype: str, is_nullable: bool) -> str:
         template = "{}"
-        if is_low_cardinality:
-            template = template.format("LowCardinality({})")
         if is_nullable:
             template = template.format("Nullable({})")
         return template.format(dtype)
@@ -125,10 +105,6 @@ def can_expand_to(self, other_column: 'Column') -> bool:
     def _inner_dtype(self, dtype) -> str:
         inner_dtype = dtype.strip()
 
-        if low_card_match := self._low_card_regex.search(inner_dtype):
-            self.is_low_cardinality = True
-            inner_dtype = low_card_match.group(1)
-
         if null_match := self._nullable_regex.search(inner_dtype):
             self.is_nullable = True
             inner_dtype = null_match.group(1)

dbt/adapters/ydb/impl.py

@@ -119,7 +119,7 @@ def get_columns_in_relation(self, relation):
             return []
 
         with dbapi_connection.cursor() as cur:
-            cur.execute(f"SELECT * FROM {relation} LIMIT 0")
+            cur.execute(f"SELECT * FROM {relation} LIMIT 1")
             return [YDBColumn(col[0], col[1]) for col in cur.description]
 
     def get_rows_different_sql(
@@ -292,6 +292,22 @@ def parse_columns_from_relation(self, relation: BaseRelation) -> List[Dict[str,
             columns.append(column)
         return columns
 
+    def valid_incremental_strategies(self):
+        """The set of standard builtin strategies which this adapter supports out-of-the-box.
+        Not used to validate custom strategies defined by end users.
+        """
+        return ["merge"]
+
+    @available.parse_none
+    def get_column_schema_from_query(self, sql: str, *_) -> List[YDBColumn]:
+        logger.info(f"Try to get column schema from query: \n{sql}")
+        connection = self.connections.get_thread_connection()
+        dbapi_connection = connection.handle
+
+        with dbapi_connection.cursor() as cur:
+            cur.execute(sql)
+            return [YDBColumn(col[0], col[1]) for col in cur.description]
+
 
 
 COLUMNS_EQUAL_SQL = '''

dbt/include/ydb/macros/adapters.sql

@@ -187,5 +187,5 @@ Example 3 of 3 of required macros that does not have a default implementation.
 
 {% macro ydb__current_timestamp() -%}
 '''Returns current UTC time'''
-{# docs show not to be implemented currently. #}
+CurrentUtcTimestamp()
 {% endmacro %}

dbt/include/ydb/macros/materializations/models/incremental/incremental.sql

@@ -44,7 +44,8 @@
       {% set relation_for_indexes = intermediate_relation %}
       {% set need_swap = true %}
   {% else %}
-    {% do run_query(get_create_table_as_sql(True, temp_relation, sql)) %}
+    {% do run_query(get_create_table_as_sql(False, temp_relation, sql)) %}
+    {% do to_drop.append(temp_relation) %}
     {% set relation_for_indexes = temp_relation %}
     {% set contract_config = config.get('contract') %}
     {% if not contract_config or not contract_config.enforced %}
@@ -104,4 +105,44 @@
 
   {{ return({'relations': [target_relation]}) }}
 
-{%- endmaterialization %}
+{%- endmaterialization %}
+
+{% macro ydb__get_merge_sql(target, source, unique_key, dest_columns, incremental_predicates=none) -%}
+    {%- set predicates = [] if incremental_predicates is none else [] + incremental_predicates -%}
+    {%- set dest_cols_csv = get_quoted_csv(dest_columns | map(attribute="name")) -%}
+    {%- set merge_update_columns = config.get('merge_update_columns') -%}
+    {%- set merge_exclude_columns = config.get('merge_exclude_columns') -%}
+    {%- set update_columns = get_merge_update_columns(merge_update_columns, merge_exclude_columns, dest_columns) -%}
+    {%- set sql_header = config.get('sql_header', none) -%}
+
+    {% if unique_key %}
+        {% if unique_key is sequence and unique_key is not mapping and unique_key is not string %}
+            {% for key in unique_key %}
+                {% set this_key_match %}
+                    DBT_INTERNAL_SOURCE.{{ key }} = DBT_INTERNAL_DEST.{{ key }}
+                {% endset %}
+                {% do predicates.append(this_key_match) %}
+            {% endfor %}
+        {% else %}
+            {% set source_unique_key = ("DBT_INTERNAL_SOURCE." ~ unique_key) | trim %}
+	    {% set target_unique_key = ("DBT_INTERNAL_DEST." ~ unique_key) | trim %}
+	    {% set unique_key_match = equals(source_unique_key, target_unique_key) | trim %}
+            {% do predicates.append(unique_key_match) %}
+        {% endif %}
+    {% else %}
+        {% do predicates.append('FALSE') %}
+    {% endif %}
+
+    {{ sql_header if sql_header is not none }}
+
+    upsert into {{ target }}
+    select {{ dest_cols_csv }} from {{ source }}
+
+{% endmacro %}
+
+
+{% macro ydb__get_incremental_default_sql(arg_dict) %}
+
+  {% do return(get_incremental_merge_sql(arg_dict)) %}
+
+{% endmacro %}

examples/jaffle_shop/models/orders.sql

@@ -18,7 +18,6 @@ select
 
 from {{ ref('stg_orders') }} as orders
 
-
 left join (
     select
         order_id,

tests/functional/adapter/basic/test_incremental.py

@@ -1,9 +1,47 @@
-# from dbt.tests.adapter.basic.test_incremental import BaseIncremental, BaseIncrementalNotSchemaChange
+import pytest
 
+from dbt.tests.adapter.basic import files
+from dbt.tests.adapter.basic.test_incremental import BaseIncremental, BaseIncrementalNotSchemaChange
 
-# class TestIncremental(BaseIncremental):
-#     pass
 
+config_materialized_incremental = """
+  {{ config(materialized="incremental", primary_key="id") }}
+"""
 
-# class TestIncrementalNotSchemaChange(BaseIncrementalNotSchemaChange):
-#     pass
+model_incremental = """
+select * from {{ source('raw', 'seed') }}
+{% if is_incremental() %}
+where id > 10
+{% endif %}
+""".strip()
+
+incremental_not_schema_change_sql = """
+{{ config(
+    materialized="incremental",
+    primary_key="user_id_current_time",
+    unique_key="user_id_current_time",
+    on_schema_change="sync_all_columns"
+) }}
+select
+    '1' || '-' || Cast(CurrentUtcTimestamp() as Text) as user_id_current_time,
+    {% if is_incremental() %}
+        'thisis18characters' as platform
+    {% else %}
+        'okthisis20characters' as platform
+    {% endif %}
+"""
+
+class TestIncremental(BaseIncremental):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "incremental.sql": config_materialized_incremental + model_incremental,
+            "schema.yml": files.schema_base_yml,
+        }
+
+
+
+class TestIncrementalNotSchemaChange(BaseIncrementalNotSchemaChange):
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"incremental_not_schema_change.sql": incremental_not_schema_change_sql}