import warnings from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional, Sequence, Tuple, Union import numpy as np import pyarrow import torch from ray._common.utils import env_bool from ray.data.collate_fn import ( TensorBatchReturnType, TensorBatchType, _is_nested_tensor_sequence, _is_tensor, _is_tensor_mapping, _is_tensor_sequence, _is_tensor_sequence_mapping, ) from ray.data.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed # Default non-blocking transfer for tensors. DEFAULT_TENSOR_NON_BLOCKING_TRANSFER = env_bool( "RAY_AIR_DEFAULT_TENSOR_NON_BLOCKING_TRANSFER", True, ) def convert_table_to_torch_tensor( data_batch: pyarrow.Table, columns: Optional[Union[List[str], List[List[str]]]] = None, column_dtypes: Optional[Union[torch.dtype, List[torch.dtype]]] = None, unsqueeze: bool = True, ) -> Union[torch.Tensor, List[torch.Tensor]]: """Converts a PyArrow table to a torch Tensor or list of torch Tensors. The return type matches the format of ``columns``: a flat list of column names produces a single tensor; a list of lists produces a list of tensors (one per group). If ``columns`` is None, all columns are used. Args: data_batch: The PyArrow table to convert. columns: Column names to include. A list of lists returns one tensor per group (useful for multi-input models). None uses all columns. column_dtypes: Torch dtype(s) for the output. A single dtype applies to all columns/groups. A list must match the number of groups when ``columns`` is a list of lists. unsqueeze: If True, reshape each per-column tensor from (N,) to (N, 1) before concatenating. Defaults to True. Returns: A single tensor of shape (N, len(columns)), or a list of tensors when ``columns`` is a list of lists. """ multi_input = columns and isinstance(columns[0], (list, tuple)) if columns is None: columns = data_batch.column_names if not multi_input and column_dtypes and not isinstance(column_dtypes, torch.dtype): raise TypeError( "If `columns` is a list of strings, " "`column_dtypes` must be None or a single `torch.dtype`." f"Got {type(column_dtypes)} instead." ) if multi_input: if not isinstance(column_dtypes, (list, tuple)): column_dtypes = [column_dtypes] * len(columns) return [ _columns_to_tensor(data_batch, group, dtype, unsqueeze) for group, dtype in zip(columns, column_dtypes) ] return _columns_to_tensor(data_batch, columns, column_dtypes, unsqueeze) def _columns_to_tensor( table: pyarrow.Table, column_names: List[str], dtype: Optional[torch.dtype], unsqueeze: bool, ) -> torch.Tensor: """Convert selected columns from a PyArrow table into a single tensor.""" from ray.data._internal.arrow_block import ArrowBlockAccessor numpy_batch = ArrowBlockAccessor(table).to_numpy(columns=column_names) tensors = [] for col in column_names: try: t = convert_ndarray_to_torch_tensor(numpy_batch[col], dtype) except Exception as e: raise ValueError( f"Failed to convert column {col} to a Torch Tensor of dtype " f"{dtype}. See above exception chain for the exact failure." ) from e if unsqueeze: t = t.unsqueeze(1) tensors.append(t) if len(tensors) > 1: return torch.cat(tensors, dim=1) return tensors[0] def convert_ndarray_to_torch_tensor( ndarray: np.ndarray, dtype: Optional[torch.dtype] = None, device: Optional[Union[str, "torch.device"]] = None, pin_memory: bool = False, ) -> torch.Tensor: """Convert a NumPy ndarray to a Torch Tensor. Args: ndarray: A NumPy ndarray that we wish to convert to a Torch Tensor. dtype: A Torch dtype for the created tensor; if None, the dtype will be inferred from the NumPy ndarray data. device: The device on which the tensor(s) should be placed; if None, the Torch tensor(s) will be constructed on the CPU. pin_memory: Whether to pin the memory of the created tensors. Returns: A Torch Tensor. """ ndarray = _unwrap_ndarray_object_type_if_needed(ndarray) # Object dtype cannot be converted into PyTorch Tensor. if ndarray.dtype.type is np.object_: raise RuntimeError( "Numpy array of object dtype cannot be converted to a Torch Tensor. This " "may because the numpy array is a ragged tensor--it contains items of " "different sizes. If using `iter_torch_batches()` API, you can pass in a " "`collate_fn` argument to specify custom logic to convert the Numpy array " "batch to a Torch tensor batch." ) # The numpy array is not always writeable as it can come from the Ray object store. # Numpy will throw a verbose warning here, which we suppress, as we don't write # to the tensors. We also don't want to copy the array to avoid memory overhead. # Original warning: https://github.com/pytorch/pytorch/blob/v1.13.0/ # torch/csrc/utils/tensor_numpy.cpp#L198-L206 with warnings.catch_warnings(): warnings.simplefilter("ignore") result = torch.as_tensor(ndarray, dtype=dtype, device=device) if pin_memory: assert result.device.type == "cpu", ( "Pin memory is only supported for CPU tensors. " f"Got device: {result.device} and pin_memory: {pin_memory}." ) result = result.pin_memory() return result def convert_ndarray_batch_to_torch_tensor_batch( ndarrays: Union[np.ndarray, Dict[str, np.ndarray]], dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, device: Optional[Union[str, "torch.device"]] = None, pin_memory: bool = False, ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: """Convert a NumPy ndarray batch to a Torch Tensor batch. Args: ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a Torch Tensor. dtypes: A (dict of) Torch dtype(s) for the created tensor; if None, the dtype will be inferred from the NumPy ndarray data. device: The device on which the tensor(s) should be placed; if None, the Torch tensor(s) will be constructed on the CPU. pin_memory: Whether to pin the memory of the created tensors. Returns: A (dict of) Torch Tensor(s). """ if isinstance(ndarrays, np.ndarray): # Single-tensor case. if isinstance(dtypes, dict): if len(dtypes) != 1: raise ValueError( "When constructing a single-tensor batch, only a single dtype " f"should be given, instead got: {dtypes}" ) dtypes = next(iter(dtypes.values())) batch = convert_ndarray_to_torch_tensor( ndarrays, dtype=dtypes, device=device, pin_memory=pin_memory, ) else: # Multi-tensor case. batch = { col_name: convert_ndarray_to_torch_tensor( col_ndarray, dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, device=device, pin_memory=pin_memory, ) for col_name, col_ndarray in ndarrays.items() } return batch def convert_ndarray_list_to_torch_tensor_list( ndarrays: Dict[str, List[np.ndarray]], dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, device: Optional[Union[str, "torch.device"]] = None, pin_memory: bool = False, ) -> Dict[str, List[torch.Tensor]]: """Convert a dict mapping column names to lists of ndarrays to Torch Tensors. Args: ndarrays: A dict mapping column names to lists of ndarrays that we wish to convert to Torch Tensors. dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype will be inferred from the NumPy ndarray data. device: The device on which the tensor(s) should be placed; if None, the Torch tensor(s) will be constructed on the CPU. pin_memory: Whether to pin the memory of the created tensors. Returns: A dict mapping column names to lists of Tensors. """ return { col_name: [ convert_ndarray_batch_to_torch_tensor_batch( ndarray, dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, device=device, pin_memory=pin_memory, ) for ndarray in col_ndarrays ] for col_name, col_ndarrays in ndarrays.items() } def arrow_batch_to_tensors( batch: pyarrow.Table, dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, combine_chunks: bool = False, pin_memory: bool = False, threadpool: Optional[ThreadPoolExecutor] = None, ) -> Union[Dict[str, torch.Tensor], Dict[str, List[torch.Tensor]]]: """Convert PyArrow batch to PyTorch tensors. Args: batch: PyArrow batch to convert dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype will be inferred from the NumPy ndarray data. combine_chunks: If True, combine chunks in Arrow batch before converting to tensors. pin_memory: Whether to pin the memory of the created tensors. threadpool: Optional ThreadPoolExecutor for parallel processing. If provided, columns/arrays will be processed in parallel. If None, processing is sequential. Returns: When combine_chunks=True: A dictionary of column name to single tensor. When combine_chunks=False: A dictionary of column name to list of tensors. """ from ray.data._internal.arrow_block import ArrowBlockAccessor from ray.data._internal.arrow_ops import transform_pyarrow if combine_chunks: numpy_batch = ArrowBlockAccessor(batch).to_batch_format("numpy") num_columns = len(numpy_batch) if num_columns > 1 and threadpool is not None: # Process columns in parallel using provided threadpool def process_column( col_name_col_array: Tuple[str, np.ndarray] ) -> Tuple[str, torch.Tensor]: col_name, col_array = col_name_col_array return col_name, convert_ndarray_batch_to_torch_tensor_batch( col_array, dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, pin_memory=pin_memory, ) # Submit all columns to threadpool and collect results processed_cols = threadpool.map(process_column, numpy_batch.items()) return dict(processed_cols) else: # Sequential processing for single column or single worker return { col_name: convert_ndarray_batch_to_torch_tensor_batch( col_array, dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, pin_memory=pin_memory, ) for col_name, col_array in numpy_batch.items() } else: numpy_list = transform_pyarrow.table_to_numpy_dict_chunked( batch, ) # Count total number of arrays across all columns total_arrays = sum(len(arrays) for arrays in numpy_list.values()) num_columns = len(numpy_list) if total_arrays > 1 and threadpool is not None: # Process arrays in parallel using provided threadpool def process_array( array_item: Tuple[str, int, np.ndarray] ) -> Tuple[str, int, torch.Tensor]: col_name, array_index, array = array_item return ( col_name, array_index, convert_ndarray_batch_to_torch_tensor_batch( array, dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes, pin_memory=pin_memory, ), ) # Flatten arrays with column name and index for parallel processing array_items = [ (col_name, idx, array) for col_name, arrays in numpy_list.items() for idx, array in enumerate(arrays) ] # Submit all arrays to threadpool and collect results processed_arrays = list(threadpool.map(process_array, array_items)) # Initialize result with all columns from numpy_list, including empty ones # Pre-allocate lists of the correct size for each column result: Dict[str, List[torch.Tensor]] = { col_name: [None] * len(arrays) for col_name, arrays in numpy_list.items() } # Populate result with processed tensors for col_name, array_index, tensor in processed_arrays: result[col_name][array_index] = tensor return result else: # Sequential processing return convert_ndarray_list_to_torch_tensor_list( numpy_list, dtypes=dtypes, pin_memory=pin_memory, ) @torch.no_grad() def concat_tensors_to_device( tensor_sequence: Sequence[torch.Tensor], device: Optional[Union[str, "torch.device"]] = None, non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER, ) -> torch.Tensor: """Stack sequence of tensors into a contiguous GPU tensor. Args: tensor_sequence: Sequence of tensors to stack device: The device to move tensors to. If None, tensors are not moved. non_blocking: If True, perform device transfer without forcing a synchronization. Returns: A contiguous tensor on the target device """ # Assumes tensors have the same shape/dtype assert ( tensor_sequence ), f"Cannot stack empty sequence of tensors. Received: {tensor_sequence}" assert all( isinstance(t, torch.Tensor) for t in tensor_sequence ), "All items must be torch.Tensor. Found invalid types: " + str( [type(t) for t in tensor_sequence if not isinstance(t, torch.Tensor)] ) # If there is only one tensor and its device already matches, return it directly. if len(tensor_sequence) == 1 and ( device is None or tensor_sequence[0].device == torch.device(device) ): return tensor_sequence[0] first_dtype = tensor_sequence[0].dtype assert all(t.dtype == first_dtype for t in tensor_sequence), ( "All tensors must have the same dtype. " f"Expected: {first_dtype}, got: {[t.dtype for t in tensor_sequence]}" ) first_shape = tensor_sequence[0].shape[1:] assert all(t.shape[1:] == first_shape for t in tensor_sequence), ( "All tensors must have the same shape[1:]. " f"Expected: {first_shape}, got: {[t.shape[1:] for t in tensor_sequence]}" ) first = tensor_sequence[0] dtype = first.dtype shape_tail = first.shape[1:] total_rows = sum(t.shape[0] for t in tensor_sequence) # Allocate an empty Tensor on device result = torch.empty((total_rows, *shape_tail), dtype=dtype, device=device) row_start = 0 for t in tensor_sequence: row_end = row_start + t.shape[0] result[row_start:row_end].copy_(t, non_blocking=non_blocking) row_start = row_end return result def _get_type_str(batch: Any) -> str: """Get a string representation of the possibly nested type of the batch. >>> import torch >>> _get_type_str([1, 2, "???"]) 'list[int | str]' >>> _get_type_str({"a": [1, 2, 3], "b": 4}) 'dict[str, int | list[int]]' >>> _get_type_str({"a": torch.tensor(1), "b": [torch.tensor(2)]}) 'dict[str, Tensor | list[Tensor]]' >>> _get_type_str({"a": torch.tensor(1), "b": {"c": torch.tensor(2)}}) 'dict[str, Tensor | dict[str, Tensor]]' """ curr_type = type(batch).__name__ if isinstance(batch, (list, tuple)): val_types = " | ".join(sorted({_get_type_str(v) for v in batch})) invalid_type_str = f"{curr_type}[{val_types}]" elif isinstance(batch, dict): val_types = " | ".join(sorted({_get_type_str(v) for v in batch.values()})) invalid_type_str = f"{curr_type}[str, {val_types}]" else: invalid_type_str = curr_type return invalid_type_str @torch.no_grad() def move_tensors_to_device( batch: TensorBatchType, device: Optional[Union[str, "torch.device"]] = None, non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER, ) -> TensorBatchReturnType: """Move tensors to the specified device. Concatenate nested lists/tuples of tensors along the first (batch) dimension. For example, for the input ((feature_0_chunk_0,), (feature_1_chunk_0, feature_1_chunk_1)) the output will be (feature_0_chunk_0, feature_1_chunk_0+1) where each feature is concatenated along the batch dimension. Args: batch: A tensor or collection of tensors to move to device. Can be: - A single tensor - A sequence of tensors - A sequence of sequences of tensors. The inner sequence of tensors is combined during GPU transfer. - A mapping (e.g., dict) of keys to tensors or sequences of tensors. The sequence of tensors is combined during GPU transfer. device: The device to move tensors to. If None, tensors are not moved. non_blocking: If True, perform device transfer without forcing a synchronization. Returns: The input tensors moved to the specified device """ if device is None: return batch if _is_tensor(batch): return batch.to(device, non_blocking=non_blocking) elif _is_tensor_sequence(batch): return type(batch)([t.to(device, non_blocking=non_blocking) for t in batch]) elif _is_nested_tensor_sequence(batch): return type(batch)( [concat_tensors_to_device(t, device, non_blocking) for t in batch] ) elif _is_tensor_mapping(batch): return {k: t.to(device, non_blocking=non_blocking) for k, t in batch.items()} elif _is_tensor_sequence_mapping(batch): return { k: concat_tensors_to_device(v, device, non_blocking) for k, v in batch.items() } else: raise ValueError( f"Invalid input type: {_get_type_str(batch)}.\n" "Expected one of the following: " "torch.Tensor, " "List/Tuple[torch.Tensor], " "Dict[str, torch.Tensor], " "Mapping[str, List/Tuple[torch.Tensor]]" )