Select DISTINCT with LIMIT 10 is doing a full-scan of the database

[shner-elmo]

Describe the bug

So I noticed that when executing the query SELECT DISTINCT my_column FROM dataset LIMIT 10 (trough the Python connector), it was taking lots of time and consuming most of my RAM.
I then looked at the query plan, and it seems like its actually doing GROUP BY my_column which causes a full-scan, what makes it even worse, is that all 10 values returned are present in the first parquet file in the dataset (pyarrow.Dataset.files[0]), so it could've just stopped scanning after the first file immediately.

To Reproduce

import datafusion
import time

start = time.perf_counter()

ctx = datafusion.SessionContext()
ctx.register_dataset('dataset', dataset)
out = ctx.sql("SELECT DISTINCT my_column FROM dataset LIMIT 10")

print('datafusion', time.perf_counter() - start)

print(repr(out.logical_plan()))
print('-' * 10)
print(repr(out.optimized_logical_plan()))
print('-' * 10)
print(repr(out.execution_plan()))
print('-' * 10)
print(repr(out.explain()))
print('-' * 10)

datafusion 0.0003965760000141927
Limit: skip=0, fetch=10
  Distinct:
    Projection: dataset.canonical_url
      TableScan: dataset
----------
Limit: skip=0, fetch=10
  Aggregate: groupBy=[[dataset.canonical_url]], aggr=[[]]
    TableScan: dataset projection=[canonical_url]
----------
GlobalLimitExec: skip=0, fetch=10
  CoalescePartitionsExec
    LocalLimitExec: fetch=10
      AggregateExec: mode=FinalPartitioned, gby=[canonical_url@0 as canonical_url], aggr=[]
        CoalesceBatchesExec: target_batch_size=8192
          RepartitionExec: partitioning=Hash([canonical_url@0], 8), input_partitions=568
            AggregateExec: mode=Partial, gby=[canonical_url@0 as canonical_url], aggr=[]
              DatasetExec: number_of_fragments=568, projection=[canonical_url]

----------
DataFrame()
+---------------+---------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                        |
+---------------+---------------------------------------------------------------------------------------------+
| logical_plan  | Limit: skip=0, fetch=10                                                                     |
|               |   Aggregate: groupBy=[[dataset.canonical_url]], aggr=[[]]                                   |
|               |     TableScan: dataset projection=[canonical_url]                                           |
| physical_plan | GlobalLimitExec: skip=0, fetch=10                                                           |
|               |   CoalescePartitionsExec                                                                    |
|               |     LocalLimitExec: fetch=10                                                                |
|               |       AggregateExec: mode=FinalPartitioned, gby=[canonical_url@0 as canonical_url], aggr=[] |
|               |         CoalesceBatchesExec: target_batch_size=8192                                         |
|               |           RepartitionExec: partitioning=Hash([canonical_url@0], 8), input_partitions=568    |
|               |             AggregateExec: mode=Partial, gby=[canonical_url@0 as canonical_url], aggr=[]    |
|               |               DatasetExec: number_of_fragments=568, projection=[canonical_url]              |
|               |                                                                                             |
+---------------+---------------------------------------------------------------------------------------------+
None
----------

Expected behavior

It should stop scanning at the first 10 unique values found.

Additional context

For context, I'm querying a Parquet dataset, and the files are stored locally:

import pyarrow.dataset as ds
dataset = ds.dataset(...)

[alamb]

I then looked at the query plan, and it seems like its actually doing GROUP BY my_column which causes a full-scan,

Yes, that is what I would expect for this kind of query

what makes it even worse, is that all 10 values returned are present in the first parquet file in the dataset (pyarrow.Dataset.files[0]), so it could've just stopped scanning after the first file immediately.

While it happens to be the case for your particular dataset that all values are present in the first tile, I don't think there is any way for datafusion to know that. To answer the query faithfully it needs to check all the files (what if there was a new value in the last file?)

[alamb]

(so in other words, I don't think this is a bug -- I think DataFusion is working as expected)

[shner-elmo]

While it happens to be the case for your particular dataset that all values are present in the first tile, I don't think there is any way for datafusion to know that. To answer the query faithfully it needs to check all the files (what if there was a new value in the last file?)

Of course it cannot know that. But it should keep scanning the data only if it didnt find 10 values yet, i.e.

unique_values = set()
for val in column:
   unique_values.add(val)

   if len(unique_values) == 10:
       rerturn unique_values

This is how it works in all the major RDBMS systems, why cant we also use this "optimization"?

I think the real issue is that we are using a GROUP BY behind the scenes to get the DISTINCT values, which while it does work, its not the appropriate solution imho, because for a GROUP BY you always need to do a full-scan before you can perform any sort of aggregation, on the other hand with a DISTINCT query you really dont have to.

The fact of the matter is that it is possible to make this query much more efficient then it currently is.
I'm willing to work on a fix if you guys are interested.

[alamb]

Of course it cannot know that. But it should keep scanning the data only if it didnt find 10 values yet, i.e.

~My point is that I am not sure how DataFusion would know there are only 10 unique values. ~

Update: I see now that DataFusion could stop because it knows there is a limit

I'm willing to work on a fix if you guys are interested.

If you have ways to improve things, that would be most appreciated. I am still not quite sure what you would do in this case that would work generally, but there may be some cleverness I am missing.

[msirek]

Is the issue that the RepartitionExec doesn't know how to stop processing early to honor the LocalLimitExec?
https://github.com/apache/arrow-datafusion/blob/7b12666ec87ea741c3f5b56ddf1647f6d794f9e3/datafusion/physical-plan/src/repartition/mod.rs#L537-L544

Maybe in cases where the partial group-by is the same as the outer group-by and there is no aggregate function applied requiring all rows in the group to be processed, it would be legal and useful to push LocalLimitExec below the RepartitionExec? That's assuming the partial group-by knows how to quit early when it's consumed by a LocalLimitExec.

[alamb]

Is the issue that the RepartitionExec doesn't know how to stop processing early to honor the LocalLimitExec?

I think you could potentially use the limit information in the GroupBy:

https://github.com/apache/arrow-datafusion/blob/e1bef86240a99bd9a1032a63833b255922065fd3/datafusion/physical-plan/src/aggregates/mod.rs#L281

And then you could teach the operators to switch to output mode once there are that may groups.

[msirek]

I'm working on a fix for this.