"""

Integration tests for adaptive joins: balancer policies in the circuit JSON profile.

Expects a running Feldera instance (see other tests under python/tests/runtime).

Adaptive join rebalancing is only compiled in when ``num_workers > 1``.

"""

from __future__ import annotations

import unittest

from feldera import PipelineBuilder

from feldera.enums import BootstrapPolicy

from feldera.runtime_config import RuntimeConfig

from feldera.testutils import FELDERA_TEST_NUM_WORKERS, FELDERA_TEST_NUM_HOSTS

from tests import TEST_CLIENT, unique_pipeline_name

from tests.platform.helper import (

API_PREFIX,

http_request,

wait_for_condition,

wait_for_program_success,

)

SQL_JOIN_AB = """

CREATE TABLE tab_a (

k BIGINT NOT NULL,

v BIGINT NOT NULL

) WITH ('materialized' = 'true');

CREATE TABLE tab_b (

k BIGINT NOT NULL,

w BIGINT NOT NULL

) WITH ('materialized' = 'true');

CREATE MATERIALIZED VIEW join_ab AS

SELECT tab_a.k, tab_a.v, tab_b.w

FROM tab_a JOIN tab_b ON tab_a.k = tab_b.k;

"""

SQL_JOIN_ABC = """

CREATE TABLE tab_a (

k BIGINT NOT NULL,

v BIGINT NOT NULL

) WITH ('materialized' = 'true');

CREATE TABLE tab_b (

k BIGINT NOT NULL,

w BIGINT NOT NULL

) WITH ('materialized' = 'true');

CREATE TABLE tab_c (

k BIGINT NOT NULL,

u BIGINT NOT NULL

) WITH ('materialized' = 'true');

CREATE MATERIALIZED VIEW join_ab AS

SELECT tab_a.k, tab_a.v, tab_b.w

FROM tab_a JOIN tab_b ON tab_a.k = tab_b.k;

CREATE MATERIALIZED VIEW join_ac AS

SELECT tab_a.k, tab_a.v, tab_c.u

FROM tab_a JOIN tab_c ON tab_a.k = tab_c.k;

"""

# Lower balancer thresholds so modest skew triggers a policy change in CI.

_ADAPTIVE_DEV_TWEAKS = {

"adaptive_joins": True,

}

def _fetch_circuit_json_profile(pipeline_name: str) -> dict:

resp = http_request(

"GET",

f"{API_PREFIX}/pipelines/{pipeline_name}/circuit_json_profile",

headers={"Accept-Encoding": "gzip"},

timeout=120,

)

assert resp.status_code == 200, (resp.status_code, resp.text[:500])

return resp.json()

def _balancer_policy_values(profile: dict) -> list[str]:

"""

Collect ``balancer_policy`` metric values from worker 0.

Values use the engine's ``Display`` spelling: ``shard``, ``broadcast``, ``balance``.

"""

workers = profile.get("worker_profiles") or []

assert workers, "circuit profile missing worker_profiles"

wp0 = workers[0]

meta = wp0.get("metadata") or {}

out: list[str] = []

for readings in meta.values():

if not isinstance(readings, list):

continue

for r in readings:

if r.get("metric_id") != "balancer_policy":

continue

val = r.get("value")

if isinstance(val, dict) and val.get("type") == "string":

out.append(str(val["value"]).lower())

return out

def _push_chunked_json(

pipeline, table: str, rows: list[dict], chunk_size: int = 4000

) -> None:

for i in range(0, len(rows), chunk_size):

pipeline.input_json(table, rows[i : i + chunk_size])

class TestAdaptiveJoins(unittest.TestCase):

def test_adaptive_joins_balancer_policies(self):

if FELDERA_TEST_NUM_WORKERS < 2:

self.skipTest("adaptive join balancing requires at least two workers")

name = unique_pipeline_name("test_adaptive_joins")

