mlflow--mlflow/tests/utils/test_search_utils.py

import base64
import json
import re

import pytest

from mlflow.entities import (
    Dataset,
    DatasetInput,
    InputTag,
    LifecycleStage,
    Metric,
    Param,
    Run,
    RunData,
    RunInfo,
    RunInputs,
    RunStatus,
    RunTag,
    TraceState,
    trace_location,
)
from mlflow.entities.trace_info import TraceInfo
from mlflow.exceptions import MlflowException
from mlflow.utils.mlflow_tags import MLFLOW_DATASET_CONTEXT
from mlflow.utils.search_utils import SearchTraceUtils, SearchUtils


@pytest.mark.parametrize(
    ("filter_string", "parsed_filter"),
    [
        (
            "metric.acc >= 0.94",
            [{"comparator": ">=", "key": "acc", "type": "metric", "value": "0.94"}],
        ),
        ("metric.acc>=100", [{"comparator": ">=", "key": "acc", "type": "metric", "value": "100"}]),
        ("params.m!='tf'", [{"comparator": "!=", "key": "m", "type": "parameter", "value": "tf"}]),
        (
            'params."m"!="tf"',
            [{"comparator": "!=", "key": "m", "type": "parameter", "value": "tf"}],
        ),
        (
            'metric."legit name" >= 0.243',
            [{"comparator": ">=", "key": "legit name", "type": "metric", "value": "0.243"}],
        ),
        ("metrics.XYZ = 3", [{"comparator": "=", "key": "XYZ", "type": "metric", "value": "3"}]),
        (
            'params."cat dog" = "pets"',
            [{"comparator": "=", "key": "cat dog", "type": "parameter", "value": "pets"}],
        ),
        (
            'metrics."X-Y-Z" = 3',
            [{"comparator": "=", "key": "X-Y-Z", "type": "metric", "value": "3"}],
        ),
        (
            'metrics."X//Y#$$@&Z" = 3',
            [{"comparator": "=", "key": "X//Y#$$@&Z", "type": "metric", "value": "3"}],
        ),
        (
            "params.model = 'LinearRegression'",
            [{"comparator": "=", "key": "model", "type": "parameter", "value": "LinearRegression"}],
        ),
        (
            "metrics.rmse < 1 and params.model_class = 'LR'",
            [
                {"comparator": "<", "key": "rmse", "type": "metric", "value": "1"},
                {"comparator": "=", "key": "model_class", "type": "parameter", "value": "LR"},
            ],
        ),
        ("", []),
        ("`metric`.a >= 0.1", [{"comparator": ">=", "key": "a", "type": "metric", "value": "0.1"}]),
        (
            "`params`.model >= 'LR'",
            [{"comparator": ">=", "key": "model", "type": "parameter", "value": "LR"}],
        ),
        (
            "tags.version = 'commit-hash'",
            [{"comparator": "=", "key": "version", "type": "tag", "value": "commit-hash"}],
        ),
        (
            "`tags`.source_name = 'a notebook'",
            [{"comparator": "=", "key": "source_name", "type": "tag", "value": "a notebook"}],
        ),
        (
            'metrics."accuracy.2.0" > 5',
            [{"comparator": ">", "key": "accuracy.2.0", "type": "metric", "value": "5"}],
        ),
        (
            "metrics.`spacey name` > 5",
            [{"comparator": ">", "key": "spacey name", "type": "metric", "value": "5"}],
        ),
        (
            'params."p.a.r.a.m" != "a"',
            [{"comparator": "!=", "key": "p.a.r.a.m", "type": "parameter", "value": "a"}],
        ),
        ('tags."t.a.g" = "a"', [{"comparator": "=", "key": "t.a.g", "type": "tag", "value": "a"}]),
        (
            "attribute.artifact_uri = '1/23/4'",
            [{"type": "attribute", "comparator": "=", "key": "artifact_uri", "value": "1/23/4"}],
        ),
        (
            "attribute.start_time >= 1234",
            [{"type": "attribute", "comparator": ">=", "key": "start_time", "value": "1234"}],
        ),
        (
            "run.status = 'RUNNING'",
            [{"type": "attribute", "comparator": "=", "key": "status", "value": "RUNNING"}],
        ),
        (
            "dataset.name = 'my_dataset'",
            [{"type": "dataset", "comparator": "=", "key": "name", "value": "my_dataset"}],
        ),
        (
            "tags.version IS NULL",
            [{"comparator": "IS NULL", "key": "version", "type": "tag", "value": None}],
        ),
        (
            "tags.version IS NOT NULL",
            [{"comparator": "IS NOT NULL", "key": "version", "type": "tag", "value": None}],
        ),
        (
            "params.lr IS NULL",
            [{"comparator": "IS NULL", "key": "lr", "type": "parameter", "value": None}],
        ),
        (
            "params.lr IS NOT NULL",
            [{"comparator": "IS NOT NULL", "key": "lr", "type": "parameter", "value": None}],
        ),
        (
            "tags.a IS NULL AND params.b = 'val'",
            [
                {"comparator": "IS NULL", "key": "a", "type": "tag", "value": None},
                {"comparator": "=", "key": "b", "type": "parameter", "value": "val"},
            ],
        ),
    ],
)
def test_filter(filter_string, parsed_filter):
    assert SearchUtils.parse_search_filter(filter_string) == parsed_filter


