[Data] Fix pandas memory calculation.

[Bye-legumes]

Why are these changes needed?

close #46785
Current the memory usage for pandas is not accurate when it's object, so we just implement to calculated it in recursion in case of nested.

Related issue number

closes #46785, closes #48506

Checks

• [√] I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• [√] I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • [√] Unit tests
  • Release tests
  • This PR is not tested :(

[Bye-legumes]

@bveeramani @c21 can you check this PR. I think the main idea is to calculate it in recursion.. I just did some test and for the example in #46785, each time the time is 20ms for calculation, I think it's acceptable?

[Bye-legumes]

Although it may not support tensor or other object.

[Bye-legumes]

hi, @bveeramani can you check this?

[richardliaw]

Hi @Bye-legumes , I can take over the review here.

I actually implemented a couple tests, but seems like this is failing here:

___________________________ test_size_bytes_bytes_object ____________________________

ray_start_regular_shared = RayContext(dashboard_url='127.0.0.1:8265', python_version='3.10.15', ray_version='3.0.0.dev0', ray_commit='9e5c5bfba9d79d1d134719cf202c7150407acbc6')

    def test_size_bytes_bytes_object(ray_start_regular_shared):
        def generate_data(batch):
            for _ in range(8):
                yield {"data": [[b"\x00" * 128 * 1024 * 128]]}
    
        ds = (
            ray.data.range(1, override_num_blocks=1)
            .map_batches(generate_data, batch_size=1)
            .map_batches(lambda batch: batch, batch_format="pandas")
        )
    
        true_value = 128 * 1024 * 128 * 8
        for bundle in ds.iter_internal_ref_bundles():
            size = bundle.size_bytes()
            # assert that true_value is within 10% of bundle.size_bytes()
>           assert true_value * 0.9 <= size <= true_value * 1.1, (true_value, size)
E           AssertionError: (134217728, 192)
E           assert (134217728 * 0.9) <= 192

[richardliaw]

A couple tests I wrote --

def test_size_bytes_small(ray_start_regular_shared):
    animals = ["Flamingo", "Centipede"]
    block = pd.DataFrame({"animals": animals})
    block["animals"] = block["animals"].astype("string")

    block_accessor = PandasBlockAccessor.for_block(block)
    bytes_size = block_accessor.size_bytes()

    # generally strings are hard, so let's use what Pandas gives us.
    # get memory usage from pandas
    memory_usage = block.memory_usage(index=True, deep=True).sum()
    # check that memory usage is within 10% of the size_bytes
    assert memory_usage * 0.9 <= bytes_size <= memory_usage * 1.1, (
        bytes_size,
        memory_usage,
    )


def test_size_bytes_large_str(ray_start_regular_shared):
    animals = [
        random.choice(["alligator", "crocodile", "centipede", "flamingo"])
        for i in range(100_000)
    ]
    block = pd.DataFrame({"animals": animals})
    block["animals"] = block["animals"].astype("string")

    block_accessor = PandasBlockAccessor.for_block(block)
    bytes_size = block_accessor.size_bytes()

    # String disk usage is wildly different from in-process memory usage
    memory_usage = block.memory_usage(index=True, deep=True).sum()
    # check that memory usage is within 10% of the size_bytes
    assert memory_usage * 0.9 <= bytes_size <= memory_usage * 1.1, (
        bytes_size,
        memory_usage,
    )


def test_size_bytes_large_floats(ray_start_regular_shared):
    animals = [random.random() for i in range(100_000)]
    block = pd.DataFrame({"animals": animals})

    block_accessor = PandasBlockAccessor.for_block(block)
    bytes_size = block_accessor.size_bytes()

    memory_usage = pickle.dumps(block).__sizeof__()
    # check that memory usage is within 10% of the size_bytes
    assert memory_usage * 0.9 <= bytes_size <= memory_usage * 1.1, (
        bytes_size,
        memory_usage,
    )


def test_size_bytes_bytes_object(ray_start_regular_shared):
    def generate_data(batch):
        for _ in range(8):
            yield {"data": [[b"\x00" * 128 * 1024 * 128]]}

    ds = (
        ray.data.range(1, override_num_blocks=1)
        .map_batches(generate_data, batch_size=1)
        .map_batches(lambda batch: batch, batch_format="pandas")
    )

    true_value = 128 * 1024 * 128 * 8
    for bundle in ds.iter_internal_ref_bundles():
        size = bundle.size_bytes()
        # assert that true_value is within 10% of bundle.size_bytes()
        assert true_value * 0.9 <= size <= true_value * 1.1, (true_value, size)


def test_size_bytes_unowned_numpy(ray_start_regular_shared):
    import numpy as np

    df = pd.DataFrame(
        {
            "data": [
                np.random.randint(size=1024, low=0, high=100, dtype=np.int8)
                for _ in range(1_000)
            ],
        }
    )

    block_accessor = PandasBlockAccessor.for_block(df)
    block_size = block_accessor.size_bytes()
    true_value = 1024 * 1000
    assert true_value * 0.9 <= block_size <= true_value * 1.1

[Bye-legumes]

Hi @Bye-legumes , I can take over the review here.

I actually implemented a couple tests, but seems like this is failing here:

___________________________ test_size_bytes_bytes_object ____________________________

ray_start_regular_shared = RayContext(dashboard_url='127.0.0.1:8265', python_version='3.10.15', ray_version='3.0.0.dev0', ray_commit='9e5c5bfba9d79d1d134719cf202c7150407acbc6')

    def test_size_bytes_bytes_object(ray_start_regular_shared):
        def generate_data(batch):
            for _ in range(8):
                yield {"data": [[b"\x00" * 128 * 1024 * 128]]}
    
        ds = (
            ray.data.range(1, override_num_blocks=1)
            .map_batches(generate_data, batch_size=1)
            .map_batches(lambda batch: batch, batch_format="pandas")
        )
    
        true_value = 128 * 1024 * 128 * 8
        for bundle in ds.iter_internal_ref_bundles():
            size = bundle.size_bytes()
            # assert that true_value is within 10% of bundle.size_bytes()
>           assert true_value * 0.9 <= size <= true_value * 1.1, (true_value, size)
E           AssertionError: (134217728, 192)
E           assert (134217728 * 0.9) <= 192

Thanks so much! I think this time it's OK now!

[Bye-legumes]

I updated the tests for readability and also made it a bit more extensive -- I found some areas that broke, could you take a look?

Hi, I just fixed the issues that you mentioned! Thx!

[richardliaw]

This looks great to me! Will ping others once tests pass.

[Bye-legumes]

This looks great to me! Will ping others once tests pass.

Hi! I think it's OK now!

[raulchen]

nit, can we directly compare the dtype without casting it to strings?

[Bye-legumes]

Just fixed now!

[richardliaw]

Some tests failing?

[Bye-legumes]

Some tests failing?

Yeah, just fix as the some of the inner blocks are TensorArrayElement, TensorDtype now. I just make some changes and also I can use type check now. In the previous it's just "python_object()" for extension array. And recent 2 days the type changed to nd.array(object)

[Bye-legumes]

Some tests failing?

All the test are OK now.

[richardliaw]

Awesome, thank you so much @Bye-legumes !!!!

[aslonnie]

This PR is reverted. many data release tests are broken.