test_metrics.py 文件源码

def test_ndcg_doesnt_completely_fail(spark_context, hive_context):
    "Mediocre test that just looks for a happy path"
    df = spark_context.parallelize([
        [4, 0, 'foo'],
        [3, 1, 'foo'],
        [0, 2, 'foo'],
        [3, 3, 'foo'],
    ]).toDF(['label', 'hit_position', 'query'])

    # Top 2 are in perfect order. Also this indirectly tests that
    # k is really top 2, and not somehow top 3 or some such
    ndcg_at_2 = mjolnir.metrics.ndcg(df, 2, query_cols=['query'])
    assert 1.0 == ndcg_at_2

    # Top 4 are slightly out. This value was checked by also
    # calculating by hand.
    ndcg_at_4 = mjolnir.metrics.ndcg(df, 4, query_cols=['query'])
    assert 0.9788 == pytest.approx(ndcg_at_4, abs=0.0001)