pipeline = PipelineBuilder(

TEST_CLIENT,

name=name,

sql=SQL_JOIN_AB,

runtime_config=RuntimeConfig(

workers=FELDERA_TEST_NUM_WORKERS,

hosts=FELDERA_TEST_NUM_HOSTS,

dev_tweaks=_ADAPTIVE_DEV_TWEAKS,

),

).create_or_replace()

try:

pipeline.start()

print(

"Single-join pipeline started; populating tables with balanced inputs"

)

# Uniform keys so both sides are ~evenly sharded (avoid k=0 used later for skew).

n_uniform = 4000

uniform_a = [{"k": i, "v": i} for i in range(1, n_uniform + 1)]

uniform_b = [{"k": i, "w": i} for i in range(1, n_uniform + 1)]

_push_chunked_json(pipeline, "tab_a", uniform_a)

_push_chunked_json(pipeline, "tab_b", uniform_b)

pipeline.wait_for_idle()

def all_sharded() -> bool:

pvals = _balancer_policy_values(_fetch_circuit_json_profile(name))

return len(pvals) > 0 and all(v == "shard" for v in pvals)

wait_for_condition(

"all balancer_policy values are shard after uniform load",

all_sharded,

timeout_s=120.0,

poll_interval_s=0.5,

)

# Heavy skew on a single join key to trigger (balance, broadcast).

print("Introducing heavy skew to trigger (balance, broadcast)")

skew_n = 25_000

skew_a = [{"k": 0, "v": i} for i in range(skew_n)]

_push_chunked_json(pipeline, "tab_a", skew_a)

pipeline.wait_for_idle()

pipeline.rebalance()

pipeline.wait_for_idle()

def skew_policies() -> bool:

pvals = _balancer_policy_values(_fetch_circuit_json_profile(name))

return "balance" in pvals and "broadcast" in pvals

wait_for_condition(

"skewed join uses balance on one stream and broadcast on another",

skew_policies,

timeout_s=180.0,

poll_interval_s=0.5,

)

print("Stopping pipeline to modify SQL")

pipeline.stop(force=False)

prev_ver = pipeline.program_version()

pipeline.modify(sql=SQL_JOIN_ABC)

wait_for_program_success(name, prev_ver + 1)

pipeline.start(bootstrap_policy=BootstrapPolicy.ALLOW)

print("Running pipeline with three-way join; populating the new table")

# Same join key cluster: third input should participate with policies compatible

# with the existing skewed cluster.

uniform_c = [{"k": i, "u": i} for i in range(1, n_uniform + 1)]

# skew_c = [{"k": 0, "u": i} for i in range(int(skew_n/2))]

_push_chunked_json(pipeline, "tab_c", uniform_c)

pipeline.rebalance()

def cluster_policies() -> bool:

pvals = _balancer_policy_values(_fetch_circuit_json_profile(name))

print("cluster_policies: ", pvals)

if not pvals:

return False

if "shard" in pvals:

return False

if "balance" not in pvals or "broadcast" not in pvals:

return False

# Extra probe input (tab_c) should stay on the broadcast side like tab_b.

return pvals.count("broadcast") == 2

wait_for_condition(

"three-way cluster: extra input stays broadcast with the other probe",

cluster_policies,

timeout_s=180.0,

poll_interval_s=0.5,

)

pipeline.stop(force=False)

pipeline.start()

print("Diluting skew with large uniform key ranges so sharding wins again.")

dilute_n = 80_000

base = 1_000_000

dilute_a = [{"k": base + i, "v": i} for i in range(dilute_n)]

dilute_b = [{"k": base + i, "w": i} for i in range(dilute_n)]

dilute_c = [{"k": base + i, "u": i} for i in range(dilute_n)]

_push_chunked_json(pipeline, "tab_a", dilute_a)

_push_chunked_json(pipeline, "tab_b", dilute_b)

_push_chunked_json(pipeline, "tab_c", dilute_c)

pipeline.rebalance()

wait_for_condition(

"all balancer_policy values return to shard after dilution",

all_sharded,

timeout_s=180.0,

poll_interval_s=0.5,

)

finally:

pipeline.stop(force=True)

pipeline.clear_storage()

if __name__ == "__main__":

unittest.main()