Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
starrocks: add some test cases for partition
Signed-off-by: jaogoy <jaogoy@gmail.com>
  • Loading branch information
jaogoy committed Jan 16, 2026
commit e6b28257b512a9446dd908eab8d608b1a8fad766
157 changes: 157 additions & 0 deletions tests/core/engine_adapter/integration/test_integration_starrocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1434,6 +1434,163 @@ def test_e2e_partition_list(self, starrocks_adapter: StarRocksEngineAdapter):
finally:
starrocks_adapter.drop_schema(db_name, ignore_if_not_exists=True)

# ========================================
# Case 6B: Expression Partitioning (table + MV)
# Covers: expression partitioning with/without functions, table vs MV
# ========================================

@pytest.mark.parametrize(
"partition_expr",
[
"(event_date, region)", # plain columns
"date_trunc('day', event_date)", # one function expression
"(from_unixtime(ts), region)", # multiple expressions
],
)
def test_e2e_partition_expression_for_table(
self,
starrocks_adapter: StarRocksEngineAdapter,
partition_expr: str,
):
"""Expression partitioning for regular tables (outer paren only when no functions)."""
db_name = "sr_e2e_part_expr_tbl_db"
table_name = f"{db_name}.sr_part_expr_table"

model_sql = f"""
MODEL (
name test.partition_expr_table,
kind FULL,
dialect starrocks,
columns (
id BIGINT,
ts BIGINT,
event_date DATE,
region VARCHAR(50)
),
partitioned_by {partition_expr},
);
SELECT *
"""

try:
starrocks_adapter.create_schema(db_name, ignore_if_exists=True)
params = self._parse_model_and_get_all_params(model_sql)
starrocks_adapter.create_table(table_name, **params)

ddl = fetchone_or_fail(starrocks_adapter, f"SHOW CREATE TABLE {table_name}")[1]
logger.info(f"Case 6B DDL:\n{ddl}")
ddl_upper = ddl.upper()
assert "PARTITION BY" in ddl_upper

before, after = ddl_upper.split("PARTITION BY", 1)
after = after.lstrip()

# Column/function presence
if "DATE" in partition_expr.upper():
assert "EVENT_DATE" in after
else:
assert "REGION" in after
if "FROM_UNIXTIME" in partition_expr.upper():
assert "FROM_UNIXTIME" in after or \
("FROM_UNIXTIME" in before and "__GENERATED_PARTITION_COLUMN" in after)
if "DATE_TRUNC" in partition_expr.upper():
assert "DATE_TRUNC" in after
finally:
starrocks_adapter.drop_schema(db_name, ignore_if_not_exists=True)

@pytest.mark.parametrize(
"partition_clause,has_func",
[
("(event_date, region)", False),
("(date_trunc('day', event_date), region)", True),
],
)
def test_e2e_partition_expression_for_mv(
self,
starrocks_adapter: StarRocksEngineAdapter,
partition_clause: str,
has_func: bool,
):
"""Expression partitioning for MVs should always keep outer parentheses."""
db_name = "sr_e2e_part_expr_mv_db"
src_table = f"{db_name}.sr_part_expr_src"
mv_table = f"{db_name}.sr_part_expr_mv"

try:
starrocks_adapter.create_schema(db_name, ignore_if_exists=True)

# Source table + data
starrocks_adapter.create_table(
src_table,
target_columns_to_types={
"id": exp.DataType.build("BIGINT"),
"ts": exp.DataType.build("BIGINT"),
"event_date": exp.DataType.build("DATE"),
"region": exp.DataType.build("VARCHAR(50)"),
},
primary_key=("id", "event_date", "region"),
table_properties={
"partitioned_by": partition_clause,
},
)
starrocks_adapter.execute(
f"""
INSERT INTO {src_table} (id, ts, event_date, region)
VALUES (1, 1700000000, '2024-01-01', 'us')
"""
)

model_sql = f"""
MODEL (
name test.partition_expr_mv,
kind VIEW (
materialized true
),
dialect starrocks,
columns (
id BIGINT,
ts BIGINT,
event_date DATE,
region VARCHAR(50)
),
partitioned_by {partition_clause},
physical_properties (
distributed_by = 'HASH(id) BUCKETS 2',
refresh_moment = 'IMMEDIATE',
refresh_scheme = 'ASYNC'
)
);
SELECT id, ts, event_date, region FROM {src_table};
"""

model = _load_sql_model(model_sql)
query = model.render_query()
assert query is not None
materialized_properties = _materialized_properties_from_model(model)

starrocks_adapter.create_view(
mv_table,
query,
replace=True,
materialized=True,
target_columns_to_types=model.columns_to_types,
materialized_properties=materialized_properties,
view_properties=model.physical_properties,
)

ddl = fetchone_or_fail(
starrocks_adapter, f"SHOW CREATE MATERIALIZED VIEW {mv_table}"
)[1]
logger.info(f"Case 6B DDL:\n{ddl}")
ddl_upper = ddl.upper()
assert "PARTITION BY" in ddl_upper
after = ddl_upper.split("PARTITION BY", 1)[1].lstrip()
assert after.startswith("("), f"MV partition should keep parentheses, got: {after[:50]}"
assert "REGION" in after
assert ("DATE_TRUNC" in after) == has_func
finally:
starrocks_adapter.drop_schema(db_name, ignore_if_not_exists=True)

