code: def test_transform_spec_support_return_tensor(
      scalar_dataset, reader_factory
  ):
      field1 = UnischemaField(
          name="abc", shape=(2, 3), numpy_dtype=np.float32
      )
  
      with pytest.raises(
          ValueError, match="field abc must be numpy array type"
      ):
          ArrowReaderWorker._check_shape_and_ravel("xyz", field1)
  
      with pytest.raises(
          ValueError, match="field abc must be the shape"
      ):
          ArrowReaderWorker._check_shape_and_ravel(
              np.zeros((2, 5)), field1
          )
  
      with pytest.raises(
          ValueError,
          match="field abc error: only support row major multi-dimensional array",
      ):
          ArrowReaderWorker._check_shape_and_ravel(
              np.zeros((2, 3), order="F"), field1
          )
  
      assert (6,) == ArrowReaderWorker._check_shape_and_ravel(
          np.zeros((2, 3)), field1
      ).shape
  
      for partial_shape in [(2, None), (None,), (None, None)]:
          field_with_unknown_dim = UnischemaField(
              name="abc", shape=partial_shape, numpy_dtype=np.float32
          )
          with pytest.raises(
              ValueError,
              match="All dimensions of a shape.*must be constant",
          ):
              ArrowReaderWorker._check_shape_and_ravel(
                  np.zeros((2, 3), order="F"), field_with_unknown_dim
              )
  
      def preproc_fn1(x):
          return pd.DataFrame(
              {
                  "tensor_col_1": x["id"].map(
                      lambda _: np.random.rand(2, 3)
                  ),
                  "tensor_col_2": x["id"].map(
                      lambda _: np.random.rand(3, 4, 5)
                  ),
              }
          )
  
      edit_fields = [
          ("tensor_col_1", np.float32, (2, 3), False),
          ("tensor_col_2", np.float32, (3, 4, 5), False),
      ]
  
      # This spec will remove all input columns and return one new column 'tensor_col_1' with shape (2, 3)
      spec1 = TransformSpec(
          preproc_fn1,
          edit_fields=edit_fields,
          removed_fields=list(scalar_dataset.data[0].keys()),
      )
  
      with reader_factory(
          scalar_dataset.url, transform_spec=spec1
      ) as reader:
          sample = next(reader)._asdict()
          assert len(sample) == 2
          assert (2, 3) == sample["tensor_col_1"].shape[1:] and (
              3,
              4,
              5,
          ) == sample["tensor_col_2"].shape[1:]
  
imports: import glob
  import itertools
  import numpy
  import pandas
  import pytest
  from pyarrow import parquet
  from petastorm import make_batch_reader, make_reader
  from petastorm.arrow_reader_worker import ArrowReaderWorker
  from petastorm.tests.test_common import create_test_scalar_dataset
  from petastorm.transform import TransformSpec
  from petastorm.unischema import UnischemaField
label_desc: Use image node_id 11 for calling a global function with example usage: test_transform_spec_support_return_tensor(scalar_dataset, reader_factory) and returns: pd