diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..d2cdd3c0dc18 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -108,6 +108,7 @@ ## Bugfixes +* (Python) Fixed `WriteToFiles` failing to write files when `file_naming` places them in a subdirectory of `path`, since only `path` itself was created before the move ([#21313](https://github.com/apache/beam/issues/21313)). * (Python) Fixed BigQuery `FILE_LOADS` silently omitting files when dynamic destinations use different spellings of the same table ([#40001](https://github.com/apache/beam/pull/40001)). * (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)). * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)). diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py index 23d0f8c10c1a..c9055c9cf932 100644 --- a/sdks/python/apache_beam/io/fileio.py +++ b/sdks/python/apache_beam/io/fileio.py @@ -874,7 +874,10 @@ def process(self, element, w=beam.DoFn.WindowParam): destination)) move_from = [f.file_name for f in temp_file_results] - move_to = [f.file_name for f in final_file_results] + move_to = [ + filesystems.FileSystems.join(self.path.get(), f.file_name) + for f in final_file_results + ] _LOGGER.info( 'Moving %d temporary files to dir: %s as %s', @@ -882,18 +885,21 @@ def process(self, element, w=beam.DoFn.WindowParam): self.path.get(), move_to) - try: - filesystems.FileSystems.mkdirs(self.path.get()) - except IOError as e: - cause = repr(e) - if 'FileExistsError' not in cause: - # Usually harmless. Especially if see FileExistsError so no need to log - _LOGGER.debug('Fail to create dir for final destination: %s', cause) + # file_naming_fn may nest the final file name under a subdirectory of + # self.path, so create the parent directory of each final destination + # rather than just the base path. + for final_dir in {filesystems.FileSystems.split(move_to_path)[0] + for move_to_path in move_to}: + try: + filesystems.FileSystems.mkdirs(final_dir) + except IOError as e: + cause = repr(e) + if 'FileExistsError' not in cause: + # Usually harmless. Especially if see FileExistsError so no need to log + _LOGGER.debug('Fail to create dir for final destination: %s', cause) try: - filesystems.FileSystems.rename( - move_from, - [filesystems.FileSystems.join(self.path.get(), f) for f in move_to]) + filesystems.FileSystems.rename(move_from, move_to) except BeamIOError: # This error is not serious, because it may happen on a retry of the # bundle. We simply log it. diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py index 3bb250a34bd2..7fce8cc4c618 100644 --- a/sdks/python/apache_beam/io/fileio_test.py +++ b/sdks/python/apache_beam/io/fileio_test.py @@ -871,6 +871,30 @@ def test_write_to_dynamic_destination(self): equal_to([('odd', list(range(1, 100, 2))), ('even', list(range(0, 100, 2)))])) + def test_write_to_nested_directory_destination(self): + dir = self._new_tempdir() + + def nested_file_naming(window, pane, shard_index, total_shards, *args): + return FileSystems.join('subdir', 'test_file_%s' % shard_index) + + with TestPipeline() as p: + _ = ( + p + | "Create" >> beam.Create(range(10)) + | beam.Map(lambda x: str(x)) + | fileio.WriteToFiles( + path=dir, sink=fileio.TextSink, file_naming=nested_file_naming)) + + with TestPipeline() as p: + result = ( + p + | fileio.MatchFiles(FileSystems.join(dir, 'subdir', '*')) + | fileio.ReadMatches() + | beam.FlatMap(lambda f: f.read_utf8().strip().split('\n')) + | beam.Map(int)) + + assert_that(result, equal_to(list(range(10)))) + def test_write_to_different_file_types_some_spilling(self): dir = self._new_tempdir()