fix(query): add partition write regression tests (#20426)

* test(query): add partition write regression coverage

* test(query): fix computed partition column type

Author: KKould

tests/sqllogictests/suites/base/09_fuse_engine/09_0054_partition_by.test

@@ -372,6 +372,13 @@ SELECT id, p FROM merge_partitioned
 2 0
 3 1
 
+# MERGE moves id 1 across the physical partition boundary. The rewritten blocks
+# must retain exact partition metadata for pruning.
+query T
+EXPLAIN SELECT id, p FROM merge_partitioned WHERE p % 2 = 0
+----
+<slt:ignore>partitions total: 3<slt:ignore>partitions scanned: 1<slt:ignore>range pruning: 3 to 1<slt:ignore>
+
 statement ok
 OPTIMIZE TABLE merge_partitioned COMPACT
 
@@ -654,5 +661,57 @@ SELECT p, length(v) FROM ctas_partitioned
 2 1
 3 1
 
+# Missing DEFAULT columns must be materialized before the hash write layout
+# evaluates the partition expression. Reversing those pipeline stages makes the
+# insert fail because p is absent from the input block.
+statement ok
+CREATE TABLE default_partitioned(id INT, p INT DEFAULT 7)
+PARTITION BY (p)
+WRITE_DISTRIBUTION_MODE='hash'
+ROW_PER_BLOCK=100000 BLOCK_PER_SEGMENT=1000
+
+statement ok
+INSERT INTO default_partitioned(id) SELECT number::INT FROM numbers(4)
+
+query III
+SELECT count(*), min(p), max(p) FROM default_partitioned
+----
+4 7 7
+
+# Nullable partition values must survive serialization as exact metadata so
+# that IS NULL can prune every non-NULL partition.
+statement ok
+CREATE TABLE nullable_partitioned(p INT NULL, id BIGINT)
+PARTITION BY (p)
+WRITE_DISTRIBUTION_MODE='hash'
+ROW_PER_BLOCK=100000 BLOCK_PER_SEGMENT=1000
+
+statement ok
+INSERT INTO nullable_partitioned
+SELECT if(number % 4 = 0, NULL, (number % 3)::INT), number FROM numbers(16)
+
+query II
+SELECT count(*), count_if(p IS NULL) FROM nullable_partitioned
+----
+16 4
+
+# NULL is a physical partition value too: no serialized block may mix it with
+# a non-NULL value or another non-NULL partition.
+query I
+SELECT count_if(partition_values != 1)
+FROM (
+    SELECT _block_name,
+           count(DISTINCT if(p IS NULL, '<NULL>', to_string(p))) AS partition_values
+    FROM nullable_partitioned
+    GROUP BY _block_name
+)
+----
+0
+
+query T
+EXPLAIN SELECT id FROM nullable_partitioned WHERE p IS NULL
+----
+<slt:ignore>partitions total: 4<slt:ignore>partitions scanned: 1<slt:ignore>range pruning: 4 to 1<slt:ignore>
+
 statement ok
 DROP DATABASE db_09_0054

tests/sqllogictests/suites/base/09_fuse_engine/09_0100_insert_multi_table.test

@@ -779,3 +779,57 @@ drop table dst1_19716;
 
 statement ok
 drop table dst2_19716;
+
+# ChunkAppendData builds each target's partition transforms independently.
+# Different target expressions must not share or misalign partition metadata.
+statement ok
+CREATE TABLE multi_partitioned_a(p INT, id INT)
+PARTITION BY (p % 2)
+ROW_PER_BLOCK=100000 BLOCK_PER_SEGMENT=1000
+
+statement ok
+CREATE TABLE multi_partitioned_b(id INT)
+PARTITION BY (id % 3)
+ROW_PER_BLOCK=100000 BLOCK_PER_SEGMENT=1000
+
+statement ok
+INSERT ALL
+  WHEN true THEN INTO multi_partitioned_a VALUES(p, id)
+  WHEN id % 2 = 0 THEN INTO multi_partitioned_b VALUES(id)
+SELECT (number % 4)::INT AS p, number::INT AS id FROM numbers(16)
+
+query I
+SELECT count(*) FROM multi_partitioned_a
+----
+16
+
+query I
+SELECT count(*) FROM multi_partitioned_b
+----
+8
+
+query I
+SELECT count_if(partition_values != 1)
+FROM (
+    SELECT _block_name, count(DISTINCT p % 2) AS partition_values
+    FROM multi_partitioned_a
+    GROUP BY _block_name
+)
+----
+0
+
+query I
+SELECT count_if(partition_values != 1)
+FROM (
+    SELECT _block_name, count(DISTINCT id % 3) AS partition_values
+    FROM multi_partitioned_b
+    GROUP BY _block_name
+)
+----
+0
+
+statement ok
+DROP TABLE multi_partitioned_a
+
+statement ok
+DROP TABLE multi_partitioned_b

tests/sqllogictests/suites/ee/02_computed_column/02_0001_partition_by.test

@@ -0,0 +1,49 @@
+## Copyright 2026 Databend Cloud
+##
+## Licensed under the Elastic License, Version 2.0 (the "License");
+## you may not use this file except in compliance with the License.
+## You may obtain a copy of the License at
+##
+##     https://www.elastic.co/licensing/elastic-license
+##
+## Unless required by applicable law or agreed to in writing, software
+## distributed under the License is distributed on an "AS IS" BASIS,
+## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+## See the License for the specific language governing permissions and
+## limitations under the License.
+
+statement ok
+DROP DATABASE IF EXISTS computed_partition_by
+
+statement ok
+CREATE DATABASE computed_partition_by
+
+statement ok
+USE computed_partition_by
+
+# Stored computed partition columns must be generated before hash routing and
+# physical partition layout.
+statement ok
+CREATE TABLE computed_partitioned(
+    id INT,
+    p SMALLINT AS (id % 2) STORED
+)
+PARTITION BY (p)
+WRITE_DISTRIBUTION_MODE='hash'
+ROW_PER_BLOCK=100000 BLOCK_PER_SEGMENT=1000
+
+statement ok
+INSERT INTO computed_partitioned(id) SELECT number::INT FROM numbers(4)
+
+query III
+SELECT count(*), min(p), max(p) FROM computed_partitioned
+----
+4 0 1
+
+query T
+EXPLAIN SELECT id FROM computed_partitioned WHERE p = 0
+----
+<slt:ignore>partitions total: 2<slt:ignore>partitions scanned: 1<slt:ignore>range pruning: 2 to 1<slt:ignore>
+
+statement ok
+DROP DATABASE computed_partition_by