test_morsels.py

from __future__ import annotations

import io
import re

import pytest

import daft
from tests.conftest import get_tests_daft_runner_name
from tests.utils import clean_explain_output

pytestmark = pytest.mark.skipif(get_tests_daft_runner_name() != "native", reason="requires Native Runner to be in use")


def make_noop_udf(batch_size: int, dtype: daft.DataType = daft.DataType.int64()):
    @daft.udf(return_dtype=dtype, batch_size=batch_size, concurrency=1)
    def noop(x: dtype) -> dtype:
        return x

    return noop


# TODO: Add snapshot tests in Rust for the explain output of the following tests.


@pytest.mark.parametrize("dynamic_batching", [True, False])
def test_batch_size_from_udf_propagated_to_scan(dynamic_batching):
    with daft.execution_config_ctx(enable_dynamic_batching=dynamic_batching):
        df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
        df = df.select(make_noop_udf(10)(daft.col("a")))
        string_io = io.StringIO()
        df.explain(True, file=string_io)
        expected = """

    * UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
    |   Expr = py_udf(col(0: a)) as a
    |   Passthrough Columns = []
    |   Properties = { batch_size = 10, concurrency = 1, async = false, scalar = false }
    |   Resource request = None
    |   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
    |   Batch Size = 10
    |
    * InMemoryScan:
    |   Schema = a#Int64
    |   Size bytes = 40
    |   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
    |   Batch Size = 10

    """
        assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(
            expected
        )


def test_batch_size_from_udf_propagated_through_ops_to_scan():
    df = daft.from_pydict(
        {
            "data": [
                "https://www.google.com",
                "https://www.yahoo.com",
                "https://www.bing.com",
                "https://www.duckduckgo.com",
                "https://www.ask.com",
            ]
        }
    )
    df = df.select(daft.functions.decode_image(daft.functions.download(daft.col("data"))))
    df = df.select(make_noop_udf(10, daft.DataType.image())(daft.col("data")))
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    captured = string_io.getvalue().split("== Physical Plan ==")[-1]
    # Match the id inside col(0: id-...)
    m = re.search(r"col\(0: (id-[a-f0-9\-]+)\)", captured)
    id_placeholder = m.group(1) if m else ""
    expected = f"""

* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: __TruncateRootUDF_0-0-0__)) as data
|   Passthrough Columns = []
|   Properties = {{ batch_size = 10, concurrency = 1, async = false, scalar = false }}
|   Resource request = None
|   Batch Size = 10
|
* Project: col(0: __TruncateRootUDF_0-0-0__) as __TruncateRootUDF_0-0-0__
|   Batch Size = Range(0, 10]
|
* Project: image_decode(col(0: {id_placeholder}), lit("raise"), lit(PyObject(RGB))) as __TruncateRootUDF_0-0-0__, col(1: data)
|   Batch Size = Range(0, 10]
|
* Project: url_download(col(0: data), lit(true), lit("raise"), lit(32), lit(PyObject(IOConfig:
|   S3Config
|       region_name: None
|       endpoint_url: None
|       key_id: None
|       session_token: None
|       access_key: None
|       credentials_provider: None
|       buffer_time: None
|       max_connections: 32
|       retry_initial_backoff_ms: 1000
|       connect_timeout_ms: 30000
|       read_timeout_ms: 30000
|       num_tries: 25
|       retry_mode: Some("adaptive")
|       anonymous: false
|       use_ssl: true
|       verify_ssl: true
|       check_hostname_ssl: true
|       requester_pays: false
|       force_virtual_addressing: false
|       multipart_size: 8388608
|       multipart_max_concurrency: 100
|       custom_retry_msgs:[]
|   AzureConfig
|       storage_account: None
|       access_key: None
|       sas_token: None
|       bearer_token: None
|       tenant_id: None
|       client_id: None
|       client_secret: None
|       use_fabric_endpoint: false
|       anonymous: false
|       endpoint_url: None
|       use_ssl: true
|       max_connections_per_io_thread: 8
|   GCSConfig
|       project_id: None
|       anonymous: false
|       max_connections_per_io_thread: 8
|       retry_initial_backoff_ms: 1000
|       connect_timeout_ms: 30000
|       read_timeout_ms: 30000
|       num_tries: 5
|   TosConfig
|       region: None
|       endpoint: None
|       access_key: None
|       secret_key: ***
|       security_token: ***
|       anonymous: false
|       max_retries: 3
|       retry_timeout_ms: 30000
|       connect_timeout_ms: 10000
|       read_timeout_ms: 30000
|       max_concurrent_requests: 50
|       max_connections_per_io_thread: 50
|   CosConfig
|       region: None
|       endpoint: None
|       secret_id: None
|       secret_key: ***
|       security_token: ***
|       anonymous: false
|       max_retries: 3
|       retry_timeout_ms: 30000
|       connect_timeout_ms: 10000
|       read_timeout_ms: 30000
|       max_concurrent_requests: 50
|       max_connections_per_io_thread: 50
|   HTTPConfig
|   User agent = daft/0.0.1
|   Retry initial backoff ms = 1000
|   Connect timeout ms = 30000
|   Read timeout ms = 30000
|   Max retries = 5
|   UnityConfig
|       endpoint: None
|       token: None
|   GravitinoConfig
|       endpoint: None
|       metalake_name: None
|       auth_type:None
|       username: None
|       password: None
|       token: None
|   HuggingFaceConfig
|   Anonymous = false
|   ))) as {id_placeholder}, col(0: data)
|   Batch Size = Range(0, 10]
|
* InMemoryScan:
|   Schema = data#Utf8
|   Size bytes = 156
|   Stats = {{ Approx num rows = 5, Approx size bytes = 156 B, Accumulated selectivity = 1.00 }}
|   Batch Size = Range(0, 10]

"""
    assert clean_explain_output(captured) == clean_explain_output(expected)


