Returning daft dataframe with ray.get() results in TypeError: cannot pickle 'builtins.LogicalPlanBuilder' object

[snowch]

Here is a minimal example:

import sys
import traceback
from typing import Any, List, Optional
import pyarrow as pa
import ray
import daft

RAY_ADDRESS = "ray://10.143.11.241:30001"

if ray.is_initialized():
    ray.shutdown()

ray.init(
    address=RAY_ADDRESS,
    runtime_env={"pip": ["daft", "pyarrow",]},
)
daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)

@ray.remote
def create_mock_arrow_data():
    daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)
    try:
        batch_data = {
            "pulocationid": pa.array([100, 101, 102, 103, 104, 105, 106, 107, 108, 109]),
            "dolocationid": pa.array([200, 201, 202, 203, 204, 205, 206, 207, 208, 209]),
            "passenger_count": pa.array([1, 2, 3, 1, 2, 4, 1, 1, 2, 3]),
            "trip_distance": pa.array([1.5, 2.3, 3.1, 4.7, 1.2, 2.8, 3.5, 4.2, 5.1, 1.8]),
            "fare_amount": pa.array([10.5, 15.3, 21.5, 25.0, 8.75, 16.5, 22.0, 28.5, 32.0, 12.5]),
            "tip_amount": pa.array([2.5, 3.0, 4.5, 5.0, 1.5, 3.5, 4.0, 6.0, 7.0, 2.0]),
        }
        
        batch = pa.RecordBatch.from_pydict(batch_data)
        table = pa.Table.from_batches([batch])
        arrow_ref = ray.put(table)
        ray_dataset = ray.data.from_arrow_refs([arrow_ref])
        return daft.from_ray_dataset(ray_dataset)
    
    except Exception as e:
        print(f"[ERROR] Error in create_mock_arrow_data: {e}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        raise ValueError(f"Error creating mock data: {e}")

def get_mock_taxi_data() -> daft.DataFrame:

    daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)
    daft_df = ray.get(create_mock_arrow_data.remote())
    return daft_df
    
    
daft_df = get_mock_taxi_data()

The exception is:

---------------------------------------------------------------------------
RayTaskError(TypeError)                   Traceback (most recent call last)
Cell In[3], line 57
     53     daft_df = ray.get(create_mock_arrow_data.remote())
     54     return daft_df
---> 57 daft_df = get_mock_taxi_data()

Cell In[3], line 53, in get_mock_taxi_data()
     49 def get_mock_taxi_data() -> daft.DataFrame:
     51     daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)
---> 53     daft_df = ray.get(create_mock_arrow_data.remote())
     54     return daft_df

File [/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/_private/auto_init_hook.py:21](http://localhost:8888/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/_private/auto_init_hook.py#line=20), in wrap_auto_init.<locals>.auto_init_wrapper(*args, **kwargs)
     18 @wraps(fn)
     19 def auto_init_wrapper(*args, **kwargs):
     20     auto_init_ray()
---> 21     return fn(*args, **kwargs)

File [/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/_private/client_mode_hook.py:102](http://localhost:8888/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/_private/client_mode_hook.py#line=101), in client_mode_hook.<locals>.wrapper(*args, **kwargs)
     98 if client_mode_should_convert():
     99     # Legacy code
    100     # we only convert init function if RAY_CLIENT_MODE=1
    101     if func.__name__ != "init" or is_client_mode_enabled_by_default:
--> 102         return getattr(ray, func.__name__)(*args, **kwargs)
    103 return func(*args, **kwargs)

File [/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/api.py:42](http://localhost:8888/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/api.py#line=41), in _ClientAPI.get(self, vals, timeout)
     35 def get(self, vals, *, timeout=None):
     36     """get is the hook stub passed on to replace `ray.get`
     37 
     38     Args:
     39         vals: [Client]ObjectRef or list of these refs to retrieve.
     40         timeout: Optional timeout in milliseconds
     41     """
---> 42     return self.worker.get(vals, timeout=timeout)

File [/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/worker.py:433](http://localhost:8888/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/worker.py#line=432), in Worker.get(self, vals, timeout)
    431     op_timeout = max_blocking_operation_time
    432 try:
--> 433     res = self._get(to_get, op_timeout)
    434     break
    435 except GetTimeoutError:

File [/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/worker.py:461](http://localhost:8888/opt/homebrew/Caskroom/miniconda/base/envs/raydaft39/lib/python3.9/site-packages/ray/util/client/worker.py#line=460), in Worker._get(self, ref, timeout)
    459         logger.exception("Failed to deserialize {}".format(chunk.error))
    460         raise
--> 461     raise err
    462 if chunk.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once(
    463     "client_object_transfer_size_warning"
    464 ):
    465     size_gb = chunk.total_size [/](http://localhost:8888/) 2**30

RayTaskError(TypeError): ray::create_mock_arrow_data() (pid=232280, ip=10.1.98.248)
  File "[/home/ray/anaconda3/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle.py", line 1479](http://localhost:8888/home/ray/anaconda3/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle.py#line=1478), in dumps
    cp.dump(obj)
  File "[/home/ray/anaconda3/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle.py", line 1245](http://localhost:8888/home/ray/anaconda3/lib/python3.9/site-packages/ray/cloudpickle/cloudpickle.py#line=1244), in dump
    return super().dump(obj)
TypeError: cannot pickle 'builtins.LogicalPlanBuilder' object

Any idea how I can fix this?

[rchowell]

I believe the issue is returning a daft DataFrame from the @ray.remote method. I don't believe a DataFrame is serializable (hence "cannot pickle") so you will have to construct the dataframe on the side that receives the mock data.

Something like,

@ray.remote
def create_mock_arrow_data():
    daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)
    try:
        batch_data = {
            "pulocationid": pa.array([100, 101, 102, 103, 104, 105, 106, 107, 108, 109]),
            "dolocationid": pa.array([200, 201, 202, 203, 204, 205, 206, 207, 208, 209]),
            "passenger_count": pa.array([1, 2, 3, 1, 2, 4, 1, 1, 2, 3]),
            "trip_distance": pa.array([1.5, 2.3, 3.1, 4.7, 1.2, 2.8, 3.5, 4.2, 5.1, 1.8]),
            "fare_amount": pa.array([10.5, 15.3, 21.5, 25.0, 8.75, 16.5, 22.0, 28.5, 32.0, 12.5]),
            "tip_amount": pa.array([2.5, 3.0, 4.5, 5.0, 1.5, 3.5, 4.0, 6.0, 7.0, 2.0]),
        }
        
        batch = pa.RecordBatch.from_pydict(batch_data)
        table = pa.Table.from_batches([batch])
        arrow_ref = ray.put(table)
        return ray.data.from_arrow_refs([arrow_ref])

def get_mock_taxi_data() -> daft.DataFrame:

    daft.context.set_runner_ray(address=RAY_ADDRESS, noop_if_initialized=True)
    ray_dataset = ray.get(create_mock_arrow_data.remote())
    return daft.from_ray_dataset(ray_dataset)
    
    
daft_df = get_mock_taxi_data()

Please let me know if these adjustments work.

[snowch]

Thanks @rchowell - unfortunately, ray dataset's don't support ray client mode:

RuntimeError: Global node is not initialized.

ray-project/ray#41333