benchbox 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

benchbox/core/dataframe/data_loader.py

@@ -50,6 +50,8 @@ from benchbox.core.dataframe.capabilities import (
 from benchbox.core.dataframe.tuning.write_config import (
     DataFrameWriteConfiguration,
 )
+from benchbox.utils.compression import CompressionError, CompressionManager
+from benchbox.utils.file_format import detect_compression, detect_data_format, strip_compression_suffix
 
 if TYPE_CHECKING:
     from benchbox.core.tpch.schema import Table
@@ -239,6 +241,27 @@ class FormatConverter:
     performance.
     """
 
+    @staticmethod
+    def _has_trailing_delimiter(source_path: Path, delimiter: str) -> bool:
+        """Detect if the first data row ends with a trailing delimiter."""
+        compression_type = detect_compression(source_path)
+        try:
+            if compression_type:
+                manager = CompressionManager()
+                compressor = manager.get_compressor(compression_type)
+                with compressor.open_for_read(source_path, mode="rt") as handle:
+                    for line in handle:
+                        if line.strip():
+                            return line.rstrip("\n").endswith(delimiter)
+                return False
+            with source_path.open("rt", encoding="utf-8", errors="replace") as handle:
+                for line in handle:
+                    if line.strip():
+                        return line.rstrip("\n").endswith(delimiter)
+            return False
+        except (CompressionError, OSError):
+            return False
+
     @staticmethod
     def convert_csv_to_parquet(
         source_path: Path,
@@ -272,39 +295,51 @@ class FormatConverter:
             return ConversionStatus.FAILED, 0
 
         try:
+            format_type = detect_data_format(source_path)
+            is_tbl_file = format_type == "tbl"
+            has_trailing = (
+                is_tbl_file and bool(column_names) and FormatConverter._has_trailing_delimiter(source_path, delimiter)
+            )
+
             # Handle TBL files with trailing delimiter (TPC spec)
-            # Add a dummy column name to account for the trailing delimiter
-            is_tbl_file = str(source_path).endswith(".tbl")
             actual_column_names = column_names
-            if is_tbl_file and column_names:
-                # TBL files have trailing delimiter, add dummy column
+            if is_tbl_file and column_names and has_trailing:
                 actual_column_names = column_names + ["_trailing_"]
 
-            # Read options for CSV
             read_options = pv.ReadOptions(
                 column_names=actual_column_names if actual_column_names else None,
             )
-
             parse_options = pv.ParseOptions(delimiter=delimiter)
-
-            # Convert options for optimal Parquet output
             convert_options = pv.ConvertOptions(
                 auto_dict_encode=True,
                 strings_can_be_null=True,
             )
 
-            # Read the CSV file
             logger.debug(f"Reading {source_path}")
-            table = pv.read_csv(
-                source_path,
-                read_options=read_options,
-                parse_options=parse_options,
-                convert_options=convert_options,
-            )
-
-            # Drop trailing column for TBL files
-            if is_tbl_file and column_names:
-                # Select only the original columns (exclude _trailing_)
+            compression_type = detect_compression(source_path)
+            if compression_type:
+                manager = CompressionManager()
+                try:
+                    compressor = manager.get_compressor(compression_type)
+                    with compressor.open_for_read(source_path, mode="rb") as stream:
+                        table = pv.read_csv(
+                            stream,
+                            read_options=read_options,
+                            parse_options=parse_options,
+                            convert_options=convert_options,
+                        )
+                except CompressionError as err:
+                    logger.error(f"Failed to read compressed file {source_path}: {err}")
+                    return ConversionStatus.FAILED, 0
+            else:
+                table = pv.read_csv(
+                    source_path,
+                    read_options=read_options,
+                    parse_options=parse_options,
+                    convert_options=convert_options,
+                )
+
+            if is_tbl_file and column_names and has_trailing:
                 table = table.select(column_names)
 
             # Apply physical layout options from write_config
@@ -485,11 +520,18 @@ class DataCache:
 
             # Verify all table files exist
             cache_dir = manifest_path.parent
-            for table_name, table_info in manifest.tables.items():
-                file_path = cache_dir / table_info["file"]
-                if not file_path.exists():
-                    logger.debug(f"Cache file missing: {file_path}")
-                    return False
+            for _table_name, table_info in manifest.tables.items():
+                if "files" in table_info:
+                    for file_name in table_info["files"]:
+                        file_path = cache_dir / file_name
+                        if not file_path.exists():
+                            logger.debug(f"Cache file missing: {file_path}")
+                            return False
+                else:
+                    file_path = cache_dir / table_info["file"]
+                    if not file_path.exists():
+                        logger.debug(f"Cache file missing: {file_path}")
+                        return False
 
             return True
 
@@ -497,7 +539,9 @@ class DataCache:
             logger.debug(f"Invalid cache manifest: {e}")
             return False
 
-    def get_cached_files(self, benchmark: str, scale_factor: float, format: DataFormat) -> dict[str, Path] | None:
+    def get_cached_files(
+        self, benchmark: str, scale_factor: float, format: DataFormat
+    ) -> dict[str, Path | list[Path]] | None:
         """Get cached data file paths.
 
         Args:
@@ -518,7 +562,13 @@ class DataCache:
                 manifest = CacheManifest.from_dict(json.load(f))
 
             cache_dir = manifest_path.parent
-            return {table_name: cache_dir / table_info["file"] for table_name, table_info in manifest.tables.items()}
+            resolved: dict[str, Path | list[Path]] = {}
+            for table_name, table_info in manifest.tables.items():
+                if "files" in table_info:
+                    resolved[table_name] = [cache_dir / file_name for file_name in table_info["files"]]
+                else:
+                    resolved[table_name] = cache_dir / table_info["file"]
+            return resolved
 
         except (json.JSONDecodeError, KeyError):
             return None
@@ -604,7 +654,7 @@ class DataCache:
         return removed
 
 
-def _compute_source_hash(source_dir: Path, tables: dict[str, Path]) -> str:
+def _compute_source_hash(source_dir: Path, tables: dict[str, Path | list[Path]]) -> str:
     """Compute hash of source files for cache validation.
 
     Uses file modification times for fast validation.
@@ -618,10 +668,12 @@ def _compute_source_hash(source_dir: Path, tables: dict[str, Path]) -> str:
     """
     hash_data = []
     for table_name in sorted(tables.keys()):
-        file_path = tables[table_name]
-        if file_path.exists():
-            stat = file_path.stat()
-            hash_data.append(f"{table_name}:{stat.st_mtime}:{stat.st_size}")
+        file_paths = tables[table_name]
+        file_list = file_paths if isinstance(file_paths, list) else [file_paths]
+        for file_path in file_list:
+            if file_path.exists():
+                stat = file_path.stat()
+                hash_data.append(f"{table_name}:{file_path.name}:{stat.st_mtime}:{stat.st_size}")
 
     combined = "|".join(hash_data)
     return hashlib.md5(combined.encode()).hexdigest()[:12]
@@ -713,7 +765,7 @@ class DataFrameDataLoader:
         scale_factor: float,
         data_dir: Path | None = None,
         write_config: DataFrameWriteConfiguration | None = None,
-    ) -> dict[str, Path]:
+    ) -> dict[str, Path | list[Path]]:
         """Prepare benchmark data files for DataFrame loading.
 
         This method:
@@ -753,8 +805,10 @@ class DataFrameDataLoader:
         effective_write_config = write_config or self.write_config
 
         # Check cache
+        first_entry = next(iter(source_files.values()))
+        first_path = first_entry[0] if isinstance(first_entry, list) else first_entry
         source_hash = _compute_source_hash(
-            data_dir or Path(list(source_files.values())[0]).parent,
+            data_dir or Path(first_path).parent,
             source_files,
         )
 
@@ -786,7 +840,7 @@ class DataFrameDataLoader:
             write_config=effective_write_config,
         )
 
-    def _get_source_files(self, benchmark: Any, data_dir: Path | None) -> dict[str, Path]:
+    def _get_source_files(self, benchmark: Any, data_dir: Path | None) -> dict[str, list[Path]]:
         """Get source data file paths.
 
         Args:
