refactor partition

Author: xinge-ji

sqlmesh/core/engine_adapter/doris.py

@@ -485,6 +485,51 @@ def _create_table_from_columns(
             **kwargs,
         )
 
+    def _parse_partition_expressions(
+        self, partitioned_by: t.List[exp.Expression]
+    ) -> t.Tuple[t.List[exp.Expression], t.Optional[str]]:
+        """Parse partition expressions and extract partition kind and normalized columns.
+
+        Returns:
+            Tuple of (normalized_partitioned_by, partition_kind)
+        """
+        parsed_partitioned_by: t.List[exp.Expression] = []
+        partition_kind: t.Optional[str] = None
+
+        for expr in partitioned_by:
+            try:
+                # Handle Anonymous function calls like RANGE(col) or LIST(col)
+                if isinstance(expr, exp.Anonymous) and expr.this:
+                    func_name = str(expr.this).upper()
+                    if func_name in ("RANGE", "LIST"):
+                        partition_kind = func_name
+                        # Extract column expressions from function arguments
+                        for arg in expr.expressions:
+                            if isinstance(arg, exp.Column):
+                                parsed_partitioned_by.append(arg)
+                            else:
+                                # Convert other expressions to columns if possible
+                                parsed_partitioned_by.append(exp.to_column(str(arg)))
+                        continue
+
+                # Handle literal strings like "RANGE(col)" or "LIST(col)"
+                if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False):
+                    text = str(expr.this)
+                    match = re.match(r"^\s*(RANGE|LIST)\s*\((.*?)\)\s*$", text, flags=re.IGNORECASE)
+                    if match:
+                        partition_kind = match.group(1).upper()
+                        inner = match.group(2)
+                        inner_cols = [c.strip().strip("`") for c in inner.split(",") if c.strip()]
+                        for col in inner_cols:
+                            parsed_partitioned_by.append(exp.to_column(col))
+                        continue
+            except Exception:
+                # If anything goes wrong, keep the original expr
+                pass
+            parsed_partitioned_by.append(expr)
+
+        return parsed_partitioned_by, partition_kind
+
     def _build_partitioned_by_exp(
         self,
         partitioned_by: t.List[exp.Expression],
@@ -505,12 +550,11 @@ def _build_partitioned_by_exp(
 
         Supports both RANGE and LIST partition syntaxes using sqlglot's doris dialect nodes.
         The partition kind is chosen by:
-        - kwargs["partition_kind"] if provided (expects 'RANGE' or 'LIST', case-insensitive)
+        - inferred from partitioned_by expressions like 'RANGE(col)' or 'LIST(col)'
         - otherwise inferred from the provided 'partitions' strings: if any contains 'VALUES IN' -> LIST; else RANGE.
         """
         partitions = kwargs.get("partitions")
         create_expressions = None
-        partition_kind: t.Optional[str] = kwargs.get("partition_kind")
 
         def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var:
             # If it's a Literal, extract the string and wrap as Var (no quotes)
@@ -525,6 +569,9 @@ def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var:
             # Fallback: return as is
             return expr
 
+        # Parse partition kind and columns from partitioned_by expressions
+        partitioned_by, partition_kind = self._parse_partition_expressions(partitioned_by)
+
         if partitions:
             if isinstance(partitions, exp.Tuple):
                 create_expressions = [
@@ -555,8 +602,8 @@ def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var:
         try:
             if is_list:
                 return exp.PartitionByListProperty(
-                    this=exp.Schema(expressions=partitioned_by),
-                    partitions=create_expressions,
+                    partition_expressions=partitioned_by,
+                    create_expressions=create_expressions,
                 )
             return exp.PartitionByRangeProperty(
                 partition_expressions=partitioned_by,
@@ -800,46 +847,43 @@ def _parse_trigger_string(
                     )
 
         # Handle duplicate_key - only handle Tuple expressions or single Column expressions
+        # Both tables and materialized views support duplicate keys in Doris
         duplicate_key = table_properties_copy.pop("duplicate_key", None)
         if duplicate_key is not None:
-            if not is_materialized_view:
-                if isinstance(duplicate_key, exp.Tuple):
-                    # Extract column names from Tuple expressions
-                    column_names = []
-                    for expr in duplicate_key.expressions:
-                        if (
-                            isinstance(expr, exp.Column)
-                            and hasattr(expr, "this")
-                            and hasattr(expr.this, "this")
-                        ):
-                            column_names.append(str(expr.this.this))
-                        elif hasattr(expr, "this"):
-                            column_names.append(str(expr.this))
-                        else:
-                            column_names.append(str(expr))
-                    properties.append(
-                        exp.DuplicateKeyProperty(
-                            expressions=[exp.to_column(k) for k in column_names]
-                        )
-                    )
-                elif isinstance(duplicate_key, exp.Column):
-                    # Handle as single column
-                    if hasattr(duplicate_key, "this") and hasattr(duplicate_key.this, "this"):
-                        column_name = str(duplicate_key.this.this)
+            if isinstance(duplicate_key, exp.Tuple):
+                # Extract column names from Tuple expressions
+                column_names = []
+                for expr in duplicate_key.expressions:
+                    if (
+                        isinstance(expr, exp.Column)
+                        and hasattr(expr, "this")
+                        and hasattr(expr.this, "this")
+                    ):
+                        column_names.append(str(expr.this.this))
+                    elif hasattr(expr, "this"):
+                        column_names.append(str(expr.this))
                     else:
-                        column_name = str(duplicate_key.this)
-                    properties.append(
-                        exp.DuplicateKeyProperty(expressions=[exp.to_column(column_name)])
-                    )
-                elif isinstance(duplicate_key, exp.Literal):
-                    properties.append(
-                        exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key.this)])
-                    )
-                elif isinstance(duplicate_key, str):
-                    properties.append(
-                        exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key)])
-                    )
-            # Note: Materialized views don't typically use duplicate_key, so we skip it
+                        column_names.append(str(expr))
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(k) for k in column_names])
+                )
+            elif isinstance(duplicate_key, exp.Column):
+                # Handle as single column
+                if hasattr(duplicate_key, "this") and hasattr(duplicate_key.this, "this"):
+                    column_name = str(duplicate_key.this.this)
+                else:
+                    column_name = str(duplicate_key.this)
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(column_name)])
+                )
+            elif isinstance(duplicate_key, exp.Literal):
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key.this)])
+                )
+            elif isinstance(duplicate_key, str):
+                properties.append(
+                    exp.DuplicateKeyProperty(expressions=[exp.to_column(duplicate_key)])
+                )
 
         if table_description:
             properties.append(
@@ -851,30 +895,8 @@ def _parse_trigger_string(
         # Handle partitioning
         add_partition = True
         if partitioned_by:
-            normalized_partitioned_by: t.List[exp.Expression] = []
-            for expr in partitioned_by:
-                try:
-                    # Handle literal strings like "RANGE(col)" or "LIST(col)"
-                    if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False):
-                        text = str(expr.this)
-                        match = re.match(
-                            r"^\s*(RANGE|LIST)\s*\((.*?)\)\s*$", text, flags=re.IGNORECASE
-                        )
-                        if match:
-                            inner = match.group(2)
-                            inner_cols = [
-                                c.strip().strip("`") for c in inner.split(",") if c.strip()
-                            ]
-                            for col in inner_cols:
-                                normalized_partitioned_by.append(exp.to_column(col))
-                            continue
-                except Exception:
-                    # If anything goes wrong, keep the original expr
-                    pass
-                normalized_partitioned_by.append(expr)
-
-            # Replace with normalized expressions
-            partitioned_by = normalized_partitioned_by
+            # Parse and normalize partition expressions
+            partitioned_by, _ = self._parse_partition_expressions(partitioned_by)
             # For tables, check if partitioned_by columns are in unique_key; for materialized views, allow regardless
             if unique_key is not None and not is_materialized_view:
                 # Extract key column names from unique_key (only Tuple or Column expressions)

tests/core/engine_adapter/test_doris.py

@@ -280,7 +280,7 @@ def test_create_table_with_partitioned_by(
     adapter.create_table(
         "test_table",
         target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("DATE")},
-        partitioned_by=[exp.to_column("b")],
+        partitioned_by=[exp.Literal.string("RANGE(b)")],
         table_properties={
             "partitions": exp.Literal.string(
                 "FROM ('2000-11-14') TO ('2021-11-14') INTERVAL 2 YEAR"
@@ -293,6 +293,168 @@ def test_create_table_with_partitioned_by(
     ]
 
 
+def test_create_table_with_range_partitioned_by_with_partitions(
+    make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(DorisEngineAdapter)
+    adapter.create_table(
+        "test_table",
+        target_columns_to_types={
+            "id": exp.DataType.build("INT"),
+            "waiter_id": exp.DataType.build("INT"),
+            "customer_id": exp.DataType.build("INT"),
+            "ds": exp.DataType.build("DATETIME"),
+        },
+        partitioned_by=[exp.Literal.string("RANGE(ds)")],
+        table_properties={
+            "partitions": exp.Tuple(
+                expressions=[
+                    exp.Literal.string('PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01"))'),
+                    exp.Literal.string('PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01"))'),
+                    exp.Literal.string('PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01"))'),
+                    exp.Literal.string("PARTITION `other` VALUES LESS THAN MAXVALUE"),
+                ]
+            ),
+            "distributed_by": exp.Tuple(
+                expressions=[
+                    exp.EQ(
+                        this=exp.Column(this=exp.Identifier(this="kind", quoted=True)),
+                        expression=exp.Literal.string("HASH"),
+                    ),
+                    exp.EQ(
+                        this=exp.Column(this=exp.Identifier(this="expressions", quoted=True)),
+                        expression=exp.Column(this=exp.Identifier(this="id", quoted=True)),
+                    ),
+                    exp.EQ(
+                        this=exp.Column(this=exp.Identifier(this="buckets", quoted=True)),
+                        expression=exp.Literal.number(10),
+                    ),
+                ]
+            ),
+            "replication_allocation": exp.Literal.string("tag.location.default: 3"),
+            "in_memory": exp.Literal.string("false"),
+            "storage_format": exp.Literal.string("V2"),
+            "disable_auto_compaction": exp.Literal.string("false"),
+        },
+    )
+
+    expected_sql = (
+        "CREATE TABLE IF NOT EXISTS `test_table` "
+        "(`id` INT, `waiter_id` INT, `customer_id` INT, `ds` DATETIME) "
+        "PARTITION BY RANGE (`ds`) "
+        '(PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01")), '
+        'PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01")), '
+        'PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01")), '
+        "PARTITION `other` VALUES LESS THAN MAXVALUE) "
+        "DISTRIBUTED BY HASH (`id`) BUCKETS 10 "
+        "PROPERTIES ("
+        "'replication_allocation'='tag.location.default: 3', "
+        "'in_memory'='false', "
+        "'storage_format'='V2', "
+        "'disable_auto_compaction'='false')"
+    )
+
+    assert to_sql_calls(adapter) == [expected_sql]
+
+
+def test_create_table_with_list_partitioned_by(
+    make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(DorisEngineAdapter)
+    adapter.create_table(
+        "test_table",
+        target_columns_to_types={
+            "id": exp.DataType.build("INT"),
+            "status": exp.DataType.build("VARCHAR(10)"),
+        },
+        partitioned_by=[exp.Literal.string("LIST(status)")],
+        table_properties={
+            "partitions": exp.Tuple(
+                expressions=[
+                    exp.Literal.string('PARTITION `active` VALUES IN ("active", "pending")'),
+                    exp.Literal.string('PARTITION `inactive` VALUES IN ("inactive", "disabled")'),
+                ]
+            ),
+        },
+    )
+
+    expected_sql = (
+        "CREATE TABLE IF NOT EXISTS `test_table` "
+        "(`id` INT, `status` VARCHAR(10)) "
+        "PARTITION BY LIST (`status`) "
+        '(PARTITION `active` VALUES IN ("active", "pending"), '
+        'PARTITION `inactive` VALUES IN ("inactive", "disabled"))'
+    )
+
+    assert to_sql_calls(adapter) == [expected_sql]
+
+
+def test_create_table_with_range_partitioned_by_anonymous_function(
+    make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
+):
+    """Test that RANGE(ds) function call syntax generates correct SQL without duplicate RANGE."""
+    adapter = make_mocked_engine_adapter(DorisEngineAdapter)
+    adapter.create_table(
+        "test_table",
+        target_columns_to_types={
+            "id": exp.DataType.build("INT"),
+            "ds": exp.DataType.build("DATETIME"),
+        },
+        # This simulates how partitioned_by RANGE(ds) gets parsed from model definition
+        partitioned_by=[exp.Anonymous(this="RANGE", expressions=[exp.to_column("ds")])],
+        table_properties={
+            "partitions": exp.Literal.string(
+                'FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH'
+            )
+        },
+    )
+
+    expected_sql = (
+        "CREATE TABLE IF NOT EXISTS `test_table` "
+        "(`id` INT, `ds` DATETIME) "
+        "PARTITION BY RANGE (`ds`) "
+        '(FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH)'
+    )
+
+    assert to_sql_calls(adapter) == [expected_sql]
+
+
+def test_create_materialized_view_with_duplicate_key(
+    make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
+):
+    adapter = make_mocked_engine_adapter(DorisEngineAdapter)
+    adapter.create_view(
+        "test_mv",
+        parse_one("SELECT id, status, COUNT(*) as cnt FROM orders GROUP BY id, status"),
+        materialized=True,
+        target_columns_to_types={
+            "id": exp.DataType.build("INT"),
+            "status": exp.DataType.build("VARCHAR(10)"),
+            "cnt": exp.DataType.build("BIGINT"),
+        },
+        view_properties={
+            "duplicate_key": exp.Tuple(
+                expressions=[
+                    exp.to_column("id"),
+                    exp.to_column("status"),
+                ]
+            ),
+        },
+    )
+
+    expected_sqls = [
+        "DROP MATERIALIZED VIEW IF EXISTS `test_mv`",
+        (
+            "CREATE MATERIALIZED VIEW `test_mv` "
+            "(`id`, `status`, `cnt`) "
+            "DUPLICATE KEY (`id`, `status`) "
+            "AS SELECT `id`, `status`, COUNT(*) AS `cnt` FROM `orders` GROUP BY `id`, `status`"
+        ),
+    ]
+
+    assert to_sql_calls(adapter) == expected_sqls
+
+
 def test_create_full_materialized_view(
     make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
 ):