(mockfs, infer_dictionary, pickled,
pickle_module)
| 2001 | @pytest.mark.parametrize( |
| 2002 | "pickled", [lambda x, m: x, lambda x, m: m.loads(m.dumps(x))]) |
| 2003 | def test_partitioning_factory_dictionary(mockfs, infer_dictionary, pickled, |
| 2004 | pickle_module): |
| 2005 | paths_or_selector = fs.FileSelector('subdir', recursive=True) |
| 2006 | format = ds.ParquetFileFormat() |
| 2007 | options = ds.FileSystemFactoryOptions('subdir') |
| 2008 | |
| 2009 | partitioning_factory = ds.DirectoryPartitioning.discover( |
| 2010 | ['group', 'key'], infer_dictionary=infer_dictionary) |
| 2011 | options.partitioning_factory = pickled(partitioning_factory, pickle_module) |
| 2012 | |
| 2013 | factory = ds.FileSystemDatasetFactory( |
| 2014 | mockfs, paths_or_selector, format, options) |
| 2015 | |
| 2016 | inferred_schema = factory.inspect() |
| 2017 | if infer_dictionary: |
| 2018 | expected_type = pa.dictionary(pa.int32(), pa.string()) |
| 2019 | assert inferred_schema.field('key').type == expected_type |
| 2020 | |
| 2021 | table = factory.finish().to_table().combine_chunks() |
| 2022 | actual = table.column('key').chunk(0) |
| 2023 | expected = pa.array(['xxx'] * 5 + ['yyy'] * 5).dictionary_encode() |
| 2024 | assert actual.equals(expected) |
| 2025 | |
| 2026 | # ARROW-9345 ensure filtering on the partition field works |
| 2027 | table = factory.finish().to_table(filter=ds.field('key') == 'xxx') |
| 2028 | actual = table.column('key').chunk(0) |
| 2029 | expected = expected.slice(0, 5) |
| 2030 | assert actual.equals(expected) |
| 2031 | else: |
| 2032 | assert inferred_schema.field('key').type == pa.string() |
| 2033 | |
| 2034 | |
| 2035 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected