Repository: dask/dask

Issue:
BUG: compute=False not working correctly in new ORC writer
Using the new ORC writer from https://github.com/dask/dask/pull/7756 with `compute=False`, and then afterwards computing the result, only the first partition gets written:

```
import pandas as pd
import dask.dataframe as dd
import pathlib

df = pd.DataFrame(np.random.randn(100, 4), columns=['a', 'b', 'c', 'd'])
ddf = dd.from_pandas(df, npartitions=4)

In [5]: ddf.to_orc("test_orc_dataset")
Out[5]: (None, None, None, None)

In [9]: list(pathlib.Path("test_orc_dataset/").glob("*"))
Out[9]: 
[PosixPath('test_orc_dataset/part.0.orc'),
 PosixPath('test_orc_dataset/part.3.orc'),
 PosixPath('test_orc_dataset/part.2.orc'),
 PosixPath('test_orc_dataset/part.1.orc')]

In [11]: dataset = ddf.to_orc("test_orc_dataset_delayed", compute=False)

In [12]: dataset.compute()

In [13]: list(pathlib.Path("test_orc_dataset_delayed/").glob("*"))
Out[13]: [PosixPath('test_orc_dataset_delayed/part.0.orc')]
```

With the Parquet writer, it's working correctly.

cc @rjzamora 


Discussion:
Just tried to reproduce this and turns out I was on pyarrow 3 which failed at `ddf.to_orc`:


```python-traceback
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-15-e3bc609ace6f> in <module>
----> 1 ddf.to_orc("test_orc_dataset")

~/dask/dask/dataframe/core.py in to_orc(self, path, *args, **kwargs)
   4443         from .io import to_orc
   4444 
-> 4445         return to_orc(self, path, *args, **kwargs)
   4446 
   4447     @derived_from(pd.DataFrame)

~/dask/dask/dataframe/io/orc/core.py in to_orc(df, path, engine, write_index, storage_options, compute, compute_kwargs)
    228         if compute_kwargs is None:
    229             compute_kwargs = dict()
--> 230         return compute_as_if_collection(DataFrame, graph, part_tasks, **compute_kwargs)
    231     return Scalar(graph, name, "")

~/dask/dask/base.py in compute_as_if_collection(cls, dsk, keys, scheduler, get, **kwargs)
    311     schedule = get_scheduler(scheduler=scheduler, cls=cls, get=get)
    312     dsk2 = optimization_function(cls)(ensure_dict(dsk), keys, **kwargs)
--> 313     return schedule(dsk2, keys, **kwargs)
    314 
    315 

~/dask/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     77             pool = MultiprocessingPoolExecutor(pool)
     78 
---> 79     results = get_async(
     80         pool.submit,
     81         pool._max_workers,

~/dask/dask/local.py in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    512                             _execute_task(task, data)  # Re-execute locally
    513                         else:
--> 514                             raise_exception(exc, tb)
    515                     res, worker_id = loads(res_info)
    516                     state["cache"][key] = res

~/dask/dask/local.py in reraise(exc, tb)
    323     if exc.__traceback__ is not tb:
    324         raise exc.with_traceback(tb)
--> 325     raise exc
    326 
    327 

~/dask/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    221     try:
    222         task, data = loads(task_info)
--> 223         result = _execute_task(task, data)
    224         id = get_id()
    225         result = dumps((result, id))

~/dask/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg

~/dask/dask/utils.py in apply(func, args, kwargs)
     35         return func(*args, **kwargs)
     36     else:
---> 37         return func(*args)
     38 
     39 

~/dask/dask/dataframe/io/orc/arrow.py in write_partition(cls, df, path, fs, filename, **kwargs)
    109         table = pa.Table.from_pandas(df)
    110         with fs.open(fs.sep.join([path, filename]), "wb") as f:
--> 111             orc.write_table(table, f)
    112 
    113 

AttributeError: module 'pyarrow.orc' has no attribute 'write_table'
```

Thanks for raising @jorisvandenbossche !  I should be able to look into this today.  Note that I am in the process of preparing a PR to dramatically change the ORC API for Dask anyway.
It seems like the optimization is being over-eager, so as a workaround, you can do `dataset.compute(optimize_graph=False)`
> Just tried to reproduce this and turns out I was on pyarrow 3 which failed at ddf.to_orc:

Yes, ORC write support was only added more recently, so maybe dask could check for the version and give a nicer error message
> Note that I am in the process of preparing a PR to dramatically change the ORC API for Dask anyway.

Interested in seeing that :) (I am basically copying a bit what you did currently for the ORC code for a IO method in dask_geopandas)
>Interested in seeing that :) (I am basically copying a bit what you did currently for the ORC code for a IO method in dask_geopandas)