def test_batch_size_from_multiple_udfs_do_not_override_each_other():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df = df.select(make_noop_udf(10)(daft.col("a")))
    df = df.select(make_noop_udf(20)(daft.col("a")))
    df = df.select(make_noop_udf(30)(daft.col("a")))
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    expected = """

* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: __TruncateRootUDF_0-0-0__)) as a
|   Passthrough Columns = []
|   Properties = { batch_size = 30, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = 30
|
* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: __TruncateRootUDF_1-0-0__)) as __TruncateRootUDF_0-0-0__
|   Passthrough Columns = []
|   Properties = { batch_size = 20, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = 20
|
* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: a)) as __TruncateRootUDF_1-0-0__
|   Passthrough Columns = []
|   Properties = { batch_size = 10, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = 10
|
* InMemoryScan:
|   Schema = a#Int64
|   Size bytes = 40
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = 10

"""
    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)


def test_batch_size_from_udf_not_propagated_through_agg():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df = df.groupby("a").sum()
    df = df.select(make_noop_udf(10)(daft.col("a")))
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    expected = """

* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: a)) as a
|   Passthrough Columns = []
|   Properties = { batch_size = 10, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 4, Approx size bytes = 32 B, Accumulated selectivity = 0.80 }
|   Batch Size = 10
|
* GroupedAggregate:
|   Group by: col(0: a)
|   Stats = { Approx num rows = 4, Approx size bytes = 32 B, Accumulated selectivity = 0.80 }
|
* InMemoryScan:
|   Schema = a#Int64
|   Size bytes = 40
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(0, 131072]

"""

    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)


def test_batch_size_from_udf_not_propagated_through_join():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df2 = daft.from_pydict({"b": [1, 2, 3, 4, 5]})
    df = df.join(df2, left_on="a", right_on="b")
    df = df.select(make_noop_udf(10)(daft.col("a")))
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    expected = """

* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: a)) as a
|   Passthrough Columns = []
|   Properties = { batch_size = 10, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 5, Approx size bytes = 37 B, Accumulated selectivity = 0.90 }
|   Batch Size = 10
|
* Project: col(0: a)
|   Stats = { Approx num rows = 5, Approx size bytes = 37 B, Accumulated selectivity = 0.90 }
|   Batch Size = Range(0, 10]
|
* HashJoin(Inner):
|   Build on left: true
|   Track Indices: true
|   Key Schema: a#Int64
|   Null equals Nulls = [false]
|   Stats = { Approx num rows = 5, Approx size bytes = 37 B, Accumulated selectivity = 0.90 }
|   Batch Size = Range(0, 10]
|\
| * Filter: not(is_null(col(0: b)))
| |   Stats = { Approx num rows = 5, Approx size bytes = 38 B, Accumulated selectivity = 0.95 }
| |   Batch Size = Range(0, 10]
| |
| * InMemoryScan:
| |   Schema = b#Int64
| |   Size bytes = 40
| |   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
| |   Batch Size = Range(0, 10]
|
| * Filter: not(is_null(col(0: a)))
| |   Stats = { Approx num rows = 5, Approx size bytes = 38 B, Accumulated selectivity = 0.95 }
| |   Batch Size = Range(0, 131072]
| |
| * InMemoryScan:
| |   Schema = a#Int64
| |   Size bytes = 40
| |   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
| |   Batch Size = Range(0, 131072]

"""
    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)


def test_batch_size_from_into_batches():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df = df.into_batches(10)
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    expected = """

* IntoBatches: 10
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(8, 10]
|
* InMemoryScan:
|   Schema = a#Int64
|   Size bytes = 40
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(8, 10]

"""
    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)


def test_batch_size_consecutive_into_batches():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df = df.into_batches(10)
    df = df.into_batches(20)
    df = df.into_batches(30)
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    expected = """

* IntoBatches: 30
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(24, 30]
|
* InMemoryScan:
|   Schema = a#Int64
|   Size bytes = 40
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(24, 30]

"""
    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)


def test_batch_size_from_into_batches_before_udf():
    df = daft.from_pydict({"a": [1, 2, 3, 4, 5]})
    df = df.into_batches(10)
    df = df.select(make_noop_udf(10)(daft.col("a")))
    string_io = io.StringIO()
    df.explain(True, file=string_io)
    print(string_io.getvalue())
    expected = """

* UDF tests.dataframe.test_morsels.make_noop_udf.<locals>.noop:
|   Expr = py_udf(col(0: a)) as a
|   Passthrough Columns = []
|   Properties = { batch_size = 10, concurrency = 1, async = false, scalar = false }
|   Resource request = None
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = 10
|
* IntoBatches: 10
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(8, 10]
|
* InMemoryScan:
|   Schema = a#Int64
|   Size bytes = 40
|   Stats = { Approx num rows = 5, Approx size bytes = 40 B, Accumulated selectivity = 1.00 }
|   Batch Size = Range(8, 10]

"""
    assert clean_explain_output(string_io.getvalue().split("== Physical Plan ==")[-1]) == clean_explain_output(expected)