def test_issue_73():
"""test file reader by column name."""
writer = FileWriter(MKV_FILE_NAME, FILES_NUM)
data = get_mkv_data("../data/mindrecord/testVehPerData/")
mkv_schema_json = {
"file_name": {
"type": "string"
},
"id": {
"type": "number"
},
"prelabel": {
"type": "string"
},
"data": {
"type": "bytes"
}
}
writer.add_schema(mkv_schema_json, "mkv_schema")
writer.add_index(["file_name", "prelabel"])
writer.write_raw_data(data)
writer.commit()
reader = FileReader(MKV_FILE_NAME + "1", 4, ["file_name"])
for index, x in enumerate(reader.get_next()):
logger.info("#item{}: {}".format(index, x))
reader.close()
paths = [
"{}{}".format(MKV_FILE_NAME,
str(x).rjust(1, '0')) for x in range(FILES_NUM)
]
for x in paths:
os.remove("{}".format(x))
os.remove("{}.db".format(x))
def test_mkv_file_reader_with_negative_num_consumer():
"""test mkv file reader when the number of consumer is negative."""
writer = FileWriter(MKV_FILE_NAME, FILES_NUM)
data = get_mkv_data("../data/mindrecord/testVehPerData/")
mkv_schema_json = {
"file_name": {
"type": "string"
},
"id": {
"type": "number"
},
"prelabel": {
"type": "string"
},
"data": {
"type": "bytes"
}
}
writer.add_schema(mkv_schema_json, "mkv_schema")
writer.add_index(["file_name", "prelabel"])
writer.write_raw_data(data)
writer.commit()
with pytest.raises(Exception) as e:
FileReader(MKV_FILE_NAME + "1", -1)
assert "Consumer number should between 1 and" in str(e.value)
paths = [
"{}{}".format(MKV_FILE_NAME,
str(x).rjust(1, '0')) for x in range(FILES_NUM)
]
for x in paths:
os.remove("{}".format(x))
os.remove("{}.db".format(x))
def test_mkv_file_writer():
"""test mkv file writer using shard api"""
data = get_mkv_data("../data/mindrecord/testVehPerData/")
schema_json = {
"file_name": {
"type": "string"
},
"id": {
"type": "number"
},
"prelabel": {
"type": "string"
}
}
header = ShardHeader()
img_schema = header.build_schema(schema_json, ["data"], "img_schema")
schema_id = header.add_schema(img_schema)
assert schema_id == 0, 'failed on building schema.'
index_fields_list = ["id", "file_name"]
ret = header.add_index_fields(index_fields_list)
assert ret == SUCCESS, 'failed on adding index fields.'
writer = ShardWriter()
paths = ["{}{}".format(MKV_FILE_NAME, x) for x in range(FILES_NUM)]
ret = writer.open(paths)
assert ret == SUCCESS, 'failed on opening files.'
writer.set_header_size(1 << 24)
writer.set_page_size(1 << 25)
ret = writer.set_shard_header(header)
assert ret == SUCCESS, 'failed on setting header.'
ret = writer.write_raw_cv_data({schema_id: data})
assert ret == SUCCESS, 'failed on writing raw data.'
ret = writer.commit()
assert ret == SUCCESS, 'failed on committing.'
generator = ShardIndexGenerator(os.path.realpath(paths[0]))
generator.build()
generator.write_to_db()