PR feedback.

Author: clarkzinzow

python/ray/data/dataset.py

@@ -990,25 +990,27 @@ def union(self, *other: List["Dataset[T]"]) -> "Dataset[T]":
         start_time = time.perf_counter()
         context = DatasetContext.get_current()
         tasks: List[ReadTask] = []
-        block_partitions: List[ObjectRef[BlockPartition]] = []
-        block_partitions_meta: List[ObjectRef[BlockPartitionMetadata]] = []
+        block_partition_refs: List[ObjectRef[BlockPartition]] = []
+        block_partition_meta_refs: List[ObjectRef[BlockPartitionMetadata]] = []
 
         datasets = [self] + list(other)
         for ds in datasets:
             bl = ds._plan.execute()
             if isinstance(bl, LazyBlockList):
                 tasks.extend(bl._tasks)
-                block_partitions.extend(bl._block_partitions)
-                block_partitions_meta.extend(bl._block_partitions_meta)
+                block_partition_refs.extend(bl._block_partition_refs)
+                block_partition_meta_refs.extend(bl._block_partition_meta_refs)
             else:
                 tasks.extend([ReadTask(lambda: None, meta) for meta in bl._metadata])
                 if context.block_splitting_enabled:
-                    block_partitions.extend(
+                    block_partition_refs.extend(
                         [ray.put([(b, m)]) for b, m in bl.get_blocks_with_metadata()]
                     )
                 else:
-                    block_partitions.extend(bl.get_blocks())
-                block_partitions_meta.extend([ray.put(meta) for meta in bl._metadata])
+                    block_partition_refs.extend(bl.get_blocks())
+                block_partition_meta_refs.extend(
+                    [ray.put(meta) for meta in bl._metadata]
+                )
 
         epochs = [ds._get_epoch() for ds in datasets]
         max_epoch = max(*epochs)
@@ -1029,7 +1031,7 @@ def union(self, *other: List["Dataset[T]"]) -> "Dataset[T]":
         dataset_stats.time_total_s = time.perf_counter() - start_time
         return Dataset(
             ExecutionPlan(
-                LazyBlockList(tasks, block_partitions, block_partitions_meta),
+                LazyBlockList(tasks, block_partition_refs, block_partition_meta_refs),
                 dataset_stats,
             ),
             max_epoch,

python/ray/data/impl/lazy_block_list.py

@@ -33,22 +33,24 @@ class LazyBlockList(BlockList):
     def __init__(
         self,
         tasks: List[ReadTask],
-        block_partitions: Optional[List[ObjectRef[MaybeBlockPartition]]] = None,
-        block_partitions_meta: Optional[List[ObjectRef[BlockPartitionMetadata]]] = None,
-        fetched_metadata: Optional[List[BlockPartitionMetadata]] = None,
+        block_partition_refs: Optional[List[ObjectRef[MaybeBlockPartition]]] = None,
+        block_partition_meta_refs: Optional[
+            List[ObjectRef[BlockPartitionMetadata]]
+        ] = None,
+        cached_metadata: Optional[List[BlockPartitionMetadata]] = None,
         ray_remote_args: Optional[Dict[str, Any]] = None,
         stats_uuid: str = None,
     ):
         """Create a LazyBlockList on the provided read tasks.
 
         Args:
             tasks: The read tasks that will produce the blocks of this lazy block list.
-            block_partitions: An optional list of already submitted read task futures
-                (i.e. block partition refs). This should be the same length as the tasks
-                argument.
-            block_partitions_meta: An optional list of block partition metadata refs.
-                This should be the same length as the tasks argument.
-            fetched_metadata: An optional list of already computed AND fetched metadata.
+            block_partition_refs: An optional list of already submitted read task
+                futures (i.e. block partition refs). This should be the same length as
+                the tasks argument.
+            block_partition_meta_refs: An optional list of block partition metadata
+                refs. This should be the same length as the tasks argument.
+            cached_metadata: An optional list of already computed AND fetched metadata.
                 This serves as a cache of fetched block metadata.
             ray_remote_args: Ray remote arguments for the read tasks.
             stats_uuid: UUID for the dataset stats, used to group and fetch read task
@@ -62,42 +64,42 @@ def __init__(
         self._execution_started = False
         self._remote_args = ray_remote_args or {}
         # Block partition metadata that have already been computed and fetched.
-        if fetched_metadata is not None:
-            self._fetched_metadata = fetched_metadata
+        if cached_metadata is not None:
+            self._cached_metadata = cached_metadata
         else:
-            self._fetched_metadata = [None] * len(tasks)
+            self._cached_metadata = [None] * len(tasks)
         # Block partition metadata that have already been computed.
-        if block_partitions_meta is not None:
-            self._block_partitions_meta = block_partitions_meta
+        if block_partition_meta_refs is not None:
+            self._block_partition_meta_refs = block_partition_meta_refs
         else:
-            self._block_partitions_meta = [None] * len(tasks)
+            self._block_partition_meta_refs = [None] * len(tasks)
         # Block partitions that have already been computed.
-        if block_partitions is not None:
-            self._block_partitions = block_partitions
+        if block_partition_refs is not None:
+            self._block_partition_refs = block_partition_refs
         else:
-            self._block_partitions = [None] * len(tasks)
-        assert len(tasks) == len(self._block_partitions), (
+            self._block_partition_refs = [None] * len(tasks)
+        assert len(tasks) == len(self._block_partition_refs), (
             tasks,
-            self._block_partitions,
+            self._block_partition_refs,
         )
-        assert len(tasks) == len(self._block_partitions_meta), (
+        assert len(tasks) == len(self._block_partition_meta_refs), (
             tasks,
-            self._block_partitions_meta,
+            self._block_partition_meta_refs,
         )
-        assert len(tasks) == len(self._fetched_metadata), (
+        assert len(tasks) == len(self._cached_metadata), (
             tasks,
-            self._fetched_metadata,
+            self._cached_metadata,
         )
 
     def get_metadata(self, fetch_if_missing: bool = False) -> List[BlockMetadata]:
         """Get the metadata for all blocks."""
-        if all(meta is not None for meta in self._fetched_metadata):
+        if all(meta is not None for meta in self._cached_metadata):
             # Always return fetched metadata if we already have it.
-            metadata = self._fetched_metadata
+            metadata = self._cached_metadata
         elif not fetch_if_missing:
             metadata = [
                 m if m is not None else t.get_metadata()
-                for m, t in zip(self._fetched_metadata, self._tasks)
+                for m, t in zip(self._cached_metadata, self._tasks)
             ]
         else:
             _, metadata = self._get_blocks_with_metadata()
@@ -136,9 +138,9 @@ def _submit_task(
     def copy(self) -> "LazyBlockList":
         return LazyBlockList(
             self._tasks.copy(),
-            block_partitions=self._block_partitions.copy(),
-            block_partitions_meta=self._block_partitions_meta.copy(),
-            fetched_metadata=self._fetched_metadata,
+            block_partition_refs=self._block_partition_refs.copy(),
+            block_partition_meta_refs=self._block_partition_meta_refs.copy(),
+            cached_metadata=self._cached_metadata,
             ray_remote_args=self._remote_args.copy(),
             stats_uuid=self._stats_uuid,
         )
@@ -147,9 +149,11 @@ def clear(self):
         """Clears all object references (block partitions and base block partitions)
         from this lazy block list.
         """
-        self._block_partitions = [None for _ in self._block_partitions]
-        self._block_partitions_meta = [None for _ in self._block_partitions_meta]
-        self._fetched_metadata = [None for _ in self._fetched_metadata]
+        self._block_partition_refs = [None for _ in self._block_partition_refs]
+        self._block_partition_meta_refs = [
+            None for _ in self._block_partition_meta_refs
+        ]
+        self._cached_metadata = [None for _ in self._cached_metadata]
 
     def _check_if_cleared(self):
         pass  # LazyBlockList can always be re-computed.
@@ -158,10 +162,12 @@ def _check_if_cleared(self):
     def split(self, split_size: int) -> List["LazyBlockList"]:
         num_splits = math.ceil(len(self._tasks) / split_size)
         tasks = np.array_split(self._tasks, num_splits)
-        block_partitions = np.array_split(self._block_partitions, num_splits)
-        block_partitions_meta = np.array_split(self._block_partitions_meta, num_splits)
+        block_partition_refs = np.array_split(self._block_partition_refs, num_splits)
+        block_partition_meta_refs = np.array_split(
+            self._block_partition_meta_refs, num_splits
+        )
         output = []
-        for t, b, m in zip(tasks, block_partitions, block_partitions_meta):
+        for t, b, m in zip(tasks, block_partition_refs, block_partition_meta_refs):
             output.append(
                 LazyBlockList(
                     t.tolist(),
@@ -179,8 +185,8 @@ def split_by_bytes(self, bytes_per_split: int) -> List["BlockList"]:
         cur_size = 0
         for t, b, bm in zip(
             self._tasks,
-            self._block_partitions,
-            self._block_partitions_meta,
+            self._block_partition_refs,
+            self._block_partition_meta_refs,
         ):
             m = t.get_metadata()
             if m.size_bytes is None:
@@ -206,13 +212,13 @@ def split_by_bytes(self, bytes_per_split: int) -> List["BlockList"]:
     def divide(self, part_idx: int) -> ("LazyBlockList", "LazyBlockList"):
         left = LazyBlockList(
             self._tasks[:part_idx],
-            self._block_partitions[:part_idx],
-            self._block_partitions_meta[:part_idx],
+            self._block_partition_refs[:part_idx],
+            self._block_partition_meta_refs[:part_idx],
         )
         right = LazyBlockList(
             self._tasks[part_idx:],
-            self._block_partitions[part_idx:],
-            self._block_partitions_meta[part_idx:],
+            self._block_partition_refs[part_idx:],
+            self._block_partition_meta_refs[part_idx:],
         )
         return left, right
 
@@ -243,26 +249,26 @@ def _get_blocks_with_metadata(
         all block metadata outputted by those tasks.
         """
         context = DatasetContext.get_current()
-        blocks, meta_refs = [], []
-        for block, meta_ref in self._iter_block_partition_refs():
-            blocks.append(block)
+        block_refs, meta_refs = [], []
+        for block_ref, meta_ref in self._iter_block_partition_refs():
+            block_refs.append(block_ref)
             meta_refs.append(meta_ref)
         if context.block_splitting_enabled:
             # If block splitting is enabled, fetch the partitions.
-            parts = ray.get(blocks)
-            blocks, metadata = [], []
+            parts = ray.get(block_refs)
+            block_refs, metadata = [], []
             for part in parts:
-                for block, meta in part:
-                    blocks.append(block)
+                for block_ref, meta in part:
+                    block_refs.append(block_ref)
                     metadata.append(meta)
-            self._fetched_metadata = metadata
-            return blocks, metadata
-        if all(meta is not None for meta in self._fetched_metadata):
+            self._cached_metadata = metadata
+            return block_refs, metadata
+        if all(meta is not None for meta in self._cached_metadata):
             # Short-circuit on cached metadata.
-            return blocks, self._fetched_metadata
+            return block_refs, self._cached_metadata
         if not meta_refs:
             # Short-circuit on empty set of block partitions.
-            assert not blocks, blocks
+            assert not block_refs, block_refs
             return [], []
         read_progress_bar = ProgressBar("Read progress", total=len(meta_refs))
         # Fetch the metadata in bulk.
@@ -273,8 +279,8 @@ def _get_blocks_with_metadata(
             meta_ref: data for meta_ref, data in zip(unique_meta_refs, metadata)
         }
         metadata = [ref_to_data[meta_ref] for meta_ref in meta_refs]
-        self._fetched_metadata = metadata
-        return blocks, metadata
+        self._cached_metadata = metadata
+        return block_refs, metadata
 
     def compute_first_block(self):
         """Kick off computation for the first block in the list.
@@ -298,7 +304,11 @@ def ensure_metadata_for_first_block(self) -> Optional[BlockMetadata]:
             return None
         metadata = self._tasks[0].get_metadata()
         if metadata.schema is not None:
+            # If pre-read schema is not null, we consider it to be "good enough" and use
+            # it.
             return metadata
+        # Otherwise, we trigger computation (if needed), wait until the task completes,
+        # and fetch the block partition metadata.
         try:
             _, metadata_ref = next(self._iter_block_partition_refs())
         except (StopIteration, ValueError):
@@ -307,16 +317,27 @@ def ensure_metadata_for_first_block(self) -> Optional[BlockMetadata]:
         else:
             # This blocks until the underlying read task is finished.
             metadata = ray.get(metadata_ref)
-            self._fetched_metadata[0] = metadata
+            self._cached_metadata[0] = metadata
         return metadata
 
     def iter_blocks_with_metadata(
         self,
+        block_for_metadata: bool = False,
     ) -> Iterator[Tuple[ObjectRef[Block], BlockMetadata]]:
         """Iterate over the blocks along with their metadata.
 
+        Note that, if block_for_metadata is False (default), this iterator returns
+        pre-read metadata from the ReadTasks given to this LazyBlockList so it doesn't
+        have to block on the execution of the read tasks. Therefore, the metadata may be
+        under-specified, e.g. missing schema or the number of rows. If fully-specified
+        block metadata is required, pass block_for_metadata=True.
+
         The length of this iterator is not known until execution.
 
+        Args:
+            block_for_metadata: Whether we should block on the execution of read tasks
+                in order to obtain fully-specified block metadata.
+
         Returns:
             An iterator of block references and the corresponding block metadata.
         """
@@ -339,11 +360,18 @@ def __next__(self):
                         part_ref, _ = next(self._base_iter)
                         partition = ray.get(part_ref)
                     else:
-                        block, _ = next(self._base_iter)
-                        metadata = outer._tasks[self._pos].get_metadata()
-                        partition = [(block, metadata)]
-                    for ref, metadata in partition:
-                        self._buffer.append((ref, metadata))
+                        block_ref, metadata_ref = next(self._base_iter)
+                        if block_for_metadata:
+                            # This blocks until the read task completes, returning
+                            # fully-specified block metadata.
+                            metadata = ray.get(metadata_ref)
+                        else:
+                            # This does not block, returning (possibly under-specified)
+                            # pre-read block metadata.
+                            metadata = outer._tasks[self._pos].get_metadata()
+                        partition = [(block_ref, metadata)]
+                    for block_ref, metadata in partition:
+                        self._buffer.append((block_ref, metadata))
                 return self._buffer.pop(0)
 
         return Iter()
@@ -379,24 +407,24 @@ def _get_or_compute(
         i: int,
     ) -> Tuple[ObjectRef[MaybeBlockPartition], ObjectRef[BlockPartitionMetadata]]:
         assert i < len(self._tasks), i
-        # Check if we need to compute more block_partitions.
-        if not self._block_partitions[i]:
+        # Check if we need to compute more block_partition_refs.
+        if not self._block_partition_refs[i]:
             # Exponentially increase the number computed per batch.
             for j in range(max(i + 1, i * 2)):
-                if j >= len(self._block_partitions):
+                if j >= len(self._block_partition_refs):
                     break
-                if not self._block_partitions[j]:
+                if not self._block_partition_refs[j]:
                     (
-                        self._block_partitions[j],
-                        self._block_partitions_meta[j],
+                        self._block_partition_refs[j],
+                        self._block_partition_meta_refs[j],
                     ) = self._submit_task(j)
-            assert self._block_partitions[i], self._block_partitions
-            assert self._block_partitions_meta[i], self._block_partitions_meta
-        return self._block_partitions[i], self._block_partitions_meta[i]
+            assert self._block_partition_refs[i], self._block_partition_refs
+            assert self._block_partition_meta_refs[i], self._block_partition_meta_refs
+        return self._block_partition_refs[i], self._block_partition_meta_refs[i]
 
     def _num_computed(self) -> int:
         i = 0
-        for b in self._block_partitions:
+        for b in self._block_partition_refs:
             if b is not None:
                 i += 1
         return i