github-actions[bot] commented on code in PR #41186:
URL: https://github.com/apache/beam/pull/41186#discussion_r4532862952
##########
sdks/python/apache_beam/yaml/yaml_io.py:
##########
@@ -99,6 +105,234 @@ def write_to_text(pcoll, path: str):
lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)
+def read_from_csv(
+ path: str,
+ *,
+ sep: str = ',',
+ comment: Optional[str] = None,
+ filename_column: Optional[str] = None,
+ error_handling: Optional[Mapping[str, Any]] = None,
+ **kwargs):
+ """Reads comma-separated values (csv) files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromCsv``.
+
+ With ``error_handling``, each record is parsed on its own, and records
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw record text.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ sep (str): The field delimiter.
+ comment (str): Lines that start with this character are skipped.
+ filename_column (str): If set, the name of a column to add to each
+ record, containing the path of the source file.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit records that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromCsv(
+ path,
+ sep=sep,
+ comment=comment,
+ filename_column=filename_column,
+ **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromCsv does not support {sorted(kwargs)} with error_handling.')
+
+ def sample(handle):
+ return pd.read_csv(
+ handle,
+ sep=sep,
+ comment=comment,
+ nrows=_SAMPLE_SIZE,
+ on_bad_lines='skip')
+
+ def parse_record(header, record, dtypes):
+ expected = len(next(csv.reader([header], delimiter=sep)))
+ actual = len(next(csv.reader(io.StringIO(record), delimiter=sep)))
+ if actual != expected:
+ raise ValueError(f'Expected {expected} fields, got {actual}.')
+ return pd.read_csv(
+ io.StringIO(header + record), sep=sep, comment=comment, dtype=dtypes)
+
+ def parse_file(lines, proxy):
+ dtypes = proxy.dtypes.drop(filename_column, errors='ignore').to_dict()
+ header = None
+ for record in _join_quoted_lines(lines):
Review Comment:
**🟡 Medium**: There are a few edge cases with `comment` and header handling
in `read_from_csv`:
1. `_join_quoted_lines(lines)` runs before comment lines are skipped. If a
comment line contains an odd number of quotes (e.g. `# don"t use this`),
`_join_quoted_lines` will treat it as an open multiline field and merge it with
the header or subsequent data records.
2. In `parse_record`, `csv.reader` does not support the `comment` character,
whereas `pd.read_csv(..., comment=comment)` does. An inline comment containing
a delimiter (e.g. `a,1 # comment, with comma`) will fail the `actual !=
expected` check in `csv.reader` even though `pd.read_csv` considers it valid.
3. `expected = len(next(csv.reader([header], delimiter=sep)))` is recomputed
on every record instead of once per file, and when reading multiple files via a
glob, `header` in subsequent files is never validated against `proxy.columns`
(nor is a UTF-8 decode error on the header line handled, which currently causes
the first data row to be consumed as `header`).
##########
sdks/python/apache_beam/yaml/yaml_io.py:
##########
@@ -99,6 +105,234 @@ def write_to_text(pcoll, path: str):
lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)
+def read_from_csv(
+ path: str,
+ *,
+ sep: str = ',',
+ comment: Optional[str] = None,
+ filename_column: Optional[str] = None,
+ error_handling: Optional[Mapping[str, Any]] = None,
+ **kwargs):
+ """Reads comma-separated values (csv) files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromCsv``.
+
+ With ``error_handling``, each record is parsed on its own, and records
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw record text.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ sep (str): The field delimiter.
+ comment (str): Lines that start with this character are skipped.
+ filename_column (str): If set, the name of a column to add to each
+ record, containing the path of the source file.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit records that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromCsv(
+ path,
+ sep=sep,
+ comment=comment,
+ filename_column=filename_column,
+ **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromCsv does not support {sorted(kwargs)} with error_handling.')
+
+ def sample(handle):
+ return pd.read_csv(
+ handle,
+ sep=sep,
+ comment=comment,
+ nrows=_SAMPLE_SIZE,
+ on_bad_lines='skip')
+
+ def parse_record(header, record, dtypes):
+ expected = len(next(csv.reader([header], delimiter=sep)))
+ actual = len(next(csv.reader(io.StringIO(record), delimiter=sep)))
+ if actual != expected:
+ raise ValueError(f'Expected {expected} fields, got {actual}.')
+ return pd.read_csv(
+ io.StringIO(header + record), sep=sep, comment=comment, dtype=dtypes)
+
+ def parse_file(lines, proxy):
+ dtypes = proxy.dtypes.drop(filename_column, errors='ignore').to_dict()
+ header = None
+ for record in _join_quoted_lines(lines):
+ if not record.strip() or (comment and record.startswith(comment)):
+ continue
+ if header is None:
+ header = record
+ else:
+ yield record, functools.partial(parse_record, header, record, dtypes)
+
+ return _ReadFilesWithErrors(
+ path, sample, parse_file, filename_column, error_handling)
+
+
+def read_from_json(
+ path: str, *, error_handling: Optional[Mapping[str, Any]] = None,
**kwargs):
+ """Reads newline-delimited JSON files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromJson``.
+
+ With ``error_handling``, each line is parsed on its own, and lines
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw line.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit lines that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromJson(path, **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromJson does not support {sorted(kwargs)} with error_handling.')
+
+ def read_json(text):
+ return pd.read_json(
+ io.StringIO(text), orient='records', lines=True, dtype=False)
+
+ def sample(handle):
+ good_lines = []
+ for line in itertools.islice(handle, _SAMPLE_SIZE):
+ try:
+ if isinstance(json.loads(line), dict):
+ good_lines.append(line)
+ except ValueError:
+ pass
+ return read_json(''.join(good_lines))
Review Comment:
**🟡 Medium**: If `good_lines` is empty (for example, if the first
`_SAMPLE_SIZE` lines in the first file are malformed or the file is empty),
`read_json('')` produces an empty DataFrame with 0 columns (`proxy.columns` is
empty), causing `parse_record` to strip all columns from every subsequent valid
record via `.reindex(columns=proxy.columns)`. If `not good_lines`, this should
raise a descriptive `ValueError` (or continue scanning past `_SAMPLE_SIZE`
until at least one valid JSON object is found).
##########
sdks/python/apache_beam/yaml/yaml_io.py:
##########
@@ -99,6 +105,234 @@ def write_to_text(pcoll, path: str):
lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)
+def read_from_csv(
+ path: str,
+ *,
+ sep: str = ',',
+ comment: Optional[str] = None,
+ filename_column: Optional[str] = None,
+ error_handling: Optional[Mapping[str, Any]] = None,
+ **kwargs):
+ """Reads comma-separated values (csv) files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromCsv``.
+
+ With ``error_handling``, each record is parsed on its own, and records
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw record text.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ sep (str): The field delimiter.
+ comment (str): Lines that start with this character are skipped.
+ filename_column (str): If set, the name of a column to add to each
+ record, containing the path of the source file.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit records that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromCsv(
+ path,
+ sep=sep,
+ comment=comment,
+ filename_column=filename_column,
+ **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromCsv does not support {sorted(kwargs)} with error_handling.')
+
+ def sample(handle):
+ return pd.read_csv(
+ handle,
+ sep=sep,
+ comment=comment,
+ nrows=_SAMPLE_SIZE,
+ on_bad_lines='skip')
+
+ def parse_record(header, record, dtypes):
+ expected = len(next(csv.reader([header], delimiter=sep)))
+ actual = len(next(csv.reader(io.StringIO(record), delimiter=sep)))
+ if actual != expected:
+ raise ValueError(f'Expected {expected} fields, got {actual}.')
+ return pd.read_csv(
+ io.StringIO(header + record), sep=sep, comment=comment, dtype=dtypes)
+
+ def parse_file(lines, proxy):
+ dtypes = proxy.dtypes.drop(filename_column, errors='ignore').to_dict()
+ header = None
+ for record in _join_quoted_lines(lines):
+ if not record.strip() or (comment and record.startswith(comment)):
+ continue
+ if header is None:
+ header = record
+ else:
+ yield record, functools.partial(parse_record, header, record, dtypes)
+
+ return _ReadFilesWithErrors(
+ path, sample, parse_file, filename_column, error_handling)
+
+
+def read_from_json(
+ path: str, *, error_handling: Optional[Mapping[str, Any]] = None,
**kwargs):
+ """Reads newline-delimited JSON files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromJson``.
+
+ With ``error_handling``, each line is parsed on its own, and lines
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw line.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit lines that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromJson(path, **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromJson does not support {sorted(kwargs)} with error_handling.')
+
+ def read_json(text):
+ return pd.read_json(
+ io.StringIO(text), orient='records', lines=True, dtype=False)
+
+ def sample(handle):
+ good_lines = []
+ for line in itertools.islice(handle, _SAMPLE_SIZE):
+ try:
+ if isinstance(json.loads(line), dict):
+ good_lines.append(line)
+ except ValueError:
+ pass
+ return read_json(''.join(good_lines))
+
+ def parse_record(line, proxy):
+ return read_json(line).reindex(columns=proxy.columns).astype(proxy.dtypes)
+
+ def parse_file(lines, proxy):
+ for line in lines:
+ if line.strip():
+ yield line, functools.partial(parse_record, line, proxy)
+
+ return _ReadFilesWithErrors(path, sample, parse_file, None, error_handling)
+
+
+# The number of records used to infer column types when error_handling is set.
+_SAMPLE_SIZE = 100
+
+
+def _join_quoted_lines(lines):
+ """Joins lines so that a quoted csv field with newlines stays in one record.
+
+ A record is complete when it has an even number of quote characters.
+ """
+ record = ''
+ for line in lines:
+ record += line
+ if record.count('"') % 2 == 0:
+ yield record
+ record = ''
+ if record:
+ yield record
+
+
+# The type of the error rows from _ReadFilesWithErrors.
+_FILE_ERROR_TYPE = RowTypeConstraint.from_fields([
+ (
+ 'element',
+ RowTypeConstraint.from_fields([('path', str), ('record', str)])),
+ ('msg', str),
+ ('stack', str),
+])
+
+
+class _ReadFilesWithErrors(beam.PTransform):
+ """Reads files record by record, and sends bad records to an error output.
+
+ Args:
+ path: The file pattern to read.
+ sample: A function that takes a text file handle and returns a
+ pandas DataFrame. Its column types are the output column types.
+ parse_file: A function that takes the lines of one file and the sample.
+ It yields (record, parse) pairs. Calling parse() returns a DataFrame
+ with zero or one rows, or raises an exception for a bad record.
+ filename_column: If set, a column to add with the source file path.
+ error_handling: The error_handling config.
+ """
+ def __init__(self, path, sample, parse_file, filename_column,
error_handling):
+ self._path = path
+ self._sample = sample
+ self._parse_file = parse_file
+ self._filename_column = filename_column
+ self._error_output = yaml_utils.SafeLineLoader.strip_metadata(
+ error_handling)['output']
+
+ def expand(self, root):
+ match = beam_io.filesystems.FileSystems.match([self._path], limits=[1])[0]
+ if not match.metadata_list:
+ raise FileNotFoundError(f'Found no files that match {self._path!r}')
+ with beam_io.filesystems.FileSystems.open(
+ match.metadata_list[0].path) as handle:
+ proxy = self._sample(
+ io.TextIOWrapper(handle, encoding='utf-8', errors='replace'))[:0]
+ if self._filename_column:
+ proxy[self._filename_column] = ''
Review Comment:
**🟠High**: `ReadViaPandas` (used when `error_handling` is not set) converts
`object` columns in the inferred DataFrame to `pd.StringDtype()`
(`objects_as_strings=True` by default in `apache_beam/dataframe/io.py`).
Without converting `object` columns (including `self._filename_column`) to
`pd.StringDtype()`, `UnbatchPandas(proxy)` maps every string column to `Any`
instead of `Optional[str]` in the resulting Beam schema (in pandas < 3.0),
breaking downstream schema-sensitive transforms such as `Sql` or
`WriteToBigQuery`.
Convert `object` dtype columns in `proxy` to `pd.StringDtype()` after adding
`self._filename_column`:
```python
if self._filename_column:
proxy[self._filename_column] = ''
for col, dtype in zip(proxy.columns, proxy.dtypes):
if dtype == object:
proxy[col] = proxy[col].astype(pd.StringDtype())
```
##########
sdks/python/apache_beam/yaml/yaml_io.py:
##########
@@ -99,6 +105,234 @@ def write_to_text(pcoll, path: str):
lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)
+def read_from_csv(
+ path: str,
+ *,
+ sep: str = ',',
+ comment: Optional[str] = None,
+ filename_column: Optional[str] = None,
+ error_handling: Optional[Mapping[str, Any]] = None,
+ **kwargs):
+ """Reads comma-separated values (csv) files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromCsv``.
+
+ With ``error_handling``, each record is parsed on its own, and records
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw record text.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ sep (str): The field delimiter.
+ comment (str): Lines that start with this character are skipped.
+ filename_column (str): If set, the name of a column to add to each
+ record, containing the path of the source file.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit records that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromCsv(
+ path,
+ sep=sep,
+ comment=comment,
+ filename_column=filename_column,
+ **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromCsv does not support {sorted(kwargs)} with error_handling.')
+
+ def sample(handle):
+ return pd.read_csv(
+ handle,
+ sep=sep,
+ comment=comment,
+ nrows=_SAMPLE_SIZE,
+ on_bad_lines='skip')
+
+ def parse_record(header, record, dtypes):
+ expected = len(next(csv.reader([header], delimiter=sep)))
+ actual = len(next(csv.reader(io.StringIO(record), delimiter=sep)))
+ if actual != expected:
+ raise ValueError(f'Expected {expected} fields, got {actual}.')
+ return pd.read_csv(
+ io.StringIO(header + record), sep=sep, comment=comment, dtype=dtypes)
+
+ def parse_file(lines, proxy):
+ dtypes = proxy.dtypes.drop(filename_column, errors='ignore').to_dict()
+ header = None
+ for record in _join_quoted_lines(lines):
+ if not record.strip() or (comment and record.startswith(comment)):
+ continue
+ if header is None:
+ header = record
+ else:
+ yield record, functools.partial(parse_record, header, record, dtypes)
+
+ return _ReadFilesWithErrors(
+ path, sample, parse_file, filename_column, error_handling)
+
+
+def read_from_json(
+ path: str, *, error_handling: Optional[Mapping[str, Any]] = None,
**kwargs):
+ """Reads newline-delimited JSON files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromJson``.
+
+ With ``error_handling``, each line is parsed on its own, and lines
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw line.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit lines that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromJson(path, **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromJson does not support {sorted(kwargs)} with error_handling.')
+
+ def read_json(text):
+ return pd.read_json(
+ io.StringIO(text), orient='records', lines=True, dtype=False)
+
+ def sample(handle):
+ good_lines = []
+ for line in itertools.islice(handle, _SAMPLE_SIZE):
+ try:
+ if isinstance(json.loads(line), dict):
+ good_lines.append(line)
+ except ValueError:
+ pass
+ return read_json(''.join(good_lines))
+
+ def parse_record(line, proxy):
+ return read_json(line).reindex(columns=proxy.columns).astype(proxy.dtypes)
Review Comment:
**🟠High**: Unlike `sample()`, `parse_record` does not verify that `line` is
a JSON object (`dict`).
1. If a line is a JSON array of objects (e.g. `[{"name": "a", "count": 1},
{"name": "b", "count": 2}]`), `pd.read_json(..., orient='records', lines=True)`
parses it without error and emits multiple rows for a single malformed line.
2. If a JSON object contains only unknown keys (e.g. `{"unexpected": 123}`)
and all columns in `proxy` are nullable/string/float,
`.reindex(columns=proxy.columns)` silently drops all keys from the record and
emits an all-`None`/`NaN` row instead of routing the record to the error output.
Consider validating `isinstance(json.loads(line), dict)` (and checking for
unknown fields or empty intersection with `proxy.columns`) before calling
`read_json(line)`.
##########
sdks/python/apache_beam/yaml/yaml_io_test.py:
##########
@@ -1060,6 +1061,99 @@ def test_read_from_kafka_avro_schema_expansion(self):
''')
+class YamlFileErrorHandlingTest(unittest.TestCase):
+ def setUp(self):
+ self.temp_dir = tempfile.mkdtemp()
Review Comment:
**🟢 Low**: `self.temp_dir` created with `tempfile.mkdtemp()` is never
cleaned up. Add a `tearDown` method (or `self.addCleanup(shutil.rmtree,
self.temp_dir)`) to remove the temporary directory after each test.
```suggestion
self.temp_dir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.temp_dir)
```
##########
sdks/python/apache_beam/yaml/yaml_io.py:
##########
@@ -99,6 +105,234 @@ def write_to_text(pcoll, path: str):
lambda x: str(getattr(x, sole_field_name))) | beam.io.WriteToText(path)
+def read_from_csv(
+ path: str,
+ *,
+ sep: str = ',',
+ comment: Optional[str] = None,
+ filename_column: Optional[str] = None,
+ error_handling: Optional[Mapping[str, Any]] = None,
+ **kwargs):
+ """Reads comma-separated values (csv) files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromCsv``.
+
+ With ``error_handling``, each record is parsed on its own, and records
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw record text.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ sep (str): The field delimiter.
+ comment (str): Lines that start with this character are skipped.
+ filename_column (str): If set, the name of a column to add to each
+ record, containing the path of the source file.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit records that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromCsv(
+ path,
+ sep=sep,
+ comment=comment,
+ filename_column=filename_column,
+ **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromCsv does not support {sorted(kwargs)} with error_handling.')
+
+ def sample(handle):
+ return pd.read_csv(
+ handle,
+ sep=sep,
+ comment=comment,
+ nrows=_SAMPLE_SIZE,
+ on_bad_lines='skip')
+
+ def parse_record(header, record, dtypes):
+ expected = len(next(csv.reader([header], delimiter=sep)))
+ actual = len(next(csv.reader(io.StringIO(record), delimiter=sep)))
+ if actual != expected:
+ raise ValueError(f'Expected {expected} fields, got {actual}.')
+ return pd.read_csv(
+ io.StringIO(header + record), sep=sep, comment=comment, dtype=dtypes)
+
+ def parse_file(lines, proxy):
+ dtypes = proxy.dtypes.drop(filename_column, errors='ignore').to_dict()
+ header = None
+ for record in _join_quoted_lines(lines):
+ if not record.strip() or (comment and record.startswith(comment)):
+ continue
+ if header is None:
+ header = record
+ else:
+ yield record, functools.partial(parse_record, header, record, dtypes)
+
+ return _ReadFilesWithErrors(
+ path, sample, parse_file, filename_column, error_handling)
+
+
+def read_from_json(
+ path: str, *, error_handling: Optional[Mapping[str, Any]] = None,
**kwargs):
+ """Reads newline-delimited JSON files into a PCollection of rows.
+
+ Without ``error_handling``, this is ``apache_beam.io.ReadFromJson``.
+
+ With ``error_handling``, each line is parsed on its own, and lines
+ that fail to parse go to the error output instead of failing the pipeline.
+ Each error row has the source file path and the raw line.
+ The column types come from the first 100 records of the first file.
+
+ Args:
+ path (str): The file path to read from. The path can contain glob
+ characters such as ``*`` and ``?``.
+ error_handling: If specified, should be a mapping giving an output into
+ which to emit lines that failed to parse, as described at
+ https://beam.apache.org/documentation/sdks/yaml-errors/
+ """
+ if not error_handling:
+ return beam_io.ReadFromJson(path, **kwargs)
+ if kwargs:
+ raise ValueError(
+ f'ReadFromJson does not support {sorted(kwargs)} with error_handling.')
+
+ def read_json(text):
+ return pd.read_json(
+ io.StringIO(text), orient='records', lines=True, dtype=False)
+
+ def sample(handle):
+ good_lines = []
+ for line in itertools.islice(handle, _SAMPLE_SIZE):
+ try:
+ if isinstance(json.loads(line), dict):
+ good_lines.append(line)
+ except ValueError:
+ pass
+ return read_json(''.join(good_lines))
+
+ def parse_record(line, proxy):
+ return read_json(line).reindex(columns=proxy.columns).astype(proxy.dtypes)
+
+ def parse_file(lines, proxy):
+ for line in lines:
+ if line.strip():
+ yield line, functools.partial(parse_record, line, proxy)
+
+ return _ReadFilesWithErrors(path, sample, parse_file, None, error_handling)
+
+
+# The number of records used to infer column types when error_handling is set.
+_SAMPLE_SIZE = 100
+
+
+def _join_quoted_lines(lines):
+ """Joins lines so that a quoted csv field with newlines stays in one record.
+
+ A record is complete when it has an even number of quote characters.
+ """
+ record = ''
+ for line in lines:
+ record += line
+ if record.count('"') % 2 == 0:
+ yield record
+ record = ''
+ if record:
+ yield record
+
+
+# The type of the error rows from _ReadFilesWithErrors.
+_FILE_ERROR_TYPE = RowTypeConstraint.from_fields([
+ (
+ 'element',
+ RowTypeConstraint.from_fields([('path', str), ('record', str)])),
+ ('msg', str),
+ ('stack', str),
+])
+
+
+class _ReadFilesWithErrors(beam.PTransform):
+ """Reads files record by record, and sends bad records to an error output.
+
+ Args:
+ path: The file pattern to read.
+ sample: A function that takes a text file handle and returns a
+ pandas DataFrame. Its column types are the output column types.
+ parse_file: A function that takes the lines of one file and the sample.
+ It yields (record, parse) pairs. Calling parse() returns a DataFrame
+ with zero or one rows, or raises an exception for a bad record.
+ filename_column: If set, a column to add with the source file path.
+ error_handling: The error_handling config.
+ """
+ def __init__(self, path, sample, parse_file, filename_column,
error_handling):
+ self._path = path
+ self._sample = sample
+ self._parse_file = parse_file
+ self._filename_column = filename_column
+ self._error_output = yaml_utils.SafeLineLoader.strip_metadata(
+ error_handling)['output']
+
+ def expand(self, root):
+ match = beam_io.filesystems.FileSystems.match([self._path], limits=[1])[0]
+ if not match.metadata_list:
+ raise FileNotFoundError(f'Found no files that match {self._path!r}')
+ with beam_io.filesystems.FileSystems.open(
+ match.metadata_list[0].path) as handle:
+ proxy = self._sample(
+ io.TextIOWrapper(handle, encoding='utf-8', errors='replace'))[:0]
+ if self._filename_column:
+ proxy[self._filename_column] = ''
+
+ results = (
+ root
+ | beam_io.fileio.MatchFiles(self._path)
+ | beam_io.fileio.ReadMatches()
Review Comment:
**🟡 Medium**: `MatchFiles` outputs all matched `FileMetadata` elements in a
single bundle. Without a `beam.Reshuffle()` before `ReadMatches()` (as done in
`_ReadFromPandas` in `apache_beam/dataframe/io.py`), runners will fuse
`MatchFiles`, `ReadMatches`, and `ParDo(self._read_file)`, causing all matched
files in a glob to be read sequentially on a single worker thread.
```suggestion
| beam.Reshuffle()
| beam_io.fileio.ReadMatches()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