@pytest.mark.parametrize(
    ("filter_string", "parsed_filter"),
    [
        ("params.m = 'LR'", [{"type": "parameter", "comparator": "=", "key": "m", "value": "LR"}]),
        ('params.m = "LR"', [{"type": "parameter", "comparator": "=", "key": "m", "value": "LR"}]),
        (
            'params.m = "L\'Hosp"',
            [{"type": "parameter", "comparator": "=", "key": "m", "value": "L'Hosp"}],
        ),
    ],
)
def test_correct_quote_trimming(filter_string, parsed_filter):
    assert SearchUtils.parse_search_filter(filter_string) == parsed_filter


@pytest.mark.parametrize(
    ("filter_string", "error_message"),
    [
        ("metric.acc >= 0.94; metrics.rmse < 1", "Search filter contained multiple expression"),
        ("m.acc >= 0.94", "Invalid entity type"),
        ("acc >= 0.94", "Invalid attribute key"),
        ("p.model >= 'LR'", "Invalid entity type"),
        ("attri.x != 1", "Invalid entity type"),
        ("a.x != 1", "Invalid entity type"),
        ("model >= 'LR'", "Invalid attribute key"),
        ("metrics.A > 0.1 OR params.B = 'LR'", "Invalid clause(s) in filter string"),
        ("metrics.A > 0.1 NAND params.B = 'LR'", "Invalid clause(s) in filter string"),
        ("metrics.A > 0.1 AND (params.B = 'LR')", "Invalid clause(s) in filter string"),
        ("`metrics.A > 0.1", "Invalid clause(s) in filter string"),
        ("param`.A > 0.1", "Invalid clause(s) in filter string"),
        ("`dummy.A > 0.1", "Invalid clause(s) in filter string"),
        ("dummy`.A > 0.1", "Invalid clause(s) in filter string"),
        ("attribute.start != 1", "Invalid attribute key"),
        ("attribute.experiment_id != 1", "Invalid attribute key"),
        ("attribute.lifecycle_stage = 'ACTIVE'", "Invalid attribute key"),
        ("attribute.name != 1", "Invalid attribute key"),
        ("attribute.time != 1", "Invalid attribute key"),
        ("attribute._status != 'RUNNING'", "Invalid attribute key"),
        ("attribute.status = true", "Invalid clause(s) in filter string"),
        ("dataset.status = 'true'", "Invalid dataset key"),
        ("dataset.profile = 'num_rows: 10'", "Invalid dataset key"),
        ("metrics.acc IS NULL", "IS NULL / IS NOT NULL is only supported for tags and params"),
        ("attribute.status IS NULL", "IS NULL / IS NOT NULL is only supported for tags and params"),
    ],
)
def test_error_filter(filter_string, error_message):
    with pytest.raises(MlflowException, match=re.escape(error_message)):
        SearchUtils.parse_search_filter(filter_string)