@@ -794,13 +848,19 @@ class DataFrameDataLoader:
             data_dir: Optional data directory
 
         Returns:
-            Dictionary mapping table name to file path
+            Dictionary mapping table name to list of file paths
         """
         # First try benchmark.tables attribute
         if hasattr(benchmark, "tables") and benchmark.tables:
             tables = benchmark.tables
             if isinstance(tables, dict):
-                return {name: Path(path) if not isinstance(path, Path) else path for name, path in tables.items()}
+                resolved: dict[str, list[Path]] = {}
+                for name, path in tables.items():
+                    if isinstance(path, list):
+                        resolved[name] = [Path(p) if not isinstance(p, Path) else p for p in path]
+                    else:
+                        resolved[name] = [Path(path) if not isinstance(path, Path) else path]
+                return resolved
 
         # Then try data_dir
         if data_dir and data_dir.exists():
@@ -810,11 +870,17 @@ class DataFrameDataLoader:
         if hasattr(benchmark, "_impl") and hasattr(benchmark._impl, "tables"):
             tables = benchmark._impl.tables
             if isinstance(tables, dict):
-                return {name: Path(path) if not isinstance(path, Path) else path for name, path in tables.items()}
+                resolved: dict[str, list[Path]] = {}
+                for name, path in tables.items():
+                    if isinstance(path, list):
+                        resolved[name] = [Path(p) if not isinstance(p, Path) else p for p in path]
+                    else:
+                        resolved[name] = [Path(path) if not isinstance(path, Path) else path]
+                return resolved
 
         return {}
 
-    def _discover_files(self, data_dir: Path) -> dict[str, Path]:
+    def _discover_files(self, data_dir: Path) -> dict[str, list[Path]]:
         """Discover data files in a directory.
 
         Args:
