[BUG] `select from table limit` reads the full dataset and persists in memory.

[ayushdg]

What happened:
When performing a SELECT * FROM table LIMIT 10, from a table read in via parquet, I notice the full dataset being read and persisted on query execution.

What you expected to happen:
Nothing to happen at query execution and when the user does decide to persist/compute the result only the relevant subset of data is read in.

Minimal Complete Verifiable Example:

from dask_cuda import LocalCUDACluster
from distributed import Client, wait
import cudf
import dask_cudf
from dask_sql import Context
import dask

write_data = False

if __name__ == "__main__":
    cluster = LocalCUDACluster()
    client = Client(cluster)
    c = Context()
    
    if write_data:
        dask.datasets.timseries(start="2022-01-01", end="2024-01-01").to_parquet("test_data.parquet")


    ddf = dask_cudf.read_parquet("test_data.parquet")
    c.create_table("test", ddf, persist=False)

    # This results in the whole dataset persisted in memory and even though `len(res)==5` all the data is in memory
    res = c.sql("SELECT * from test LIMIT 10")

Anything else we need to know?:

Environment:

• dask-sql version: 2022.1.0
• Python version: 3.8
• Operating System: ubuntu 18.04
• Install method (conda, pip, source): conda

[charlesbluca]

Looks like this is happening because of a direct call to persist() in _apply_offset:

dask-sql/dask_sql/physical/rel/logical/limit.py

Line 61 in b571a60

df = df.persist()

Which is used by the DaskLimitPlugin to handle offsets and limits on select statements; I can try removing it to see if there are major performance implications.

[ayushdg]

Thanks for triaging. So there's a couple of things I'd like to discuss.

• Explore if not persisting is an option.
• Irrespective of the persist, do we need to read in the full data or can we read in a subset similar to doing ddf.head()

[charlesbluca]

do we need to read in the full data or can we read in a subset similar to doing ddf.head()

It looks like that's what's being done immediately after the persist; we check if the requested data is located in the first partition, and if so, simply use df.head():

dask-sql/dask_sql/physical/rel/logical/limit.py

Lines 62 to 66 in b571a60

	if not offset:
	# We do a (hopefully) very quick check: if the first partition
	# is already enough, we will just use this
	first_partition_length = len(df.partitions[0])
	if first_partition_length >= end:

Failing this, we end up computing the length of each partition, which IIUC ends up persisting the dataframe anyways?

dask-sql/dask_sql/physical/rel/logical/limit.py

Lines 69 to 71 in b571a60

	# First, we need to find out which partitions we want to use.
	# Therefore we count the total number of entries
	partition_borders = df.map_partitions(lambda x: len(x))

If that's the case, then I see nothing wrong with removing this persist call, as all it really does is ensure that we persist the dataframe even when we don't need to - although perhaps the idea here was similar to the persist calls in sort where we want to compute potentially complex operations before computing the subset of data we want?

[charlesbluca]

Relevant PR attempting to add ioc to dask dataframes dask/dask#6661

[ayushdg]

Relevant PR attempting to add ioc to dask dataframes dask/dask#6661

Nice! Once this lands do you think we can leverage some of this logic within our limit plugin?

[charlesbluca]

Yeah, I think even some smaller changes being discussed there (such as the addition of a pre-computed partition_sizes attribute) could be used to reduce the amount of cases where we would be forced to partially / fully persist a frame

[ayushdg]

Thanks for that context. I think we should be good to close this issue but should raise another issue in dask-sql to explore operations where we could leverage iloc and/or just precomputing partition_sizes instead of persisting once that PR lands.