@pytest.mark.parametrize(
    ("filter_string", "error_message"),
    [
        ("metric.model = 'LR'", "Expected numeric value type for metric"),
        ("metric.model = '5'", "Expected numeric value type for metric"),
        ("params.acc = 5", "Expected a quoted string value for param"),
        ("tags.acc = 5", "Expected a quoted string value for tag"),
        ("metrics.acc != metrics.acc", "Expected numeric value type for metric"),
        ("1.0 > metrics.acc", "Expected 'Identifier' found"),
        ("attribute.status = 1", "Expected a quoted string value for attributes"),
    ],
)
def test_error_comparison_clauses(filter_string, error_message):
    with pytest.raises(MlflowException, match=error_message):
        SearchUtils.parse_search_filter(filter_string)


@pytest.mark.parametrize(
    ("filter_string", "error_message"),
    [
        ("params.acc = LR", "value is either not quoted or unidentified quote types"),
        ("tags.acc = LR", "value is either not quoted or unidentified quote types"),
        ("params.acc = `LR`", "value is either not quoted or unidentified quote types"),
        ("params.'acc = LR", "Invalid clause(s) in filter string"),
        ("params.acc = 'LR", "Invalid clause(s) in filter string"),
        ("params.acc = LR'", "Invalid clause(s) in filter string"),
        ("params.acc = \"LR'", "Invalid clause(s) in filter string"),
        ("tags.acc = \"LR'", "Invalid clause(s) in filter string"),
        ("tags.acc = = 'LR'", "Invalid clause(s) in filter string"),
        ("attribute.status IS 'RUNNING'", "Invalid clause(s) in filter string"),
    ],
)
def test_bad_quotes(filter_string, error_message):
    with pytest.raises(MlflowException, match=re.escape(error_message)):
        SearchUtils.parse_search_filter(filter_string)


@pytest.mark.parametrize(
    ("filter_string", "error_message"),
    [
        ("params.acc LR !=", "Invalid clause(s) in filter string"),
        ("params.acc LR", "Invalid clause(s) in filter string"),
        ("metric.acc !=", "Invalid clause(s) in filter string"),
        ("acc != 1.0", "Invalid attribute key"),
        ("foo is null", "Invalid attribute key"),
        ("1=1", "Expected 'Identifier' found"),
        ("1==2", "Expected 'Identifier' found"),
    ],
)
def test_invalid_clauses(filter_string, error_message):
    with pytest.raises(MlflowException, match=re.escape(error_message)):
        SearchUtils.parse_search_filter(filter_string)


@pytest.mark.parametrize(
    ("entity_type", "bad_comparators", "key", "entity_value"),
    [
        ("metrics", ["~", "~="], "abc", 1.0),
        ("params", [">", "<", ">=", "<=", "~"], "abc", "'my-param-value'"),
        ("tags", [">", "<", ">=", "<=", "~"], "abc", "'my-tag-value'"),
        ("attributes", [">", "<", ">=", "<=", "~"], "status", "'my-tag-value'"),
        ("attributes", ["LIKE", "ILIKE"], "start_time", 1234),
        ("datasets", [">", "<", ">=", "<=", "~"], "name", "'my-dataset-name'"),
    ],
)
def test_bad_comparators(entity_type, bad_comparators, key, entity_value):
    run = Run(
        run_info=RunInfo(
            run_id="hi",
            experiment_id=0,
            user_id="user-id",
            status=RunStatus.to_string(RunStatus.FAILED),
            start_time=0,
            end_time=1,
            lifecycle_stage=LifecycleStage.ACTIVE,
        ),
        run_data=RunData(metrics=[], params=[], tags=[]),
    )
    for bad_comparator in bad_comparators:
        bad_filter = f"{entity_type}.{key} {bad_comparator} {entity_value}"
        with pytest.raises(MlflowException, match="Invalid comparator"):
            SearchUtils.filter([run], bad_filter)