Yes - Sorry for all the churn here. I have a rough version of what I want to do in #8004 - Note that I decided to go with a slightly different design from what we have for parquet, because I realized that the parquet-engine model will make it a bit difficult to deal with both cudf-filtering support and the upcomming high-level expression work.  The primary change is that, instead of having an engine define `read_metadata` and `read_partition`,  the engine should define `get_dataset_info`, `construct_output_meta`, and `construct_partition_plan`.
> Yes - Sorry for all the churn here. I have a rough version of what I want to do in #8004 -

No need to be sorry! I am just genuinely interested to see how it could be organized / if I want to adapt my implementation in https://github.com/geopandas/dask-geopandas/pull/91 as well. Will try to take a look at #8004 one of the coming days.

Resolving patch:
diff --git a/dask/dataframe/io/orc/core.py b/dask/dataframe/io/orc/core.py
--- a/dask/dataframe/io/orc/core.py
+++ b/dask/dataframe/io/orc/core.py
@@ -2,6 +2,7 @@
 
 from fsspec.core import get_fs_token_paths
 from fsspec.utils import stringify_path
+from packaging.version import parse as parse_version
 
 from ....base import compute_as_if_collection, tokenize
 from ....highlevelgraph import HighLevelGraph
@@ -46,11 +47,16 @@ def __call__(self, parts):
         return _df
 
 
-def _get_engine(engine):
+def _get_engine(engine, write=False):
     # Get engine
     if engine == "pyarrow":
+        import pyarrow as pa
+
         from .arrow import ArrowORCEngine
 
+        if write and parse_version(pa.__version__) < parse_version("4.0.0"):
+            raise ValueError("to_orc is not supported for pyarrow<4.0.0")
+
         return ArrowORCEngine
     elif not isinstance(engine, ORCEngine):
         raise TypeError("engine must be 'pyarrow', or an ORCEngine object")
@@ -181,7 +187,7 @@ def to_orc(
     """
 
     # Get engine
-    engine = _get_engine(engine)
+    engine = _get_engine(engine, write=True)
 
     if hasattr(path, "name"):
         path = stringify_path(path)
@@ -207,7 +213,7 @@ def to_orc(
         write_index,
         storage_options,
     )
-    part_tasks = []
+    final_name = name + "-final"
     for d, filename in enumerate(filenames):
         dsk[(name, d)] = (
             apply,
@@ -219,13 +225,13 @@ def to_orc(
                 filename,
             ],
         )
-        part_tasks.append((name, d))
-    dsk[name] = (lambda x: None, part_tasks)
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[df])
+    part_tasks = list(dsk.keys())
+    dsk[(final_name, 0)] = (lambda x: None, part_tasks)
+    graph = HighLevelGraph.from_collections((final_name, 0), dsk, dependencies=[df])
 
     # Compute or return future
     if compute:
         if compute_kwargs is None:
             compute_kwargs = dict()
         return compute_as_if_collection(DataFrame, graph, part_tasks, **compute_kwargs)
-    return Scalar(graph, name, "")
+    return Scalar(graph, final_name, "")


Test patch:
diff --git a/dask/dataframe/io/tests/test_orc.py b/dask/dataframe/io/tests/test_orc.py
--- a/dask/dataframe/io/tests/test_orc.py
+++ b/dask/dataframe/io/tests/test_orc.py
@@ -1,3 +1,4 @@
+import glob
 import os
 import shutil
 import tempfile
@@ -143,8 +144,31 @@ def test_orc_aggregate_files_offset(orc_files):
     assert len(df2.partitions[0].index) > len(df2.index) // 2
 
 
+@pytest.mark.skipif(
+    parse_version(pa.__version__) < parse_version("4.0.0"),
+    reason=("PyArrow>=4.0.0 required for ORC write support."),
+)
 def test_orc_names(orc_files, tmp_path):
     df = dd.read_orc(orc_files)
     assert df._name.startswith("read-orc")
     out = df.to_orc(tmp_path, compute=False)
     assert out._name.startswith("to-orc")
+
+
+@pytest.mark.skipif(
+    parse_version(pa.__version__) < parse_version("4.0.0"),
+    reason=("PyArrow>=4.0.0 required for ORC write support."),
+)
+def test_to_orc_delayed(tmp_path):
+    # See: https://github.com/dask/dask/issues/8022
+    df = pd.DataFrame(np.random.randn(100, 4), columns=["a", "b", "c", "d"])
+    ddf = dd.from_pandas(df, npartitions=4)
+
+    eager_path = os.path.join(tmp_path, "eager_orc_dataset")
+    ddf.to_orc(eager_path)
+    assert len(glob.glob(os.path.join(eager_path, "*"))) == 4
+
+    delayed_path = os.path.join(tmp_path, "delayed_orc_dataset")
+    dataset = ddf.to_orc(delayed_path, compute=False)
+    dataset.compute()
+    assert len(glob.glob(os.path.join(delayed_path, "*"))) == 4