fix(iqb/pipeline): always create parquet file (#41)

It definitely feels less error prone to always create a file and default
to creating an empty file if there's no data.

Author: bassosimone

data/run_query.py

@@ -91,7 +91,7 @@ def run_bq_query(
 
     # Step 1: Execute query and save results using IQBPipeline
     # This creates: ./iqb/data/cache/v1/{start}/{end}/{query_name}/
-    #   - data.parquet: query results
+    #   - data.parquet: query results (empty file if no results)
     #   - stats.json: query metadata
     pipeline = IQBPipeline(project_id=project_id, data_dir=data_dir)
     result = pipeline.execute_query_template(
@@ -101,14 +101,6 @@ def run_bq_query(
     )
     info = result.save_parquet()
 
-    if info.no_content:
-        print("⚠ Query returned no results", file=sys.stderr)
-        if output_file:
-            output_file.write_text("[]")
-        else:
-            print("[]")
-        return
-
     print(f"✓ Parquet saved: {info.file_path}", file=sys.stderr)
 
     # Save query statistics (timing, bytes processed, template hash)
@@ -119,6 +111,11 @@ def run_bq_query(
     print("Converting parquet to JSON...", file=sys.stderr)
     table = pq.read_table(info.file_path)
     records = table.to_pylist()
+
+    # Check if query returned no results
+    if len(records) <= 0:
+        print("⚠ Query returned no results", file=sys.stderr)
+
     json_output = json.dumps(records, indent=2)
 
     # Step 3: Write JSON to output file or stdout

library/src/iqb/pipeline.py

@@ -83,6 +83,7 @@
 from pathlib import Path
 from typing import Final
 
+import pyarrow as pa
 import pyarrow.parquet as pq
 from google.cloud import bigquery, bigquery_storage_v1
 from google.cloud.bigquery import job, table
@@ -120,23 +121,17 @@ class CacheEntry:
     stats_path: Path
 
 
-# TODO(bassosimone): create an empty parquet file rather than
-# returning "no_content=True" in a subsequent diff. I realized
-# that the current approach is easy to get wrong.
-
-
 @dataclass(frozen=True)
 class ParquetFileInfo:
     """
     Result of serializing a query result into the cache using parquet.
 
+    An empty parquet file is written if the query returns no rows.
+
     Attributes:
-        no_content: true if the query returned no content, in which case
-            no parquet file is actually being written.
         file_path: full path to the written file.
     """
 
-    no_content: bool
     file_path: Path
 
 
@@ -162,25 +157,30 @@ class QueryResult:
     template_hash: str
 
     def save_parquet(self) -> ParquetFileInfo:
-        """Streams and saves the query results to data.parquet in cache_dir."""
+        """Streams and saves the query results to data.parquet in cache_dir.
+
+        If the query returns no rows, an empty parquet file is written.
+        """
         self.cache_dir.mkdir(parents=True, exist_ok=True)
         parquet_path = self.cache_dir / CACHE_DATA_FILENAME
 
+        # Note: using .as_posix to avoid paths with backslashes
+        # that can cause issues with PyArrow on Windows
+        posix_path = parquet_path.as_posix()
+
         # Access the first batch to obtain the schema
         batches = self.rows.to_arrow_iterable(bqstorage_client=self.bq_read_client)
         first_batch = next(batches, None)
-        if first_batch is None:
-            return ParquetFileInfo(no_content=True, file_path=parquet_path)
+        schema = first_batch.schema if first_batch is not None else pa.schema([])
 
-        # Note: using .as_posix to avoid paths with backslashes
-        # that can cause issues with PyArrow on Windows
-        posix_path = parquet_path.as_posix()
-        with pq.ParquetWriter(posix_path, first_batch.schema) as writer:
-            writer.write_batch(first_batch)
+        # Write the possibly-empty parquet file
+        with pq.ParquetWriter(posix_path, schema) as writer:
+            if first_batch is not None:
+                writer.write_batch(first_batch)
             for batch in batches:
                 writer.write_batch(batch)
 
-        return ParquetFileInfo(no_content=False, file_path=parquet_path)
+        return ParquetFileInfo(file_path=parquet_path)
 
     def save_stats(self) -> Path:
         """Writes query statistics to stats.json in cache_dir.

library/tests/iqb/pipeline_test.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 from unittest.mock import MagicMock, Mock, patch
 
+import pyarrow as pa
 import pytest
 
 from iqb.pipeline import (
@@ -299,38 +300,53 @@ def test_save_parquet_with_data(self, tmp_path):
 
             info = result.save_parquet()
 
-            # Verify
-            assert info.no_content is False
+            # Verify file path and directory creation
             expected_path = cache_dir / "data.parquet"
             assert info.file_path == expected_path
             assert cache_dir.exists()
             mock_writer_instance.write_batch.assert_called_once_with(mock_batch)
 
     def test_save_parquet_empty_results(self, tmp_path):
-        """Test handling of empty query results."""
+        """Test handling of empty query results - writes empty parquet file."""
         cache_dir = tmp_path / "cache"
 
         # Mock empty iterator
         mock_rows = Mock()
         mock_rows.to_arrow_iterable.return_value = iter([])
 
-        result = QueryResult(
-            bq_read_client=Mock(),
-            job=Mock(),
-            rows=mock_rows,
-            cache_dir=cache_dir,
-            query_start_time="2024-11-27T10:00:00.000000Z",
-            template_hash="abc123",
-        )
+        # Mock ParquetWriter to verify it's called with empty schema
+        with patch("iqb.pipeline.pq.ParquetWriter") as mock_writer:
+            mock_writer_instance = MagicMock()
+            mock_writer.return_value.__enter__.return_value = mock_writer_instance
 
-        info = result.save_parquet()
+            result = QueryResult(
+                bq_read_client=Mock(),
+                job=Mock(),
+                rows=mock_rows,
+                cache_dir=cache_dir,
+                query_start_time="2024-11-27T10:00:00.000000Z",
+                template_hash="abc123",
+            )
 
-        # Verify no file created, but directory exists
-        assert info.no_content is True
-        expected_path = cache_dir / "data.parquet"
-        assert info.file_path == expected_path
-        assert cache_dir.exists()
-        assert not expected_path.exists()
+            info = result.save_parquet()
+
+            # Verify empty parquet file would be created
+            expected_path = cache_dir / "data.parquet"
+            assert info.file_path == expected_path
+            assert cache_dir.exists()
+
+            # Verify ParquetWriter was called with empty schema
+            mock_writer.assert_called_once()
+            call_args = mock_writer.call_args
+            assert call_args[0][0] == expected_path.as_posix()
+
+            # Verify schema is empty (no fields)
+            schema_arg = call_args[0][1]
+            assert isinstance(schema_arg, pa.Schema)
+            assert len(schema_arg) == 0
+
+            # Verify no batches were written (first_batch is None, for loop has nothing)
+            mock_writer_instance.write_batch.assert_not_called()
 
     def test_save_parquet_multiple_batches(self, tmp_path):
         """Test saving multiple Arrow batches."""
@@ -362,7 +378,8 @@ def test_save_parquet_multiple_batches(self, tmp_path):
 
             # Verify all batches written
             assert mock_writer_instance.write_batch.call_count == 3
-            assert info.no_content is False
+            expected_path = cache_dir / "data.parquet"
+            assert info.file_path == expected_path
 
     def test_save_parquet_creates_nested_directories(self, tmp_path):
         """Test that save_parquet creates nested cache directory."""
@@ -392,7 +409,8 @@ def test_save_parquet_creates_nested_directories(self, tmp_path):
 
             # Verify cache directory created
             assert cache_dir.exists()
-            assert info.no_content is False
+            expected_path = cache_dir / "data.parquet"
+            assert info.file_path == expected_path
 
 
 class TestQueryResultSaveStats:
@@ -672,18 +690,9 @@ def test_get_cache_entry_validation_before_fs_check(self, mock_storage, mock_cli
 class TestParquetFileInfo:
     """Test ParquetFileInfo dataclass."""
 
-    def test_parquet_file_info_with_content(self, tmp_path):
-        """Test ParquetFileInfo creation with content."""
-        test_file = tmp_path / "test.parquet"
-        info = ParquetFileInfo(no_content=False, file_path=test_file)
-
-        assert info.no_content is False
-        assert info.file_path == test_file
-
-    def test_parquet_file_info_no_content(self, tmp_path):
-        """Test ParquetFileInfo creation without content."""
+    def test_parquet_file_info_creation(self, tmp_path):
+        """Test ParquetFileInfo creation."""
         test_file = tmp_path / "test.parquet"
-        info = ParquetFileInfo(no_content=True, file_path=test_file)
+        info = ParquetFileInfo(file_path=test_file)
 
-        assert info.no_content is True
         assert info.file_path == test_file