def test_apply(self):
extractor = ChildrenBagExtractor()
create_or_apply_quant(paths.QUANTLEVELS, [extractor])
self.assertIsNotNone(extractor.levels)
model_levels = QuantizationLevels().load(source=paths.QUANTLEVELS)._levels["children"]
for key in model_levels:
self.assertListEqual(list(model_levels[key]), list(extractor.levels[key]))
def test_create(self):
session = create_spark("test_quant_util")
extractor = ChildrenBagExtractor()
with tempfile.NamedTemporaryFile(mode="r+b", suffix="-quant.asdf") as tmp:
path = tmp.name
uast_extractor = ParquetLoader(session, paths.PARQUET_DIR) \
.link(Moder("file")) \
.link(UastRow2Document()) \
.link(UastDeserializer())
create_or_apply_quant(path, [extractor], uast_extractor)
self.assertIsNotNone(extractor.levels)
self.assertTrue(os.path.exists(path))
model_levels = QuantizationLevels().load(source=path)._levels["children"]
for key in model_levels:
self.assertListEqual(list(model_levels[key]), list(extractor.levels[key]))