code: def test_with_batch_reader(
      scalar_dataset, shuffling_queue_capacity, data_loader_type
  ):
      """See if we are getting correct batch sizes when using DataLoader with make_batch_reader"""
      pytorch_compatible_fields = [
          k
          for k, v in scalar_dataset.data[0].items()
          if not isinstance(v, (np.datetime64, np.unicode_))
      ]
      with data_loader_type(
          make_batch_reader(
              scalar_dataset.url,
              schema_fields=pytorch_compatible_fields,
          ),
          batch_size=3,
          shuffling_queue_capacity=shuffling_queue_capacity,
      ) as loader:
          batches = list(loader)
          assert len(scalar_dataset.data) == sum(
              batch["id"].shape[0] for batch in batches
          )
  
          # list types are broken in pyarrow 0.15.0. Don't test list-of-int field
          if pa.__version__ != "0.15.0":
              assert len(scalar_dataset.data) == sum(
                  batch["int_fixed_size_list"].shape[0]
                  for batch in batches
              )
              assert batches[0]["int_fixed_size_list"].shape[1] == len(
                  scalar_dataset.data[0]["int_fixed_size_list"]
              )
  
imports: from decimal import Decimal
  from packaging import version
  import numpy
  import pyarrow
  import pytest
  import torch
  from petastorm import make_reader, TransformSpec, make_batch_reader
  from petastorm.pytorch import _sanitize_pytorch_types, DataLoader, BatchedDataLoader, decimal_friendly_collate, InMemBatchedDataLoader, _load_rows_into_mem
  from petastorm.tests.test_common import TestSchema
label_desc: Use image node_id 14 for calling a global function with example usage: test_with_batch_reader(scalar_dataset, shuffling_queue_capacity, data_loader_type) without return types