Fix download_demo for data.zip files (#2699)

fealho · web-flow · commit 18ece9e5a775 · 2025-10-16T09:41:06.000-07:00
diff --git a/sdv/datasets/demo.py b/sdv/datasets/demo.py
@@ -24,6 +24,7 @@
 BUCKET_URL = f'https://{BUCKET}.s3.amazonaws.com'
 SIGNATURE_VERSION = UNSIGNED
 METADATA_FILENAME = 'metadata.json'
+FALLBACK_ENCODING = 'latin-1'
 
 
 def _validate_modalities(modality):
@@ -197,21 +198,67 @@ def _extract_data(bytes_io, output_folder_name):
             return in_memory_directory
 
 
-def _get_data(modality, output_folder_name, in_memory_directory):
+def _get_data_with_output_folder(output_folder_name):
+    """Load CSV tables from an extracted folder on disk.
+
+    Returns a tuple of (data_dict, skipped_files).
+    Non-CSV files are ignored.
+    """
     data = {}
-    if output_folder_name:
-        for root, _dirs, files in os.walk(output_folder_name):
-            for filename in files:
-                if filename.endswith('.csv'):
-                    table_name = Path(filename).stem
-                    data_path = os.path.join(root, filename)
-                    data[table_name] = pd.read_csv(data_path)
+    skipped_files = []
+    for root, _dirs, files in os.walk(output_folder_name):
+        for filename in files:
+            if not filename.lower().endswith('.csv'):
+                skipped_files.append(filename)
+                continue
+
+            table_name = Path(filename).stem
+            data_path = os.path.join(root, filename)
+            try:
+                data[table_name] = pd.read_csv(data_path)
+            except UnicodeDecodeError:
+                data[table_name] = pd.read_csv(data_path, encoding=FALLBACK_ENCODING)
+            except Exception as e:
+                rel = os.path.relpath(data_path, output_folder_name)
+                skipped_files.append(f'{rel}: {e}')
+
+    return data, skipped_files
+
+
+def _get_data_without_output_folder(in_memory_directory):
+    """Load CSV tables directly from in-memory zip contents.
+
+    Returns a tuple of (data_dict, skipped_files).
+    Non-CSV entries are ignored.
+    """
+    data = {}
+    skipped_files = []
+    for filename, file_ in in_memory_directory.items():
+        if not filename.lower().endswith('.csv'):
+            skipped_files.append(filename)
+            continue
+
+        table_name = Path(filename).stem
+        try:
+            data[table_name] = pd.read_csv(io.BytesIO(file_), low_memory=False)
+        except UnicodeDecodeError:
+            data[table_name] = pd.read_csv(
+                io.BytesIO(file_), low_memory=False, encoding=FALLBACK_ENCODING
+            )
+        except Exception as e:
+            skipped_files.append(f'{filename}: {e}')
 
+    return data, skipped_files
+
+
+def _get_data(modality, output_folder_name, in_memory_directory):
+    if output_folder_name:
+        data, skipped_files = _get_data_with_output_folder(output_folder_name)
     else:
-        for filename, file_ in in_memory_directory.items():
-            if filename.endswith('.csv'):
-                table_name = Path(filename).stem
-                data[table_name] = pd.read_csv(io.StringIO(file_.decode()), low_memory=False)
+        data, skipped_files = _get_data_without_output_folder(in_memory_directory)
+
+    if skipped_files:
+        warnings.warn('Skipped files: ' + ', '.join(sorted(skipped_files)))
 
     if not data:
         raise DemoResourceNotFoundError(
diff --git a/tests/unit/datasets/test_demo.py b/tests/unit/datasets/test_demo.py
@@ -1019,3 +1019,181 @@ def test_download_demo_raises_when_no_csv_in_zip_single_table(mock_list, mock_ge
     msg = 'Demo data could not be downloaded because no csv files were found in data.zip'
     with pytest.raises(DemoResourceNotFoundError, match=re.escape(msg)):
         download_demo('single_table', 'word')
+
+
+@patch('sdv.datasets.demo._get_data_from_bucket')
+@patch('sdv.datasets.demo._list_objects')
+def test_download_demo_skips_non_csv_in_memory_no_warning(mock_list, mock_get):
+    """In-memory path: ignore non-CSV files silently; load valid CSVs."""
+    # Setup
+    mock_list.return_value = [
+        {'Key': 'single_table/mix/data.zip'},
+        {'Key': 'single_table/mix/metadata.json'},
+    ]
+
+    df = pd.DataFrame({'id': [1, 2], 'name': ['a', 'b']})
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.writestr('good.csv', df.to_csv(index=False))
+        zf.writestr('note.txt', 'hello world')
+        zf.writestr('nested/readme.md', '# readme')
+        # Add a directory entry explicitly
+        zf.writestr('empty_dir/', '')
+    zip_bytes = buf.getvalue()
+
+    meta_bytes = json.dumps({
+        'METADATA_SPEC_VERSION': 'V1',
+        'tables': {
+            'good': {
+                'columns': {
+                    'id': {'sdtype': 'numerical', 'computer_representation': 'Int64'},
+                    'name': {'sdtype': 'categorical'},
+                }
+            }
+        },
+        'relationships': [],
+    }).encode()
+
+    mock_get.side_effect = lambda key: zip_bytes if key.endswith('data.zip') else meta_bytes
+
+    # Run and Assert
+    warn_msg = 'Skipped files: empty_dir/, nested/readme.md, note.txt'
+    with pytest.warns(UserWarning, match=warn_msg) as rec:
+        data, _ = download_demo('single_table', 'mix')
+
+    assert len(rec) == 1
+    expected = pd.DataFrame({'id': [1, 2], 'name': ['a', 'b']})
+    pd.testing.assert_frame_equal(data, expected)
+
+
+@patch('sdv.datasets.demo._get_data_from_bucket')
+@patch('sdv.datasets.demo._list_objects')
+def test_download_demo_on_disk_warns_failed_csv_only(mock_list, mock_get, tmp_path, monkeypatch):
+    """On-disk path: warn only for failed CSVs; non-CSV are skipped silently."""
+    # Setup
+    mock_list.return_value = [
+        {'Key': 'single_table/mix/data.zip'},
+        {'Key': 'single_table/mix/metadata.json'},
+    ]
+
+    good = pd.DataFrame({'x': [1, 2]})
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.writestr('good.csv', good.to_csv(index=False))
+        zf.writestr('bad.csv', 'will_fail')
+        zf.writestr('info.txt', 'ignore me')
+    zip_bytes = buf.getvalue()
+
+    meta_bytes = json.dumps({
+        'METADATA_SPEC_VERSION': 'V1',
+        'tables': {
+            'good': {
+                'columns': {
+                    'x': {'sdtype': 'numerical', 'computer_representation': 'Int64'},
+                }
+            }
+        },
+        'relationships': [],
+    }).encode()
+
+    mock_get.side_effect = lambda key: zip_bytes if key.endswith('data.zip') else meta_bytes
+
+    # Force read_csv to fail on bad.csv only
+    orig_read_csv = pd.read_csv
+
+    def fake_read_csv(path_or_buf, *args, **kwargs):
+        if isinstance(path_or_buf, str) and path_or_buf.endswith('bad.csv'):
+            raise ValueError('bad-parse')
+        return orig_read_csv(path_or_buf, *args, **kwargs)
+
+    monkeypatch.setattr('pandas.read_csv', fake_read_csv)
+
+    out_dir = tmp_path / 'mix_out'
+
+    # Run and Assert
+    warn_msg = 'Skipped files: bad.csv: bad-parse, info.txt'
+    with pytest.warns(UserWarning, match=warn_msg) as rec:
+        data, _ = download_demo('single_table', 'mix', out_dir)
+
+    assert len(rec) == 1
+    pd.testing.assert_frame_equal(data, good)
+
+
+@patch('sdv.datasets.demo._get_data_from_bucket')
+@patch('sdv.datasets.demo._list_objects')
+def test_download_demo_handles_non_utf8_in_memory(mock_list, mock_get):
+    """It should successfully read Latin-1 encoded CSVs from in-memory extraction."""
+    # Setup
+    mock_list.return_value = [
+        {'Key': 'single_table/nonutf/data.zip'},
+        {'Key': 'single_table/nonutf/metadata.json'},
+    ]
+
+    df = pd.DataFrame({'id': [1], 'name': ['café']})
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.writestr('nonutf.csv', df.to_csv(index=False).encode('latin-1'))
+    zip_bytes = buf.getvalue()
+
+    meta_bytes = json.dumps({
+        'METADATA_SPEC_VERSION': 'V1',
+        'tables': {
+            'nonutf': {
+                'columns': {
+                    'id': {'sdtype': 'numerical', 'computer_representation': 'Int64'},
+                    'name': {'sdtype': 'categorical'},
+                }
+            }
+        },
+        'relationships': [],
+    }).encode()
+
+    mock_get.side_effect = lambda key: zip_bytes if key.endswith('data.zip') else meta_bytes
+
+    # Run
+    data, _ = download_demo('single_table', 'nonutf')
+
+    # Assert
+    expected = pd.DataFrame({'id': [1], 'name': ['café']})
+    pd.testing.assert_frame_equal(data, expected)
+
+
+@patch('sdv.datasets.demo._get_data_from_bucket')
+@patch('sdv.datasets.demo._list_objects')
+def test_download_demo_handles_non_utf8_on_disk(mock_list, mock_get, tmp_path):
+    """It should successfully read Latin-1 encoded CSVs when extracted to disk."""
+    # Setup
+    mock_list.return_value = [
+        {'Key': 'single_table/nonutf/data.zip'},
+        {'Key': 'single_table/nonutf/metadata.json'},
+    ]
+
+    df = pd.DataFrame({'id': [1], 'name': ['café']})
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.writestr('nonutf.csv', df.to_csv(index=False).encode('latin-1'))
+    zip_bytes = buf.getvalue()
+
+    meta_bytes = json.dumps({
+        'METADATA_SPEC_VERSION': 'V1',
+        'tables': {
+            'nonutf': {
+                'columns': {
+                    'id': {'sdtype': 'numerical', 'computer_representation': 'Int64'},
+                    'name': {'sdtype': 'categorical'},
+                }
+            }
+        },
+        'relationships': [],
+    }).encode()
+
+    mock_get.side_effect = lambda key: zip_bytes if key.endswith('data.zip') else meta_bytes
+
+    out_dir = tmp_path / 'latin_out'
+
+    # Run
+    data, _ = download_demo('single_table', 'nonutf', out_dir)
+
+    # Assert
+    expected = pd.DataFrame({'id': [1], 'name': ['café']})
+    pd.testing.assert_frame_equal(data, expected)