@@ -823,7 +889,7 @@ class DataFrameDataLoader:
         Returns:
             Dictionary mapping table name to file path
         """
-        files = {}
+        files: dict[str, list[Path]] = {}
 
         for pattern in ["*.tbl", "*.csv", "*.parquet"]:
             for file_path in data_dir.glob(pattern):
@@ -837,19 +903,17 @@ class DataFrameDataLoader:
                         for pq_file in parquet_files:
                             if pq_file.is_file():
                                 table_name = pq_file.stem.lower()
-                                if table_name not in files:
-                                    files[table_name] = pq_file
+                                files.setdefault(table_name, []).append(pq_file)
                     continue
 
                 # Table name is stem of filename
                 table_name = file_path.stem.lower()
                 # Prefer existing entry if already found (TBL > CSV priority)
-                if table_name not in files:
-                    files[table_name] = file_path
+                files.setdefault(table_name, []).append(file_path)
 
         return files
 
-    def _detect_source_format(self, source_files: dict[str, Path]) -> DataFormat:
+    def _detect_source_format(self, source_files: dict[str, Path | list[Path]]) -> DataFormat:
         """Detect the format of source files.
 
         Args:
@@ -859,13 +923,16 @@ class DataFrameDataLoader:
             Detected DataFormat
         """
         for path in source_files.values():
-            suffix = path.suffix.lower()
-            if suffix == ".parquet":
-                return DataFormat.PARQUET
-            elif suffix == ".tbl":
-                return DataFormat.CSV  # TBL is CSV with different delimiter
-            elif suffix == ".arrow":
-                return DataFormat.ARROW
+            paths = path if isinstance(path, list) else [path]
+            for entry in paths:
+                format_type = detect_data_format(entry)
+                if format_type == "parquet":
+                    return DataFormat.PARQUET
+                if format_type == "tbl":
+                    return DataFormat.CSV
+                stripped_suffix = strip_compression_suffix(entry).suffix.lower()
+                if stripped_suffix == ".arrow":
+                    return DataFormat.ARROW
 
         return DataFormat.CSV
 
@@ -874,12 +941,12 @@ class DataFrameDataLoader:
         benchmark: Any,
         benchmark_name: str,
         scale_factor: float,
-        source_files: dict[str, Path],
+        source_files: dict[str, Path | list[Path]],
         source_format: DataFormat,
         target_format: DataFormat,
         source_hash: str,
         write_config: DataFrameWriteConfiguration | None = None,
-    ) -> dict[str, Path]:
+    ) -> dict[str, Path | list[Path]]:
         """Convert data files to target format.
 
         Args:
