code: def test_load_rows_into_mem(
      synthetic_dataset, capacity, reader_factory
  ):
      """
      See if inmem buffer length is minimum of capacity and input data length.
      See if keys are correct.
      Also see reader is stopped.
      """
      with reader_factory(
          synthetic_dataset.url,
          schema_fields=[
              TestSchema.id,
              TestSchema.id_float,
              TestSchema.id_odd,
          ],
      ) as reader:
          keys, buffer = _load_rows_into_mem(
              reader, torch.as_tensor, capacity
          )
          assert len(buffer[0]) == min(
              capacity, len(synthetic_dataset.data)
          )
          assert list(keys)[0] == ("id")
          assert list(keys)[1] == ("id_float")
          assert list(keys)[2] == ("id_odd")
          assert reader.stopped
  
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 18 for calling a global function with example usage: test_load_rows_into_mem(synthetic_dataset, capacity, reader_factory) without return types