def test_field_projection_existence_only(self, file_spec, random_access,
parallelism):
with contextlib.closing(file_spec(self.create_tempfile,
random_access)) as files:
with riegeli.RecordWriter(
files.writing_open(),
owns_dest=files.writing_should_close,
assumed_pos=files.writing_assumed_pos,
options=f'{record_writer_options(parallelism)},transpose'
) as writer:
for i in range(23):
writer.write_message(sample_message(i, 10000))
with riegeli.RecordReader(
files.reading_open(),
owns_src=files.reading_should_close,
assumed_pos=files.reading_assumed_pos,
field_projection=[
[
records_test_pb2.SimpleMessage.DESCRIPTOR.
fields_by_name['id'].number
],
[
records_test_pb2.SimpleMessage.DESCRIPTOR.
fields_by_name['payload'].number,
riegeli.EXISTENCE_ONLY,
],
]) as reader:
for i in range(23):
self.assertEqual(
reader.read_message(records_test_pb2.SimpleMessage),
records_test_pb2.SimpleMessage(id=i, payload=b''))
self.assertIsNone(
reader.read_message(records_test_pb2.SimpleMessage))
def test_field_projection(self, file_spec, random_access, parallelism):
with contextlib.closing(file_spec(self.create_tempfile,
random_access)) as files:
with riegeli.RecordWriter(
files.writing_open(),
close=files.writing_should_close,
assumed_pos=files.writing_assumed_pos,
options=record_writer_options(parallelism) +
',transpose') as writer:
for i in range(23):
writer.write_message(sample_message(i, 10000))
with riegeli.RecordReader(
files.reading_open(),
close=files.reading_should_close,
assumed_pos=files.reading_assumed_pos,
field_projection=[[
records_test_pb2.SimpleMessage.DESCRIPTOR.
fields_by_name['id'].number
]]) as reader:
for i in range(23):
self.assertEqual(
reader.read_message(records_test_pb2.SimpleMessage),
records_test_pb2.SimpleMessage(id=i))
self.assertIsNone(
reader.read_message(records_test_pb2.SimpleMessage))
def sample_message_id_only(i):
return records_test_pb2.SimpleMessage(id=i)
def sample_message(i, size):
return records_test_pb2.SimpleMessage(id=i, payload=sample_string(i, size))
