From 9ae8cc6c38d9c7e691e89bce23ff3093a3eed610 Mon Sep 17 00:00:00 2001 From: "stephen.ibuoye" Date: Tue, 15 Sep 2026 00:59:58 +0000 Subject: [PATCH 1/2] Prefer diverse paths for schema inference Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../stream/default_file_based_stream.py | 51 +++++++- .../stream/test_default_file_based_stream.py | 115 +++++++++++++++++- 2 files changed, 162 insertions(+), 4 deletions(-) diff --git a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py b/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py index 588c4a18e4..9a6da249c2 100644 --- a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +++ b/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py @@ -4,6 +4,7 @@ import asyncio import itertools +import re import traceback from collections import defaultdict from copy import deepcopy @@ -307,9 +308,12 @@ def _get_raw_json_schema(self) -> JsonSchema: self.logger.warning( msg=f"Refusing to infer schema for {first_n_files} files; using {max_n_files_for_schema_inference} files." ) - first_n_files = max_n_files_for_schema_inference - - files = sorted(files, key=lambda x: x.last_modified, reverse=True)[:first_n_files] + files = self._select_files_for_schema_inference(files, max_n_files_for_schema_inference) + self.logger.info( + msg=f"Files selected for schema inference for stream {self.name}: {[file.uri for file in files]}" + ) + else: + files = sorted(files, key=lambda x: x.last_modified, reverse=True)[:first_n_files] inferred_schema = self.infer_schema(files) @@ -324,6 +328,47 @@ def _get_raw_json_schema(self) -> JsonSchema: return schema + @staticmethod + def _schema_inference_group_key(uri: str) -> str: + """ + Collapse a file URI to its path shape: directory segments containing digits (dates, ids, uuids, hashes) + become "*", and the file name is reduced to its leading alphabetic prefix plus extension. + """ + segments = uri.strip("/").split("/") + directory = ["*" if re.search(r"\d", segment) else segment for segment in segments[:-1]] + stem, ext = path.splitext(segments[-1]) + prefix_match = re.match(r"[A-Za-z]*", stem) + prefix = prefix_match.group(0) if prefix_match else "" + return "/".join(directory + [prefix + ext]) + + @staticmethod + def _select_files_for_schema_inference(files: List[RemoteFile], n: int) -> List[RemoteFile]: + """ + Pick up to n files for schema inference, newest first, preferring one file per distinct path shape + so heterogeneous file sets contribute all of their columns. Remaining slots are filled with the next-newest files. + """ + newest_first = sorted(files, key=lambda x: x.last_modified, reverse=True) + if len(newest_first) <= n: + return newest_first + selected: List[RemoteFile] = [] + seen_groups: Set[str] = set() + for file in newest_first: + group = DefaultFileBasedStream._schema_inference_group_key(file.uri) + if group in seen_groups: + continue + seen_groups.add(group) + selected.append(file) + if len(selected) == n: + return selected + selected_uris = {file.uri for file in selected} + for file in newest_first: + if len(selected) == n: + break + if file.uri not in selected_uris: + selected.append(file) + selected_uris.add(file.uri) + return selected + def get_files(self) -> Iterable[RemoteFile]: """ Return all files that belong to the stream as defined by the stream's globs. diff --git a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py b/unit_tests/sources/file_based/stream/test_default_file_based_stream.py index 54394a36d8..3c9a2f7b47 100644 --- a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py +++ b/unit_tests/sources/file_based/stream/test_default_file_based_stream.py @@ -5,7 +5,7 @@ import traceback import unittest from copy import deepcopy -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any, Iterable, Iterator, Mapping from unittest import mock from unittest.mock import Mock @@ -88,6 +88,90 @@ def test_fill_nulls(input_schema: Mapping[str, Any], expected_output: Mapping[st assert DefaultFileBasedStream._fill_nulls(input_schema) == expected_output +@pytest.mark.parametrize( + "uri, expected_group_key", + [ + ( + "sync/datasync/6723152521986048/-323232/fcb7eea7-c953-4971-a493-56dbcfa3326b/matchedEvents/Page/Page-yyeAFHOyeip34fjJKnI58_LUqEw.avro", + "sync/datasync/*/*/*/matchedEvents/Page/Page.avro", + ), + ( + "sync/datasync/6723152521986048/-323232/fcb7eea7-c953-4971-a493-56dbcfa3326b/allguides.avro", + "sync/datasync/*/*/*/allguides.avro", + ), + ("data/2024/01/05/events_20240105.csv", "data/*/*/*/events.csv"), + ("customers.csv", "customers.csv"), + ("/bucket/2024-01-01.jsonl", "bucket/.jsonl"), + ], +) +def test_schema_inference_group_key(uri: str, expected_group_key: str) -> None: + assert DefaultFileBasedStream._schema_inference_group_key(uri) == expected_group_key + + +def test_select_files_for_schema_inference_prefers_one_file_per_path_shape() -> None: + files = [ + RemoteFile( + uri=f"matchedEvents/Page/Page-{i}.avro", + last_modified=datetime(2022, 10, 22, tzinfo=timezone.utc) - timedelta(minutes=i), + ) + for i in range(30) + ] + files.extend( + RemoteFile( + uri=f"matchedEvents/Feature/Feature-{i}.avro", + last_modified=datetime(2022, 10, 22, tzinfo=timezone.utc) - timedelta(minutes=30 + i), + ) + for i in range(3) + ) + files.append( + RemoteFile( + uri="allguides.avro", + last_modified=datetime(2022, 10, 22, tzinfo=timezone.utc) - timedelta(minutes=100), + ) + ) + + selected = DefaultFileBasedStream._select_files_for_schema_inference(files, n=5) + selected_uris = [file.uri for file in selected] + + assert len(selected) == 5 + assert len(selected_uris) == len(set(selected_uris)) + assert selected_uris == [ + "matchedEvents/Page/Page-0.avro", + "matchedEvents/Feature/Feature-0.avro", + "allguides.avro", + "matchedEvents/Page/Page-1.avro", + "matchedEvents/Page/Page-2.avro", + ] + + +def test_select_files_for_schema_inference_falls_back_to_newest_when_homogeneous() -> None: + files = [ + RemoteFile( + uri=f"data/2024/01/{i:02d}/events_{i}.csv", + last_modified=datetime(2022, 10, 22, tzinfo=timezone.utc) + timedelta(minutes=i), + ) + for i in range(20) + ] + + selected = DefaultFileBasedStream._select_files_for_schema_inference(files, n=10) + + assert selected == list(reversed(files[-10:])) + + +def test_select_files_for_schema_inference_returns_all_when_under_limit() -> None: + files = [ + RemoteFile( + uri=f"file-{i}.csv", + last_modified=datetime(2022, 10, 22, tzinfo=timezone.utc) + timedelta(minutes=i), + ) + for i in range(3) + ] + + selected = DefaultFileBasedStream._select_files_for_schema_inference(files, n=10) + + assert selected == list(reversed(files)) + + class DefaultFileBasedStreamTest(unittest.TestCase): _NOW = datetime(2022, 10, 22, tzinfo=timezone.utc) _A_RECORD = {"a_record": 1} @@ -227,6 +311,35 @@ def test_override_max_n_files_for_schema_inference_is_respected(self) -> None: } assert self._parser.infer_schema.call_count == 3 + def test_when_too_many_files_then_schema_inference_covers_each_path_shape(self) -> None: + self._stream_config.input_schema = None + self._stream_config.schemaless = False + self._stream_config.use_first_found_file_for_schema_discovery = False + self._stream_config.recent_n_files_to_read_for_schema_discovery = None + self._discovery_policy.get_max_n_files_for_schema_inference.return_value = 10 + self._stream_reader.get_matching_files.return_value = [ + *[ + RemoteFile( + uri=f"matchedEvents/Page/Page-{i}.avro", + last_modified=self._NOW - timedelta(minutes=i), + ) + for i in range(11) + ], + RemoteFile( + uri="allguides.avro", + last_modified=self._NOW - timedelta(minutes=100), + ), + ] + self._stream.infer_schema = Mock(return_value={"a": {"type": "string"}}) + + schema = self._stream._get_raw_json_schema() + + self._stream.infer_schema.assert_called_once() + selected_files = self._stream.infer_schema.call_args.args[0] + assert len(selected_files) == 10 + assert "allguides.avro" in [file.uri for file in selected_files] + assert schema == {"type": "object", "properties": {"a": {"type": "string"}}} + def test_use_first_found_file_for_schema_discovery(self) -> None: self._stream.config.use_first_found_file_for_schema_discovery = True From 14b8466abcb596e0529682a17fea11cc6b4ea2d4 Mon Sep 17 00:00:00 2001 From: "stephen.ibuoye" Date: Tue, 15 Sep 2026 01:06:58 +0000 Subject: [PATCH 2/2] Restrict diverse selection to the configured recent-file window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../stream/default_file_based_stream.py | 8 +++-- .../stream/test_default_file_based_stream.py | 31 +++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py b/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py index 9a6da249c2..9fdd009b7a 100644 --- a/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py +++ b/airbyte_cdk/sources/file_based/stream/default_file_based_stream.py @@ -303,17 +303,21 @@ def _get_raw_json_schema(self) -> JsonSchema: self._discovery_policy.get_max_n_files_for_schema_inference(self.get_parser()) ) + newest_first = sorted(files, key=lambda x: x.last_modified, reverse=True)[:first_n_files] + if first_n_files > max_n_files_for_schema_inference: # Use the most recent files for schema inference, so we pick up schema changes during discovery. self.logger.warning( msg=f"Refusing to infer schema for {first_n_files} files; using {max_n_files_for_schema_inference} files." ) - files = self._select_files_for_schema_inference(files, max_n_files_for_schema_inference) + files = self._select_files_for_schema_inference( + newest_first, max_n_files_for_schema_inference + ) self.logger.info( msg=f"Files selected for schema inference for stream {self.name}: {[file.uri for file in files]}" ) else: - files = sorted(files, key=lambda x: x.last_modified, reverse=True)[:first_n_files] + files = newest_first inferred_schema = self.infer_schema(files) diff --git a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py b/unit_tests/sources/file_based/stream/test_default_file_based_stream.py index 3c9a2f7b47..5bf11c5b3e 100644 --- a/unit_tests/sources/file_based/stream/test_default_file_based_stream.py +++ b/unit_tests/sources/file_based/stream/test_default_file_based_stream.py @@ -340,6 +340,37 @@ def test_when_too_many_files_then_schema_inference_covers_each_path_shape(self) assert "allguides.avro" in [file.uri for file in selected_files] assert schema == {"type": "object", "properties": {"a": {"type": "string"}}} + def test_when_recent_n_files_exceeds_cap_then_selection_stays_within_recent_window( + self, + ) -> None: + self._stream_config.input_schema = None + self._stream_config.schemaless = False + self._stream_config.use_first_found_file_for_schema_discovery = False + self._stream_config.recent_n_files_to_read_for_schema_discovery = 20 + self._discovery_policy.get_max_n_files_for_schema_inference.return_value = 10 + self._stream_reader.get_matching_files.return_value = [ + *[ + RemoteFile( + uri=f"matchedEvents/Page/Page-{i}.avro", + last_modified=self._NOW - timedelta(minutes=i), + ) + for i in range(25) + ], + RemoteFile( + uri="allguides.avro", + last_modified=self._NOW - timedelta(minutes=100), + ), + ] + self._stream.infer_schema = Mock(return_value={"a": {"type": "string"}}) + + self._stream._get_raw_json_schema() + + self._stream.infer_schema.assert_called_once() + selected_uris = {file.uri for file in self._stream.infer_schema.call_args.args[0]} + assert len(selected_uris) == 10 + assert "allguides.avro" not in selected_uris + assert selected_uris <= {f"matchedEvents/Page/Page-{i}.avro" for i in range(20)} + def test_use_first_found_file_for_schema_discovery(self) -> None: self._stream.config.use_first_found_file_for_schema_discovery = True