From bfdb6962a4cd7374194a5051329f65fe31f68b37 Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Mon, 27 Apr 2026 08:50:27 -0600 Subject: [PATCH 1/3] add webshart format exporter --- README.md | 3 +- src/caption_flow/cli.py | 12 ++++- src/caption_flow/processors/webdataset.py | 34 ++++++++---- src/caption_flow/storage/exporter.py | 66 ++++++++++++++++++++++- tests/test_exporter.py | 34 ++++++++++++ tests/test_webdataset_ranges.py | 42 +++++++-------- 6 files changed, 158 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 76759d2..aabccda 100644 --- a/README.md +++ b/README.md @@ -76,13 +76,14 @@ Usage: caption-flow export [OPTIONS] Export caption data to various formats. Options: - --format [jsonl|json|csv|txt|huggingface_hub|all] Export format (default: jsonl) + --format [jsonl|json|csv|txt|parquet|webshart|lance|huggingface_hub|all] Export format (default: jsonl) ``` * **jsonl**: create JSON line file in the specified `--output` path * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions +* **webshart**: writes captions into existing webshart shard metadata JSON files under the plural `captions` key * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary * **all**: creates all export formats in a specified `--output` directory diff --git a/src/caption_flow/cli.py b/src/caption_flow/cli.py index e127bc6..34cf687 100644 --- a/src/caption_flow/cli.py +++ b/src/caption_flow/cli.py @@ -1498,7 +1498,17 @@ async def _run_export_process( @click.option( "--format", type=click.Choice( - ["jsonl", "json", "csv", "txt", "parquet", "lance", "huggingface_hub", "all"], + [ + "jsonl", + "json", + "csv", + "txt", + "parquet", + "webshart", + "lance", + "huggingface_hub", + "all", + ], case_sensitive=False, ), default="jsonl", diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index 788f997..1e81e76 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -227,7 +227,7 @@ def _create_units_background(self) -> None: continue shard_name = shard_info["name"] - shard_files = shard_info["num_files"] + shard_files = shard_info.get("num_samples", shard_info["num_files"]) # Check if we need to move to next shard if current_file_idx >= shard_files: @@ -681,17 +681,18 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict # Use webshart to process unprocessed ranges for start_idx, end_idx in unprocessed_ranges: try: - # Jump to shard and starting position - if shard_idx is not None: - self.loader.shard(shard_idx=shard_idx, cursor_idx=start_idx) - else: - # Try to find shard by name - self.loader.shard(filename=shard_name, cursor_idx=start_idx) - # Iterate through the range for idx in range(start_idx, end_idx + 1): try: - entry = webshart.next_with_cache_wait(self.loader) + if shard_idx is not None and hasattr(self.loader, "load_sample"): + entry = self.loader.load_sample(shard_idx, idx) + else: + # Fallback for older webshart versions. + if shard_idx is not None: + self.loader.shard(shard_idx=shard_idx, cursor_idx=idx) + else: + self.loader.shard(filename=shard_name, cursor_idx=idx) + entry = webshart.next_with_cache_wait(self.loader) # Decode image image = None @@ -723,6 +724,7 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict shard_id=shard_name, chunk_id=str(chunk_index), sample_id=str(idx) ) job_id_str = job_id.get_sample_str() + entry_metadata = getattr(entry, "metadata", {}) or {} yield { "image": image, @@ -735,7 +737,21 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict "_job_id": job_id_str, "_filename": entry.path, "_file_size": entry.size, + "_json_path": entry_metadata.get("json_path"), "_processed_indices": processed_indices, + **{ + k: v + for k, v in entry_metadata.items() + if k + not in { + "path", + "offset", + "size", + "width", + "height", + "aspect", + } + }, }, "job_id": job_id_str, } diff --git a/src/caption_flow/storage/exporter.py b/src/caption_flow/storage/exporter.py index fd5ebfe..125e594 100644 --- a/src/caption_flow/storage/exporter.py +++ b/src/caption_flow/storage/exporter.py @@ -47,7 +47,7 @@ async def export_shard( Args: ---- shard_name: Name of the shard to export - format: Export format ('jsonl', 'json', 'csv', 'parquet', 'txt') + format: Export format ('jsonl', 'json', 'csv', 'parquet', 'txt', 'webshart') output_path: Output file or directory path columns: Specific columns to export limit: Maximum number of rows to export @@ -80,6 +80,12 @@ async def export_shard( ) else: output_file = output_path / f"{shard_name}.{format}" + elif format == "webshart": + # webshart format updates the existing shard metadata JSON. + if output_path.suffix.lower() == ".json" and not output_path.is_dir(): + output_file = output_path + else: + output_file = output_path / f"{shard_name}.json" else: # Directory-based formats output_file = output_path / shard_name @@ -99,6 +105,12 @@ async def export_shard( kwargs.get("filename_column", "filename"), kwargs.get("export_column", "captions"), ) + elif format == "webshart": + return exporter.to_webshart_metadata( + output_file, + kwargs.get("filename_column", "filename"), + kwargs.get("export_column", "captions"), + ) else: raise ValueError(f"Unsupported format: {format}") @@ -629,3 +641,55 @@ def to_txt( logger.info(f"Created {files_created} text files in: {output_dir}") return files_created + + def to_webshart_metadata( + self, + metadata_path: Union[str, Path], + filename_column: str = "filename", + export_column: str = "captions", + ) -> int: + """Store captions in an existing webshart metadata JSON file.""" + if export_column not in self.contents.columns: + if export_column not in self.contents.output_fields: + raise ExportError(f"Column '{export_column}' not found in data") + + captions_by_sample = {} + skipped_no_filename = 0 + skipped_no_content = 0 + + for row in self.contents.rows: + filename = self._get_filename_from_row(row, filename_column) or row.get("item_key") + if not filename: + skipped_no_filename += 1 + continue + + content = row.get(export_column) + if content is None: + skipped_no_content += 1 + continue + + captions_by_sample[str(filename)] = content + + if skipped_no_filename > 0: + logger.warning(f"Skipped {skipped_no_filename} rows with no extractable filename") + if skipped_no_content > 0: + logger.warning(f"Skipped {skipped_no_content} rows with no {export_column} content") + + if not captions_by_sample: + return 0 + + try: + import webshart + except ImportError as exc: + raise ExportError("webshart is required for webshart metadata export") from exc + + if not hasattr(webshart, "write_captions_to_metadata"): + raise ExportError( + "Installed webshart does not support write_captions_to_metadata; " + "upgrade webshart to export captions into metadata listings." + ) + + metadata_path = Path(metadata_path) + updated = webshart.write_captions_to_metadata(metadata_path, captions_by_sample) + logger.info(f"Updated {updated} captions in webshart metadata: {metadata_path}") + return updated diff --git a/tests/test_exporter.py b/tests/test_exporter.py index 8e874c4..74d6881 100644 --- a/tests/test_exporter.py +++ b/tests/test_exporter.py @@ -3,7 +3,9 @@ import csv import json import logging +import sys import tempfile +import types from datetime import datetime from pathlib import Path @@ -438,6 +440,38 @@ def test_to_csv_export(self, sample_storage_contents, temp_storage_dir): assert len(rows) == 2 assert "job_id" in rows[0] + def test_to_webshart_metadata_uses_webshart_api( + self, sample_storage_contents, temp_storage_dir, monkeypatch + ): + """Test webshart export delegates metadata writes to the webshart API.""" + calls = [] + + def write_captions_to_metadata(metadata_path, captions_by_sample): + calls.append((Path(metadata_path), captions_by_sample)) + return len(captions_by_sample) + + monkeypatch.setitem( + sys.modules, + "webshart", + types.SimpleNamespace(write_captions_to_metadata=write_captions_to_metadata), + ) + + exporter = StorageExporter(sample_storage_contents) + metadata_path = temp_storage_dir / "shard-0000.json" + + count = exporter.to_webshart_metadata(metadata_path) + + assert count == 2 + assert calls == [ + ( + metadata_path, + { + "image1.jpg": ["Test caption 1", "Test caption 2"], + "image2.png": ["Another test caption"], + }, + ) + ] + def test_empty_contents_handling(self): """Test handling of empty storage contents.""" empty_contents = StorageContents( diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index 60a28b2..264dde9 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -849,13 +849,14 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): mock_entry.data = b"fake_image_data" mock_entry.path = "image_005.jpg" mock_entry.size = 1024 + mock_entry.metadata = { + "captions": "existing caption", + "json_path": "image_005.json", + "json_metadata": {"caption": "existing caption"}, + } # Mock the loader methods - worker_processor_real.loader.shard = Mock() - - # Create a simple iterator that returns our mock entry - def mock_next_with_cache_wait(loader): - return mock_entry + worker_processor_real.loader.load_sample = Mock(side_effect=[mock_entry] * 3) unit = WorkUnit( unit_id="shard_0:chunk:0", @@ -874,21 +875,17 @@ def mock_next_with_cache_wait(loader): # Mock the image decoding test_image = Image.new("RGB", (100, 100), color="red") - with patch( - "caption_flow.processors.webdataset.webshart.next_with_cache_wait", - side_effect=[mock_entry] * 3, - ): - with patch("caption_flow.processors.webdataset.cv2.imdecode") as mock_decode: - with patch("caption_flow.processors.webdataset.cv2.cvtColor") as mock_convert: - with patch( - "caption_flow.processors.webdataset.Image.fromarray", - return_value=test_image, - ): - # Mock cv2 processing chain - mock_decode.return_value = "fake_cv2_image" - mock_convert.return_value = "fake_rgb_array" - - results = list(worker_processor_real.process_unit(unit, {})) + with patch("caption_flow.processors.webdataset.cv2.imdecode") as mock_decode: + with patch("caption_flow.processors.webdataset.cv2.cvtColor") as mock_convert: + with patch( + "caption_flow.processors.webdataset.Image.fromarray", + return_value=test_image, + ): + # Mock cv2 processing chain + mock_decode.return_value = "fake_cv2_image" + mock_convert.return_value = "fake_rgb_array" + + results = list(worker_processor_real.process_unit(unit, {})) assert len(results) == 3 @@ -900,10 +897,13 @@ def mock_next_with_cache_wait(loader): assert result["image_data"] == b"fake_image_data" assert result["metadata"]["_filename"] == "image_005.jpg" assert result["metadata"]["_file_size"] == 1024 + assert result["metadata"]["captions"] == "existing caption" + assert result["metadata"]["json_metadata"] == {"caption": "existing caption"} + assert result["metadata"]["_json_path"] == "image_005.json" assert not result["metadata"].get("_mock", False) # Should not have mock flag # Verify loader was called correctly - worker_processor_real.loader.shard.assert_called_once_with(shard_idx=0, cursor_idx=5) + worker_processor_real.loader.load_sample.assert_any_call(0, 5) def test_process_unit_real_mode_shard_by_name(self, worker_processor_real): """Test processing when shard_idx is None (fallback to name lookup).""" From 297ea1fa2421e4f3e069b1d74183032c32715c65 Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Mon, 27 Apr 2026 09:30:17 -0600 Subject: [PATCH 2/3] address issues --- README.md | 6 ++++-- src/caption_flow/processors/webdataset.py | 19 +++++++++++++------ src/caption_flow/storage/exporter.py | 21 ++++++++++++++++++--- tests/test_exporter.py | 13 ++++++++++++- tests/test_webdataset_ranges.py | 1 + 5 files changed, 48 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index aabccda..7d45731 100644 --- a/README.md +++ b/README.md @@ -83,9 +83,11 @@ Options: * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions -* **webshart**: writes captions into existing webshart shard metadata JSON files under the plural `captions` key +* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. for this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. if you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard. * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary -* **all**: creates all export formats in a specified `--output` directory +* **all**: creates the directory/file-generating export formats in a specified `--output` directory. prefer a directory here; `webshart` is a special case that expects existing per-shard metadata `.json` files rather than creating new metadata files. + +> note: `--output` paths ending in `.json` are treated specially for `webshart`. use a directory for normal multi-format exports and an existing shard metadata JSON file only when intentionally updating a `webshart` shard. --- diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index 1e81e76..1d26c78 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -681,17 +681,21 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict # Use webshart to process unprocessed ranges for start_idx, end_idx in unprocessed_ranges: try: + use_sample_loader = shard_idx is not None and hasattr(self.loader, "load_sample") + if not use_sample_loader: + # Fallback for older webshart versions. Seek once per contiguous range, + # then advance with next_with_cache_wait for each item. + if shard_idx is not None: + self.loader.shard(shard_idx=shard_idx, cursor_idx=start_idx) + else: + self.loader.shard(filename=shard_name, cursor_idx=start_idx) + # Iterate through the range for idx in range(start_idx, end_idx + 1): try: - if shard_idx is not None and hasattr(self.loader, "load_sample"): + if use_sample_loader: entry = self.loader.load_sample(shard_idx, idx) else: - # Fallback for older webshart versions. - if shard_idx is not None: - self.loader.shard(shard_idx=shard_idx, cursor_idx=idx) - else: - self.loader.shard(filename=shard_name, cursor_idx=idx) entry = webshart.next_with_cache_wait(self.loader) # Decode image @@ -725,6 +729,8 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict ) job_id_str = job_id.get_sample_str() entry_metadata = getattr(entry, "metadata", {}) or {} + if not isinstance(entry_metadata, dict): + entry_metadata = {} yield { "image": image, @@ -750,6 +756,7 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict "width", "height", "aspect", + "json_path", } }, }, diff --git a/src/caption_flow/storage/exporter.py b/src/caption_flow/storage/exporter.py index 125e594..0b92aca 100644 --- a/src/caption_flow/storage/exporter.py +++ b/src/caption_flow/storage/exporter.py @@ -81,9 +81,13 @@ async def export_shard( else: output_file = output_path / f"{shard_name}.{format}" elif format == "webshart": - # webshart format updates the existing shard metadata JSON. + # webshart metadata must be shard-specific to avoid multiple shards + # reusing the same JSON file when called repeatedly. if output_path.suffix.lower() == ".json" and not output_path.is_dir(): - output_file = output_path + if output_path.stem == shard_name: + output_file = output_path + else: + output_file = output_path.parent / f"{shard_name}.json" else: output_file = output_path / f"{shard_name}.json" else: @@ -678,6 +682,18 @@ def to_webshart_metadata( if not captions_by_sample: return 0 + metadata_path = Path(metadata_path) + if not metadata_path.exists(): + raise ExportError( + f"Webshart metadata file does not exist: {metadata_path}. " + "Expected an existing metadata JSON file to update." + ) + if not metadata_path.is_file(): + raise ExportError( + f"Webshart metadata path is not a file: {metadata_path}. " + "Expected an existing metadata JSON file to update." + ) + try: import webshart except ImportError as exc: @@ -689,7 +705,6 @@ def to_webshart_metadata( "upgrade webshart to export captions into metadata listings." ) - metadata_path = Path(metadata_path) updated = webshart.write_captions_to_metadata(metadata_path, captions_by_sample) logger.info(f"Updated {updated} captions in webshart metadata: {metadata_path}") return updated diff --git a/tests/test_exporter.py b/tests/test_exporter.py index 74d6881..7d3ca0b 100644 --- a/tests/test_exporter.py +++ b/tests/test_exporter.py @@ -15,7 +15,7 @@ import pytest_asyncio from caption_flow.models import Caption, StorageContents from caption_flow.storage import StorageManager -from caption_flow.storage.exporter import LanceStorageExporter, StorageExporter +from caption_flow.storage.exporter import ExportError, LanceStorageExporter, StorageExporter # Set up logging to avoid logger not defined errors logging.basicConfig(level=logging.INFO) @@ -458,6 +458,7 @@ def write_captions_to_metadata(metadata_path, captions_by_sample): exporter = StorageExporter(sample_storage_contents) metadata_path = temp_storage_dir / "shard-0000.json" + metadata_path.write_text('{"files": {}}', encoding="utf-8") count = exporter.to_webshart_metadata(metadata_path) @@ -472,6 +473,16 @@ def write_captions_to_metadata(metadata_path, captions_by_sample): ) ] + def test_to_webshart_metadata_requires_existing_file( + self, sample_storage_contents, temp_storage_dir + ): + """Test webshart export fails clearly when metadata JSON is missing.""" + exporter = StorageExporter(sample_storage_contents) + metadata_path = temp_storage_dir / "missing.json" + + with pytest.raises(ExportError, match="does not exist"): + exporter.to_webshart_metadata(metadata_path) + def test_empty_contents_handling(self): """Test handling of empty storage contents.""" empty_contents = StorageContents( diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index 264dde9..e3dbec3 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -900,6 +900,7 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): assert result["metadata"]["captions"] == "existing caption" assert result["metadata"]["json_metadata"] == {"caption": "existing caption"} assert result["metadata"]["_json_path"] == "image_005.json" + assert "json_path" not in result["metadata"] assert not result["metadata"].get("_mock", False) # Should not have mock flag # Verify loader was called correctly From 0b835a5c5b166c2e2fe0f6319f31041a8da004a3 Mon Sep 17 00:00:00 2001 From: bghira <59658056+bghira@users.noreply.github.com> Date: Mon, 27 Apr 2026 09:53:16 -0600 Subject: [PATCH 3/3] address issues --- src/caption_flow/processors/webdataset.py | 30 ++++---- src/caption_flow/storage/exporter.py | 31 +++++++- tests/test_exporter.py | 86 +++++++++++++++++++++++ tests/test_webdataset_ranges.py | 4 ++ 4 files changed, 135 insertions(+), 16 deletions(-) diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index 1d26c78..1e12135 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -731,6 +731,21 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict entry_metadata = getattr(entry, "metadata", {}) or {} if not isinstance(entry_metadata, dict): entry_metadata = {} + filtered_entry_metadata = { + k: v + for k, v in entry_metadata.items() + if not k.startswith("_") + and k + not in { + "path", + "offset", + "size", + "width", + "height", + "aspect", + "json_path", + } + } yield { "image": image, @@ -738,6 +753,7 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict "item_key": Path(entry.path).stem, "item_index": idx, "metadata": { + **filtered_entry_metadata, "_item_index": idx, "_chunk_relative_index": idx - unit.data["start_index"], "_job_id": job_id_str, @@ -745,20 +761,6 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict "_file_size": entry.size, "_json_path": entry_metadata.get("json_path"), "_processed_indices": processed_indices, - **{ - k: v - for k, v in entry_metadata.items() - if k - not in { - "path", - "offset", - "size", - "width", - "height", - "aspect", - "json_path", - } - }, }, "job_id": job_id_str, } diff --git a/src/caption_flow/storage/exporter.py b/src/caption_flow/storage/exporter.py index 0b92aca..89d32f4 100644 --- a/src/caption_flow/storage/exporter.py +++ b/src/caption_flow/storage/exporter.py @@ -87,7 +87,11 @@ async def export_shard( if output_path.stem == shard_name: output_file = output_path else: - output_file = output_path.parent / f"{shard_name}.json" + raise ExportError( + "Invalid webshart output path " + f"'{output_path}': explicit JSON output files must be named " + f"'{shard_name}.json' for shard '{shard_name}'." + ) else: output_file = output_path / f"{shard_name}.json" else: @@ -646,6 +650,24 @@ def to_txt( logger.info(f"Created {files_created} text files in: {output_dir}") return files_created + def _normalize_webshart_captions(self, value: Any, export_column: str) -> List[str]: + """Normalize a caption export value to webshart's plural captions list.""" + if isinstance(value, str): + return [value] if value else [] + if isinstance(value, (list, tuple)): + captions = [] + for item in value: + if item is None: + continue + text = item if isinstance(item, str) else str(self._serialize_value(item)) + if text: + captions.append(text) + return captions + raise ExportError( + f"Column '{export_column}' must contain a string or list of strings " + "for webshart metadata export" + ) + def to_webshart_metadata( self, metadata_path: Union[str, Path], @@ -672,7 +694,12 @@ def to_webshart_metadata( skipped_no_content += 1 continue - captions_by_sample[str(filename)] = content + captions = self._normalize_webshart_captions(content, export_column) + if not captions: + skipped_no_content += 1 + continue + + captions_by_sample[str(filename)] = captions if skipped_no_filename > 0: logger.warning(f"Skipped {skipped_no_filename} rows with no extractable filename") diff --git a/tests/test_exporter.py b/tests/test_exporter.py index 7d3ca0b..8f4309b 100644 --- a/tests/test_exporter.py +++ b/tests/test_exporter.py @@ -182,6 +182,44 @@ async def test_export_shard_csv(self, populated_storage_manager, temp_storage_di assert "job_id" in df.columns assert "filename" in df.columns + @pytest.mark.asyncio + async def test_export_shard_webshart_directory_output( + self, populated_storage_manager, temp_storage_dir, monkeypatch + ): + """Test per-shard webshart export uses existing metadata in output directory.""" + calls = [] + + def write_captions_to_metadata(metadata_path, captions_by_sample): + calls.append((Path(metadata_path), captions_by_sample)) + return len(captions_by_sample) + + monkeypatch.setitem( + sys.modules, + "webshart", + types.SimpleNamespace(write_captions_to_metadata=write_captions_to_metadata), + ) + + metadata_path = temp_storage_dir / "default.json" + metadata_path.write_text('{"files": {}}', encoding="utf-8") + exporter = LanceStorageExporter(populated_storage_manager) + + count = await exporter.export_shard("default", "webshart", temp_storage_dir) + + assert count == 3 + assert calls + assert calls[0][0] == metadata_path + assert all(isinstance(value, list) for value in calls[0][1].values()) + + @pytest.mark.asyncio + async def test_export_shard_webshart_rejects_mismatched_json_path( + self, populated_storage_manager, temp_storage_dir + ): + """Test explicit webshart JSON path must match the shard name.""" + exporter = LanceStorageExporter(populated_storage_manager) + + with pytest.raises(ExportError, match="explicit JSON output files must be named"): + await exporter.export_shard("default", "webshart", temp_storage_dir / "wrong.json") + @pytest.mark.asyncio async def test_export_shard_json_directory(self, populated_storage_manager, temp_storage_dir): """Test exporting shard to JSON directory format.""" @@ -483,6 +521,54 @@ def test_to_webshart_metadata_requires_existing_file( with pytest.raises(ExportError, match="does not exist"): exporter.to_webshart_metadata(metadata_path) + def test_to_webshart_metadata_normalizes_string_captions( + self, sample_storage_contents, temp_storage_dir, monkeypatch + ): + """Test webshart export wraps single string captions in a list.""" + calls = [] + + def write_captions_to_metadata(metadata_path, captions_by_sample): + calls.append(captions_by_sample) + return len(captions_by_sample) + + monkeypatch.setitem( + sys.modules, + "webshart", + types.SimpleNamespace(write_captions_to_metadata=write_captions_to_metadata), + ) + + contents = StorageContents( + rows=[{"filename": "image1.jpg", "captions": "single caption"}], + columns=["filename", "captions"], + output_fields=["captions"], + total_rows=1, + metadata={}, + ) + exporter = StorageExporter(contents) + metadata_path = temp_storage_dir / "shard-0000.json" + metadata_path.write_text('{"files": {}}', encoding="utf-8") + + count = exporter.to_webshart_metadata(metadata_path) + + assert count == 1 + assert calls == [{"image1.jpg": ["single caption"]}] + + def test_to_webshart_metadata_rejects_unexpected_caption_type( + self, sample_storage_contents, temp_storage_dir + ): + """Test webshart export rejects values that are not strings or lists.""" + contents = StorageContents( + rows=[{"filename": "image1.jpg", "captions": {"bad": "shape"}}], + columns=["filename", "captions"], + output_fields=["captions"], + total_rows=1, + metadata={}, + ) + exporter = StorageExporter(contents) + + with pytest.raises(ExportError, match="string or list of strings"): + exporter.to_webshart_metadata(temp_storage_dir / "shard-0000.json") + def test_empty_contents_handling(self): """Test handling of empty storage contents.""" empty_contents = StorageContents( diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index e3dbec3..29c75db 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -853,6 +853,8 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): "captions": "existing caption", "json_path": "image_005.json", "json_metadata": {"caption": "existing caption"}, + "_filename": "wrong.jpg", + "_job_id": "wrong_job", } # Mock the loader methods @@ -896,11 +898,13 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): assert result["image"] == test_image assert result["image_data"] == b"fake_image_data" assert result["metadata"]["_filename"] == "image_005.jpg" + assert result["metadata"]["_job_id"] == result["job_id"] assert result["metadata"]["_file_size"] == 1024 assert result["metadata"]["captions"] == "existing caption" assert result["metadata"]["json_metadata"] == {"caption": "existing caption"} assert result["metadata"]["_json_path"] == "image_005.json" assert "json_path" not in result["metadata"] + assert result["metadata"]["_filename"] != "wrong.jpg" assert not result["metadata"].get("_mock", False) # Should not have mock flag # Verify loader was called correctly