# ========================================
# Case 7: Other Key Types (test_design.md Case 7)
# Covers: duplicate_key, unique_key, aggregate_key
Expand Down
71 changes: 64 additions & 7 deletions tests/core/engine_adapter/test_starrocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -809,9 +809,11 @@ class TestPartitionPropertyBuilding:
"partition_expr,expected_clause",
[
# Expression partitioning - single column
("'dt'", "PARTITION BY (dt)"),
("'dt'", "PARTITION BY (`dt`)"),
# Expression partitioning - multi-column
("(year, month)", "PARTITION BY (year, month)"),
("(year, month)", "PARTITION BY (`year`, `month`)"),
# Expression partitioning - multi-column with func
("(date_trunc('day', dt), region)", "PARTITION BY DATE_TRUNC('DAY', `dt`), `region`"),
# RANGE partitioning
("RANGE (dt)", "PARTITION BY RANGE (`dt`) ()"),
# LIST partitioning
Expand Down Expand Up @@ -852,6 +854,61 @@ def test_partitioned_by_forms(
sql = to_sql_calls(adapter)[0]
assert expected_clause in sql

@pytest.mark.parametrize(
"partition_expr,expected_clause",
[
("(year, month)", "PARTITION BY (`year`, `month`)"),
(
"(date_trunc('day', dt), region)",
"PARTITION BY (DATE_TRUNC('DAY', `dt`), `region`)",
),
(
"(from_unixtime(dt))",
"PARTITION BY (FROM_UNIXTIME(`dt`))",
),
],
)
def test_partitioned_by_forms_for_mv(
self,
make_mocked_engine_adapter: t.Callable[..., StarRocksEngineAdapter],
partition_expr: str,
expected_clause: str,
):
"""MV partition_by should keep outer parentheses when rendering partition tuples."""
model_sql = f"""
MODEL (
name test_schema.test_mv_partition,
kind VIEW (
materialized true,
),
dialect starrocks,
columns (dt DATE, region STRING, year INT, month INT),
physical_properties (
partition_by = {partition_expr}
)
);
SELECT dt, region, year, month FROM src;
"""

model = _load_sql_model(model_sql)
materialized_properties = (
{"partitioned_by": model.partitioned_by} if model.partitioned_by else None
)

adapter = make_mocked_engine_adapter(StarRocksEngineAdapter)
adapter.create_view(
model.name,
model.render_query(),
materialized=True,
replace=False,
target_columns_to_types=_columns(model),
materialized_properties=materialized_properties,
view_properties=model.physical_properties,
)

sql = to_sql_calls(adapter)[0]
assert expected_clause in sql

def test_partition_by_alias(
self, make_mocked_engine_adapter: t.Callable[..., StarRocksEngineAdapter]
):
Expand Down Expand Up @@ -881,7 +938,7 @@ def test_partition_by_alias(
)

sql = to_sql_calls(adapter)[0]
assert "PARTITION BY (year, month)" in sql
assert "PARTITION BY (year, month)" in sql or "PARTITION BY (`year`, `month`)" in sql

def test_partitioned_by_as_model_parameter(
self, make_mocked_engine_adapter: t.Callable[..., StarRocksEngineAdapter]
Expand Down Expand Up @@ -909,7 +966,7 @@ def test_partitioned_by_as_model_parameter(
)

sql = to_sql_calls(adapter)[0]
assert "PARTITION BY (year, month)" in sql
assert "PARTITION BY (year, month)" in sql or "PARTITION BY (`year`, `month`)" in sql

def test_partitions_value_forms(
self, make_mocked_engine_adapter: t.Callable[..., StarRocksEngineAdapter]
Expand Down Expand Up @@ -1354,10 +1411,10 @@ def _build_mv_model(self, property_sql: str) -> SqlModel:
model_sql = f"""
MODEL (
name test_schema.test_mv_refresh_model,
kind VIEW,
kind VIEW (materialized true),
dialect starrocks,
columns (a INT),
virtual_properties (
physical_properties (
{property_sql}
)
);
Expand All @@ -1377,7 +1434,7 @@ def _create_simple_mv(
replace=False,
materialized=True,
target_columns_to_types=_columns(model),
view_properties=model.virtual_properties,
view_properties=model.physical_properties,
)
# replace=False → only CREATE statement is emitted
return to_sql_calls(adapter)[-1]
Expand Down
2 changes: 1 addition & 1 deletion tests/core/test_connection_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -1971,7 +1971,7 @@ def test_starrocks(make_config):
assert config.database == "testdb"
assert config.DIALECT == "starrocks"
assert config.DISPLAY_NAME == "StarRocks"
assert config.DISPLAY_ORDER == 19
assert config.DISPLAY_ORDER == 18
assert config.is_recommended_for_state_sync is False

# Test with minimal configuration (using default port)
Expand Down