@pytest.mark.parametrize(
    ("filter_string", "matching_runs"),
    [
        (None, [0, 1, 2]),
        ("", [0, 1, 2]),
        ("attributes.status = 'FAILED'", [0, 2]),
        ("metrics.key1 = 123", [1]),
        ("metrics.key1 != 123", [0, 2]),
        ("metrics.key1 >= 123", [1, 2]),
        ("params.my_param = 'A'", [0, 1]),
        ("tags.tag1 = 'D'", [2]),
        ("tags.tag1 != 'D'", [1]),
        ("params.my_param = 'A' AND attributes.status = 'FAILED'", [0]),
        ("datasets.name = 'name1'", [0, 1]),
        ("datasets.name IN ('name1', 'name2')", [0, 1, 2]),
        ("datasets.digest IN ('digest1', 'digest2')", [0, 1, 2]),
        ("datasets.name = 'name1' AND datasets.digest = 'digest2'", []),
        ("datasets.context = 'train'", [0]),
        ("datasets.name = 'name1' AND datasets.context = 'train'", [0]),
    ],
)
def test_correct_filtering(filter_string, matching_runs):
    runs = [
        Run(
            run_info=RunInfo(
                run_id="hi",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 121, 1, 0)], params=[Param("my_param", "A")], tags=[]
            ),
            run_inputs=RunInputs(
                dataset_inputs=[
                    DatasetInput(
                        dataset=Dataset(
                            name="name1",
                            digest="digest1",
                            source_type="my_source_type",
                            source="source",
                        ),
                        tags=[InputTag(MLFLOW_DATASET_CONTEXT, "train")],
                    )
                ]
            ),
        ),
        Run(
            run_info=RunInfo(
                run_id="hi2",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 123, 1, 0)],
                params=[Param("my_param", "A")],
                tags=[RunTag("tag1", "C")],
            ),
            run_inputs=RunInputs(
                dataset_inputs=[
                    DatasetInput(
                        dataset=Dataset(
                            name="name1",
                            digest="digest1",
                            source_type="my_source_type",
                            source="source",
                        ),
                        tags=[],
                    )
                ]
            ),
        ),
        Run(
            run_info=RunInfo(
                run_id="hi3",
                experiment_id=1,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 125, 1, 0)],
                params=[Param("my_param", "B")],
                tags=[RunTag("tag1", "D")],
            ),
            run_inputs=RunInputs(
                dataset_inputs=[
                    DatasetInput(
                        dataset=Dataset(
                            name="name2",
                            digest="digest2",
                            source_type="my_source_type",
                            source="source",
                        ),
                        tags=[],
                    )
                ]
            ),
        ),
    ]
    filtered_runs = SearchUtils.filter(runs, filter_string)
    assert set(filtered_runs) == {runs[i] for i in matching_runs}


