code: def get_type_spec(
      schema: Union["pyarrow.lib.Schema", "PandasBlockSchema"],
      columns: Union[str, List[str]],
  ) -> Union[tf.TypeSpec, Dict[str, tf.TypeSpec]]:
      import pyarrow as pa
      from ray.data.extensions import TensorDtype, ArrowTensorType
  
      assert not isinstance(schema, type)
  
      dtypes: Dict[str, Union[np.dtype, pa.DataType]] = dict(
          zip(schema.names, schema.types)
      )
  
      def get_dtype(
          dtype: Union[np.dtype, pa.DataType]
      ) -> tf.dtypes.DType:
          if isinstance(dtype, pa.DataType):
              dtype = dtype.to_pandas_dtype()
          if isinstance(dtype, TensorDtype):
              dtype = dtype.element_dtype
          return tf.dtypes.as_dtype(dtype)
  
      def get_shape(
          dtype: Union[np.dtype, pa.DataType]
      ) -> Tuple[int, ...]:
          shape = (None,)
          if isinstance(dtype, ArrowTensorType):
              dtype = dtype.to_pandas_dtype()
          if isinstance(dtype, TensorDtype):
              shape += dtype.element_shape
          return shape
  
      def get_tensor_spec(
          dtype: Union[np.dtype, pa.DataType], *, name: str
      ) -> tf.TypeSpec:
          shape, dtype = get_shape(dtype), get_dtype(dtype)
          # Batch dimension is always `None`. So, if there's more than one `None`-valued
          # dimension, then the tensor is ragged.
          is_ragged = sum(dim is None for dim in shape) > 1
          if is_ragged:
              type_spec = tf.RaggedTensorSpec(shape, dtype=dtype)
          else:
              type_spec = tf.TensorSpec(shape, dtype=dtype, name=name)
          return type_spec
  
      if isinstance(columns, str):
          name, dtype = columns, dtypes[columns]
          return get_tensor_spec(dtype, name=name)
  
      return {
          name: get_tensor_spec(dtype, name=name)
          for name, dtype in dtypes.items()
          if name in columns
      }
  
file_path: ray/python/ray/air/_internal/tensorflow_utils.py
imports: from typing import TYPE_CHECKING, Dict, List, Optional, Union, Tuple
  import numpy
  import pyarrow
  import tensorflow
  from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
label_desc: Use image node_id 3 for calling a global function with example usage: get_type_spec(schema, columns) and returns: unknown, tf, shape, type_spec, get_tensor_spec