@@ -917,43 +984,77 @@ class DataFrameDataLoader:
         # Get schema info for column names
         schema_info = self._get_schema_info(benchmark)
 
-        converted_files = {}
-        table_metadata = {}
+        converted_files: dict[str, Path | list[Path]] = {}
+        table_metadata: dict[str, dict[str, Any]] = {}
 
         for table_name, source_path in source_files.items():
-            target_path = cache_path / f"{table_name}.parquet"
+            source_list = source_path if isinstance(source_path, list) else [source_path]
 
             # Get column names from schema
             column_names = schema_info.get(table_name)
 
-            # Determine delimiter
-            delimiter = "|" if source_path.suffix.lower() == ".tbl" else ","
-
             # Get table-specific write config (filter sort columns that exist in this table)
             table_write_config = self._get_table_write_config(write_config, table_name, column_names)
 
-            # Convert
-            status, row_count = FormatConverter.convert_csv_to_parquet(
-                source_path=source_path,
-                target_path=target_path,
-                column_names=column_names,
-                delimiter=delimiter,
-                write_config=table_write_config,
-            )
-
-            if status == ConversionStatus.SUCCESS:
-                converted_files[table_name] = target_path
+            converted_list: list[Path] = []
+            table_entries: list[dict[str, Any]] = []
+
+            for source_entry in source_list:
+                format_type = detect_data_format(source_entry)
+                delimiter = "|" if format_type == "tbl" else ","
+
+                stripped_name = strip_compression_suffix(source_entry).name
+                if format_type == "tbl":
+                    base_name = stripped_name.replace(".tbl", "", 1)
+                elif format_type == "csv":
+                    base_name = stripped_name.replace(".csv", "", 1)
+                else:
+                    base_name = Path(stripped_name).stem
+
+                target_path = cache_path / f"{base_name}.parquet"
+
+                status, row_count = FormatConverter.convert_csv_to_parquet(
+                    source_path=source_entry,
+                    target_path=target_path,
+                    column_names=column_names,
+                    delimiter=delimiter,
+                    write_config=table_write_config,
+                )
+
+                if status == ConversionStatus.SUCCESS:
+                    converted_list.append(target_path)
+                    table_entries.append(
+                        {
+                            "file": target_path.name,
+                            "row_count": row_count,
+                            "source": source_entry.name,
+                        }
+                    )
+                else:
+                    logger.error(f"Failed to convert {table_name}, using source file")
+                    converted_list.append(source_entry)
+                    table_entries.append(
+                        {
+                            "file": source_entry.name,
+                            "row_count": 0,
+                            "source": source_entry.name,
+                            "status": "FAILED",
+                        }
+                    )
+
+            if len(converted_list) == 1:
+                converted_files[table_name] = converted_list[0]
+                table_metadata[table_name] = table_entries[0]
+            else:
+                converted_files[table_name] = converted_list
                 table_metadata[table_name] = {
-                    "file": f"{table_name}.parquet",
-                    "row_count": row_count,
-                    "source": source_path.name,
+                    "files": [entry["file"] for entry in table_entries],
+                    "row_counts": [entry["row_count"] for entry in table_entries],
+                    "sources": [entry["source"] for entry in table_entries],
                 }
-                # Record write config if applied
-                if table_write_config and not table_write_config.is_default():
-                    table_metadata[table_name]["write_config"] = table_write_config.to_dict()
-            else:
-                logger.error(f"Failed to convert {table_name}, using source file")
-                converted_files[table_name] = source_path
+
+            if table_write_config and not table_write_config.is_default():
+                table_metadata[table_name]["write_config"] = table_write_config.to_dict()
 
         # Save manifest
         self.cache.save_manifest(