def test_filter_runs_by_start_time():
    runs = [
        Run(
            run_info=RunInfo(
                run_id=run_id,
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=idx,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(),
        )
        for idx, run_id in enumerate(["a", "b", "c"])
    ]
    assert SearchUtils.filter(runs, "attribute.start_time >= 0") == runs
    assert SearchUtils.filter(runs, "attribute.start_time > 1") == runs[2:]
    assert SearchUtils.filter(runs, "attribute.start_time = 2") == runs[2:]


def test_filter_runs_by_user_id():
    runs = [
        Run(
            run_info=RunInfo(
                run_id="a",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=1,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(),
        ),
        Run(
            run_info=RunInfo(
                run_id="b",
                experiment_id=0,
                user_id="user-id2",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=1,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(),
        ),
    ]
    assert SearchUtils.filter(runs, "attribute.user_id = 'user-id2'")[0] == runs[1]


def test_filter_runs_by_end_time():
    runs = [
        Run(
            run_info=RunInfo(
                run_id=run_id,
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=idx,
                end_time=idx,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(),
        )
        for idx, run_id in enumerate(["a", "b", "c"])
    ]
    assert SearchUtils.filter(runs, "attribute.end_time >= 0") == runs
    assert SearchUtils.filter(runs, "attribute.end_time > 1") == runs[2:]
    assert SearchUtils.filter(runs, "attribute.end_time = 2") == runs[2:]


@pytest.mark.parametrize(
    ("order_bys", "matching_runs"),
    [
        (None, [2, 1, 0]),
        ([], [2, 1, 0]),
        (["tags.noSuchTag"], [2, 1, 0]),
        (["attributes.status"], [2, 0, 1]),
        (["attributes.start_time"], [0, 2, 1]),
        (["metrics.key1 asc"], [0, 1, 2]),
        (['metrics."key1"  desc'], [2, 1, 0]),
        (["params.my_param"], [1, 0, 2]),
        (["params.my_param aSc", "attributes.status  ASC"], [0, 1, 2]),
        (["params.my_param", "attributes.status DESC"], [1, 0, 2]),
        (["params.my_param DESC", "attributes.status   DESC"], [2, 1, 0]),
        (["params.`my_param` DESC", "attributes.status DESC"], [2, 1, 0]),
        (["tags.tag1"], [1, 2, 0]),
        (["tags.tag1    DESC"], [2, 1, 0]),
    ],
)
def test_correct_sorting(order_bys, matching_runs):
    runs = [
        Run(
            run_info=RunInfo(
                run_id="9",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 121, 1, 0)], params=[Param("my_param", "A")], tags=[]
            ),
        ),
        Run(
            run_info=RunInfo(
                run_id="8",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=1,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 123, 1, 0)],
                params=[Param("my_param", "A")],
                tags=[RunTag("tag1", "C")],
            ),
        ),
        Run(
            run_info=RunInfo(
                run_id="7",
                experiment_id=1,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=1,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(
                metrics=[Metric("key1", 125, 1, 0)],
                params=[Param("my_param", "B")],
                tags=[RunTag("tag1", "D")],
            ),
        ),
    ]
    sorted_runs = SearchUtils.sort(runs, order_bys)
    sorted_run_indices = []
    for run in sorted_runs:
        for i, r in enumerate(runs):
            if r == run:
                sorted_run_indices.append(i)
                break
    assert sorted_run_indices == matching_runs


def test_order_by_metric_with_nans_infs_nones():
    metric_vals_str = ["nan", "inf", "-inf", "-1000", "0", "1000", "None"]
    runs = [
        Run(
            run_info=RunInfo(
                run_id=x,
                experiment_id=0,
                user_id="user",
                status=RunStatus.to_string(RunStatus.FINISHED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData(metrics=[Metric("x", None if x == "None" else float(x), 1, 0)]),
        )
        for x in metric_vals_str
    ]
    sorted_runs_asc = [x.info.run_id for x in SearchUtils.sort(runs, ["metrics.x asc"])]
    sorted_runs_desc = [x.info.run_id for x in SearchUtils.sort(runs, ["metrics.x desc"])]
    # asc
    assert sorted_runs_asc == ["-inf", "-1000", "0", "1000", "inf", "nan", "None"]
    # desc
    assert sorted_runs_desc == ["inf", "1000", "0", "-1000", "-inf", "nan", "None"]


@pytest.mark.parametrize(
    ("order_by", "error_message"),
    [
        ("m.acc", "Invalid entity type"),
        ("acc", "Invalid attribute key"),
        ("attri.x", "Invalid entity type"),
        ("`metrics.A", "Invalid order_by clause"),
        ("`metrics.A`", "Invalid entity type"),
        ("attribute.start", "Invalid attribute key"),
        ("attribute.experiment_id", "Invalid attribute key"),
        ("metrics.A != 1", "Invalid order_by clause"),
        ("params.my_param ", "Invalid order_by clause"),
        ("attribute.run_id ACS", "Invalid ordering key"),
        ("attribute.run_id decs", "Invalid ordering key"),
    ],
)
def test_invalid_order_by_search_runs(order_by, error_message):
    with pytest.raises(MlflowException, match=error_message):
        SearchUtils.parse_order_by_for_search_runs(order_by)


@pytest.mark.parametrize(
    ("order_by", "ascending_expected"),
    [
        ("metrics.`Mean Square Error`", True),
        ("metrics.`Mean Square Error` ASC", True),
        ("metrics.`Mean Square Error` DESC", False),
    ],
)
def test_space_order_by_search_runs(order_by, ascending_expected):
    identifier_type, identifier_name, ascending = SearchUtils.parse_order_by_for_search_runs(
        order_by
    )
    assert identifier_type == "metric"
    assert identifier_name == "Mean Square Error"
    assert ascending == ascending_expected


@pytest.mark.parametrize(
    ("order_by", "error_message"),
    [
        ("creation_timestamp DESC", "Invalid order by key"),
        ("last_updated_timestamp DESC blah", "Invalid order_by clause"),
        ("", "Invalid order_by clause"),
        ("timestamp somerandomstuff ASC", "Invalid order_by clause"),
        ("timestamp somerandomstuff", "Invalid order_by clause"),
        ("timestamp decs", "Invalid order_by clause"),
        ("timestamp ACS", "Invalid order_by clause"),
        ("name aCs", "Invalid ordering key"),
    ],
)
def test_invalid_order_by_search_registered_models(order_by, error_message):
    with pytest.raises(MlflowException, match=re.escape(error_message)):
        SearchUtils.parse_order_by_for_search_registered_models(order_by)


@pytest.mark.parametrize(
    ("page_token", "max_results", "matching_runs", "expected_next_page_token"),
    [
        (None, 1, [0], {"offset": 1}),
        (None, 2, [0, 1], {"offset": 2}),
        (None, 3, [0, 1, 2], None),
        (None, 5, [0, 1, 2], None),
        ({"offset": 1}, 1, [1], {"offset": 2}),
        ({"offset": 1}, 2, [1, 2], None),
        ({"offset": 1}, 3, [1, 2], None),
        ({"offset": 2}, 1, [2], None),
        ({"offset": 2}, 2, [2], None),
        ({"offset": 2}, 0, [], {"offset": 2}),
        ({"offset": 3}, 1, [], None),
    ],
)
def test_pagination(page_token, max_results, matching_runs, expected_next_page_token):
    runs = [
        Run(
            run_info=RunInfo(
                run_id="0",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData([], [], []),
        ),
        Run(
            run_info=RunInfo(
                run_id="1",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData([], [], []),
        ),
        Run(
            run_info=RunInfo(
                run_id="2",
                experiment_id=0,
                user_id="user-id",
                status=RunStatus.to_string(RunStatus.FAILED),
                start_time=0,
                end_time=1,
                lifecycle_stage=LifecycleStage.ACTIVE,
            ),
            run_data=RunData([], [], []),
        ),
    ]
    encoded_page_token = None
    if page_token:
        encoded_page_token = base64.b64encode(json.dumps(page_token).encode("utf-8"))
    paginated_runs, next_page_token = SearchUtils.paginate(runs, encoded_page_token, max_results)

    paginated_run_indices = []
    for run in paginated_runs:
        for i, r in enumerate(runs):
            if r == run:
                paginated_run_indices.append(i)
                break
    assert paginated_run_indices == matching_runs

    decoded_next_page_token = None
    if next_page_token:
        decoded_next_page_token = json.loads(base64.b64decode(next_page_token))
    assert decoded_next_page_token == expected_next_page_token


@pytest.mark.parametrize(
    ("page_token", "error_message"),
    [
        (base64.b64encode(json.dumps({}).encode("utf-8")), "Invalid page token"),
        (base64.b64encode(json.dumps({"offset": "a"}).encode("utf-8")), "Invalid page token"),
        (base64.b64encode(json.dumps({"offsoot": 7}).encode("utf-8")), "Invalid page token"),
        (base64.b64encode(b"not json"), "Invalid page token"),
        ("not base64", "Invalid page token"),
    ],
)
def test_invalid_page_tokens(page_token, error_message):
    with pytest.raises(MlflowException, match=error_message):
        SearchUtils.paginate([], page_token, 1)


def test_like_pattern_with_plus_character():
    import mlflow

    name = "jamie-foo C+W bar"
    mlflow.create_experiment(name)

    exps = mlflow.search_experiments(filter_string=f'name LIKE "{name}"')
    assert len(exps) == 1

    exps = mlflow.search_experiments(filter_string='name LIKE "jamie-foo C+%"')
    assert len(exps) == 1


def test_filter_runs_by_tag_and_param_is_null():
    run_with_tag = Run(
        run_info=RunInfo(
            run_id="run1",
            experiment_id=0,
            user_id="user",
            status=RunStatus.to_string(RunStatus.FINISHED),
            start_time=0,
            end_time=1,
            lifecycle_stage=LifecycleStage.ACTIVE,
        ),
        run_data=RunData(tags=[RunTag("env", "prod")], params=[], metrics=[]),
    )
    run_with_param = Run(
        run_info=RunInfo(
            run_id="run2",
            experiment_id=0,
            user_id="user",
            status=RunStatus.to_string(RunStatus.FINISHED),
            start_time=0,
            end_time=1,
            lifecycle_stage=LifecycleStage.ACTIVE,
        ),
        run_data=RunData(tags=[], params=[Param("lr", "0.01")], metrics=[]),
    )
    runs = [run_with_tag, run_with_param]

    assert [r.info.run_id for r in SearchUtils.filter(runs, "tags.env IS NOT NULL")] == ["run1"]
    assert [r.info.run_id for r in SearchUtils.filter(runs, "tags.env IS NULL")] == ["run2"]
    assert [r.info.run_id for r in SearchUtils.filter(runs, "params.lr IS NOT NULL")] == ["run2"]
    assert [r.info.run_id for r in SearchUtils.filter(runs, "params.lr IS NULL")] == ["run1"]


def test_search_trace_utils_filter_tag_is_null():
    loc = trace_location.TraceLocation.from_experiment_id("0")
    trace1 = TraceInfo(
        trace_id="t1",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        tags={"env": "prod", "region": "us"},
    )
    trace2 = TraceInfo(
        trace_id="t2",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        tags={"env": "staging"},
    )
    trace3 = TraceInfo(
        trace_id="t3",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        tags={},
    )
    traces = [trace1, trace2, trace3]

    result = SearchTraceUtils.filter(traces, "tag.region IS NULL")
    assert {t.trace_id for t in result} == {"t2", "t3"}

    result = SearchTraceUtils.filter(traces, "tag.region IS NOT NULL")
    assert {t.trace_id for t in result} == {"t1"}

    result = SearchTraceUtils.filter(traces, "tag.env IS NULL")
    assert {t.trace_id for t in result} == {"t3"}

    result = SearchTraceUtils.filter(traces, "tag.env IS NOT NULL")
    assert {t.trace_id for t in result} == {"t1", "t2"}

    result = SearchTraceUtils.filter(traces, 'tag.region IS NULL AND tag.env = "staging"')
    assert {t.trace_id for t in result} == {"t2"}


def test_search_trace_utils_filter_metadata_is_null():
    loc = trace_location.TraceLocation.from_experiment_id("0")
    trace1 = TraceInfo(
        trace_id="t1",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        trace_metadata={"user": "alice", "session": "s1"},
    )
    trace2 = TraceInfo(
        trace_id="t2",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        trace_metadata={"user": "bob"},
    )
    trace3 = TraceInfo(
        trace_id="t3",
        trace_location=loc,
        request_time=0,
        state=TraceState.OK,
        trace_metadata={},
    )
    traces = [trace1, trace2, trace3]

    result = SearchTraceUtils.filter(traces, "metadata.session IS NULL")
    assert {t.trace_id for t in result} == {"t2", "t3"}

    result = SearchTraceUtils.filter(traces, "metadata.session IS NOT NULL")
    assert {t.trace_id for t in result} == {"t1"}