fsspeckit.datasets API Reference

datasets

Dataset-level operations for fsspeckit.

This package contains dataset-specific functionality including: - DuckDB parquet handlers for high-performance dataset operations - PyArrow handlers for dataset I/O operations - PyArrow utilities for schema management and type conversion - Dataset merging and optimization tools

Classes

fsspeckit.datasets.DatasetError

DatasetError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: Exception

Base exception for fsspeckit dataset operations.

Attributes:

Name	Type	Description
message		Human-readable error description.
operation		The operation that failed (e.g., 'read', 'write', 'merge').
details		Additional structured context about the error.

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetFileError

DatasetFileError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when file I/O operations fail.

Use this for file read/write errors, permission issues, etc.

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetMergeError

DatasetMergeError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when merge operations fail.

Use this for merge-specific failures (key column issues, schema mismatches, etc.).

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetOperationError

DatasetOperationError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when a dataset operation fails.

Use this for general operation failures that don't fit more specific categories.

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetPathError

DatasetPathError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when path-related operations fail.

Used for path normalization failures, missing paths, and invalid protocols surfaced through core path validation. The datasets layer preserves and enriches this error type at the public boundary.

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetSchemaError

DatasetSchemaError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when schema-related operations fail.

Use this for schema validation, casting, and compatibility issues.

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.DatasetValidationError

DatasetValidationError(
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
)

Bases: DatasetError

Raised when input validation fails.

Use this when user-provided input is invalid (e.g., invalid mode, missing columns).

Source code in src/fsspeckit/core/exceptions.py

def __init__(
    self,
    message: str,
    operation: str | None = None,
    details: dict[str, Any] | None = None,
) -> None:
    super().__init__(message)
    self.operation = operation
    self.details = details or {}

fsspeckit.datasets.PyarrowDatasetIO

PyarrowDatasetIO(
    filesystem: AbstractFileSystem | None = None,
)

Bases: BaseDatasetHandler

PyArrow-based dataset I/O operations.

This class provides methods for reading and writing parquet files and datasets using PyArrow's high-performance parquet engine.

The class inherits from BaseDatasetHandler to leverage shared implementations while providing PyArrow-specific optimizations.

Parameters:

Name	Type	Description	Default
filesystem	AbstractFileSystem | None	Optional fsspec filesystem instance. If None, uses local filesystem.	None

Example

from fsspeckit.datasets.pyarrow import PyarrowDatasetIO

io = PyarrowDatasetIO()

# Read parquet
table = io.read_parquet("/path/to/data.parquet")

# Write dataset
io.write_dataset(table, "/path/to/dataset/", mode="append")

# Merge into dataset
io.merge(table, "/path/to/dataset/", strategy="upsert", key_columns=["id"])

Initialize PyArrow dataset I/O.

Parameters:

Name	Type	Description	Default
filesystem	AbstractFileSystem | None	Optional fsspec filesystem. If None, uses local filesystem.	None

Source code in src/fsspeckit/datasets/pyarrow/io.py

def __init__(
    self,
    filesystem: AbstractFileSystem | None = None,
) -> None:
    """Initialize PyArrow dataset I/O.

    Args:
        filesystem: Optional fsspec filesystem. If None, uses local filesystem.
    """
    from fsspeckit.common import optional as optional_module

    if not optional_module._PYARROW_AVAILABLE:
        raise ImportError(
            "pyarrow is required for PyarrowDatasetIO. "
            "Install with: pip install fsspeckit[datasets]"
        )

    if filesystem is None:
        filesystem = fsspec_filesystem("file")

    self._filesystem = filesystem

Attributes

fsspeckit.datasets.PyarrowDatasetIO.filesystem property

filesystem: AbstractFileSystem

Return the filesystem instance.

Functions

fsspeckit.datasets.PyarrowDatasetIO.__enter__

__enter__() -> PyarrowDatasetIO

Enter context manager.

Provided for API symmetry with DuckDB. PyArrow has no connection to close, so this is a no-op context manager.

Source code in src/fsspeckit/datasets/pyarrow/io.py

def __enter__(self) -> PyarrowDatasetIO:
    """Enter context manager.

    Provided for API symmetry with DuckDB. PyArrow has no connection
    to close, so this is a no-op context manager.
    """
    return self

fsspeckit.datasets.PyarrowDatasetIO.__exit__

__exit__(
    exc_type: type[BaseException] | None,
    exc_value: BaseException | None,
    traceback: Any,
) -> None

Exit context manager (no-op for PyArrow).

Source code in src/fsspeckit/datasets/pyarrow/io.py

def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc_value: BaseException | None,
    traceback: Any,
) -> None:
    """Exit context manager (no-op for PyArrow)."""
    return None

fsspeckit.datasets.PyarrowDatasetIO.merge

merge(
    data: Table | list[Table],
    path: str,
    strategy: Literal["insert", "update", "upsert"],
    key_columns: list[str] | str,
    *,
    partition_columns: list[str] | str | None = None,
    schema: Schema | None = None,
    compression: str | None = "snappy",
    max_rows_per_file: int | None = 5000000,
    row_group_size: int | None = 500000,
    merge_chunk_size_rows: int = 100000,
    enable_streaming_merge: bool = True,
    merge_max_memory_mb: int = 1024,
    merge_max_process_memory_mb: int | None = None,
    merge_min_system_available_mb: int = 512,
    merge_progress_callback: Callable[[int, int], None]
    | None = None,
    use_merge: bool | None = None,
) -> MergeResult

Merge data into an existing parquet dataset.

This method performs an incremental merge (insert, update, or upsert) of the provided data into the target dataset. It uses PyArrow's high-performance operations and supports both in-memory and streaming merge strategies.

Parameters:

Name	Type	Description	Default
data	Table | list[Table]	PyArrow table or list of tables to merge	required
path	str	Target dataset path	required
strategy	Literal['insert', 'update', 'upsert']	Merge strategy ('insert', 'update', or 'upsert')	required
key_columns	list[str] | str	Column(s) used as unique identifiers	required
partition_columns	list[str] | str | None	Optional column(s) to partition by	None
schema	Schema | None	Optional schema to enforce on the source data	None
compression	str | None	Compression codec (default: snappy)	'snappy'
max_rows_per_file	int | None	Max rows per file in output (default: 5,000,000)	5000000
row_group_size	int | None	Rows per row group (default: 500_000)	500000
merge_chunk_size_rows	int	Rows per processing chunk (default: 100_000)	100000
enable_streaming_merge	bool	Whether to use streaming merge (default: True)	True
merge_max_memory_mb	int	Max PyArrow memory in MB (default: 1024)	1024
merge_max_process_memory_mb	int | None	Optional max process RSS in MB	None
merge_min_system_available_mb	int	Min system available memory in MB (default: 512)	512
merge_progress_callback	Callable[[int, int], None] | None	Optional callback for progress updates	None
use_merge	bool | None	Reserved for backward compatibility (ignored by current implementations)	None

Returns:

Type	Description
MergeResult	MergeResult with detailed statistics

Source code in src/fsspeckit/datasets/pyarrow/io.py

def merge(
    self,
    data: pa.Table | list[pa.Table],
    path: str,
    strategy: Literal["insert", "update", "upsert"],
    key_columns: list[str] | str,
    *,
    partition_columns: list[str] | str | None = None,
    schema: pa.Schema | None = None,
    compression: str | None = "snappy",
    max_rows_per_file: int | None = 5_000_000,
    row_group_size: int | None = 500_000,
    merge_chunk_size_rows: int = 100_000,
    enable_streaming_merge: bool = True,
    merge_max_memory_mb: int = 1024,
    merge_max_process_memory_mb: int | None = None,
    merge_min_system_available_mb: int = 512,
    merge_progress_callback: Callable[[int, int], None] | None = None,
    use_merge: bool | None = None,
) -> MergeResult:
    """Merge data into an existing parquet dataset.

    This method performs an incremental merge (insert, update, or upsert)
    of the provided data into the target dataset. It uses PyArrow's
    high-performance operations and supports both in-memory and
    streaming merge strategies.

    Args:
        data: PyArrow table or list of tables to merge
        path: Target dataset path
        strategy: Merge strategy ('insert', 'update', or 'upsert')
        key_columns: Column(s) used as unique identifiers
        partition_columns: Optional column(s) to partition by
        schema: Optional schema to enforce on the source data
        compression: Compression codec (default: snappy)
        max_rows_per_file: Max rows per file in output (default: 5,000,000)
        row_group_size: Rows per row group (default: 500_000)
        merge_chunk_size_rows: Rows per processing chunk (default: 100_000)
        enable_streaming_merge: Whether to use streaming merge (default: True)
        merge_max_memory_mb: Max PyArrow memory in MB (default: 1024)
        merge_max_process_memory_mb: Optional max process RSS in MB
        merge_min_system_available_mb: Min system available memory in MB (default: 512)
        merge_progress_callback: Optional callback for progress updates
        use_merge: Reserved for backward compatibility (ignored by current implementations)

    Returns:
        MergeResult with detailed statistics
    """
    import pyarrow.parquet as pq
    import pyarrow.dataset as ds

    from fsspeckit.core.incremental import (
        IncrementalFileManager,
        MergeFileMetadata,
        MergeResult,
        confirm_affected_files,
        extract_source_partition_values,
        list_dataset_files,
        parse_hive_partition_path,
        plan_incremental_rewrite,
        validate_no_null_keys,
        validate_partition_column_immutability,
    )
    from fsspeckit.common.security import validate_compression_codec, validate_path
    from fsspeckit.datasets.pyarrow.dataset import (
        PerformanceMonitor,
        _filter_by_key_membership,
        _make_struct_safe,
    )
    from fsspeckit.datasets.pyarrow.adaptive_tracker import AdaptiveKeyTracker

    monitor = PerformanceMonitor(
        max_pyarrow_mb=merge_max_memory_mb,
        max_process_memory_mb=merge_max_process_memory_mb,
        min_system_available_mb=merge_min_system_available_mb,
    )
    monitor.start_op("initialization")

    if use_merge is not None:
        logger.debug("pyarrow_merge_use_merge_ignored", use_merge=use_merge)

    pa_mod = _import_pyarrow()
    validate_path(path)
    validate_compression_codec(compression)
    row_group_size = self._validate_write_parameters(
        max_rows_per_file,
        row_group_size,
    )

    # Convert data to source_table
    source_table = self._combine_tables(data)

    if schema is not None:
        from fsspeckit.datasets.schema import cast_schema

        source_table = cast_schema(source_table, schema)

    key_cols = self._validate_key_columns(
        key_columns,
        source_table.column_names,
        context="source",
    )
    partition_cols = self._validate_partition_columns(
        partition_columns,
        source_table.column_names,
    )

    # Validate no null keys in source before planning/deduplication can mask rows.
    validate_no_null_keys(source_table, key_cols)

    target_files = list_dataset_files(path, filesystem=self._filesystem)
    target_metadata = MergeTargetMetadata(
        exists=bool(target_files),
        files=target_files,
        row_count=sum(
            pq.read_metadata(f, filesystem=self._filesystem).num_rows
            for f in target_files
        ),
    )

    monitor.start_op("source_deduplication")
    plan = plan_merge_operation(
        source_table=source_table,
        strategy=strategy,
        key_columns=key_cols,
        target_metadata=target_metadata,
        partition_columns=partition_cols,
    )
    monitor.end_op()

    source_table = plan.source_table
    key_cols = plan.key_columns
    partition_cols = plan.partition_columns
    target_files = plan.target_files
    target_exists = plan.target_exists
    target_count_before = plan.target_count_before

    # Keep source keys as PyArrow Table/Array for vectorized operations.
    source_key_table = source_table.select(key_cols)
    source_key_tracker = AdaptiveKeyTracker()
    if len(key_cols) == 1:
        # For single column, keep as PyArrow array for vectorized operations.
        source_key_array = source_key_table.column(0)
        for key in plan.source_keys:
            source_key_tracker.add(key)
    else:
        # For multi-column keys, use vectorized conversion.
        source_key_array = None
        for key in plan.source_keys:
            source_key_tracker.add(key)

    early_result = resolve_merge_plan_early_exit(plan)
    if early_result is not None:
        return early_result

    if not target_exists:
        self._filesystem.mkdirs(path, exist_ok=True)
        write_res = self.write_dataset(
            source_table,
            path,
            mode="append",
            partition_by=partition_cols or None,
            partitioning_flavor="hive" if partition_cols else None,
            compression=compression,
            max_rows_per_file=max_rows_per_file,
            row_group_size=row_group_size,
        )

        inserted_files = [m.path for m in write_res.files]
        files_meta = [
            MergeFileMetadata(
                path=m.path,
                row_count=m.row_count,
                operation="inserted",
                size_bytes=m.size_bytes,
            )
            for m in write_res.files
        ]

        return MergeResult(
            strategy=strategy,
            source_count=source_table.num_rows,
            target_count_before=0,
            target_count_after=write_res.total_rows,
            inserted=write_res.total_rows,
            updated=0,
            deleted=0,
            files=files_meta,
            rewritten_files=[],
            inserted_files=inserted_files,
            preserved_files=[],
        )

    def _select_rows_by_keys(
        table: pa.Table, keys: set[object] | AdaptiveKeyTracker
    ) -> pa.Table:
        is_tracker = isinstance(keys, AdaptiveKeyTracker)
        if is_tracker:
            if keys.get_metrics()["unique_keys_estimate"] == 0:
                return table.slice(0, 0)
        elif not keys:
            return table.slice(0, 0)

        if is_tracker:
            # If tracker has exact keys, use them for efficiency
            if (
                hasattr(keys, "_tier")
                and keys._tier == "EXACT"
                and keys._exact_keys is not None
            ):
                key_set = keys._exact_keys
            else:
                # Probabilistic or LRU: must filter manually
                mask = []
                if len(key_cols) == 1:
                    col_values = table.column(key_cols[0]).to_pylist()
                    for val in col_values:
                        mask.append(val in keys)
                else:
                    struct_keys = _make_struct_safe(table, key_cols).to_pylist()
                    for d in struct_keys:
                        mask.append(tuple(d.values()) in keys)
                return table.filter(pa_mod.array(mask))
        else:
            key_set = keys

        return self._select_rows_by_keys(
            table,
            key_cols,
            key_set,
        )

    source_partition_values: set[tuple[object, ...]] | None = None
    if partition_cols:
        source_partition_values = extract_source_partition_values(
            source_table, partition_cols
        )

    # Prepare keys for planning (using lists for now as the planning utilities expect sequences)
    # TODO: Update planning utilities to support trackers or PyArrow tables directly
    if len(key_cols) == 1:
        source_keys_seq = source_key_table.column(0).to_pylist()
    else:
        source_keys_seq = [
            tuple(d.values())
            for d in _make_struct_safe(source_key_table, key_cols).to_pylist()
        ]

    rewrite_plan = plan_incremental_rewrite(
        dataset_path=path,
        source_keys=source_keys_seq,
        key_columns=key_cols,
        filesystem=self._filesystem,
        partition_columns=partition_cols or None,
        source_partition_values=source_partition_values,
    )

    affected_files = confirm_affected_files(
        candidate_files=rewrite_plan.affected_files,
        key_columns=key_cols,
        source_keys=source_keys_seq,
        filesystem=self._filesystem,
    )

    matched_keys = AdaptiveKeyTracker()
    matched_keys_by_file: dict[str, AdaptiveKeyTracker] = {}
    for file_path in affected_files:
        try:
            key_table = pq.read_table(
                file_path, columns=key_cols, filesystem=self._filesystem
            )

            # Use vectorized matching helper
            matched_rows = _filter_by_key_membership(
                key_table, key_cols, source_key_table, keep_matches=True
            )
            if matched_rows.num_rows > 0:
                file_matched = AdaptiveKeyTracker()
                if len(key_cols) == 1:
                    # Get the actual matched keys
                    for key in matched_rows.column(0).to_pylist():
                        file_matched.add(key)
                        matched_keys.add(key)
                else:
                    # Use struct array for vectorized multi-key extraction
                    for d in _make_struct_safe(matched_rows, key_cols).to_pylist():
                        key = tuple(d.values())
                        file_matched.add(key)
                        matched_keys.add(key)
                matched_keys_by_file[file_path] = file_matched
        except (IOError, RuntimeError, ValueError) as e:
            logger.error(
                "failed_to_check_file_for_matching_keys",
                path=file_path,
                error=str(e),
                operation="merge",
                exc_info=True,
            )
            # Conservative: if we can't confirm, treat all source keys as matched
            matched_keys_by_file[file_path] = source_key_tracker
            if len(key_cols) == 1:
                for key in source_key_array.to_pylist():
                    matched_keys.add(key)
            else:
                for d in _make_struct_safe(source_key_table, key_cols).to_pylist():
                    matched_keys.add(tuple(d.values()))

    # Calculate inserted keys using trackers
    inserted_key_tracker = AdaptiveKeyTracker()
    if len(key_cols) == 1:
        for key in source_key_array.to_pylist():
            if key not in matched_keys:
                inserted_key_tracker.add(key)
    else:
        for d in _make_struct_safe(source_key_table, key_cols).to_pylist():
            key = tuple(d.values())
            if key not in matched_keys:
                inserted_key_tracker.add(key)

    if strategy == "insert":
        preserved_files = list(target_files)
        if inserted_key_tracker.get_metrics()["unique_keys_estimate"] == 0:
            return MergeResult(
                strategy="insert",
                source_count=source_table.num_rows,
                target_count_before=target_count_before,
                target_count_after=target_count_before,
                inserted=0,
                updated=0,
                deleted=0,
                files=[
                    MergeFileMetadata(path=f, row_count=0, operation="preserved")
                    for f in preserved_files
                ],
                rewritten_files=[],
                inserted_files=[],
                preserved_files=preserved_files,
            )
        insert_table = _select_rows_by_keys(source_table, inserted_key_tracker)
        write_res = self.write_dataset(
            insert_table,
            path,
            mode="append",
            partition_by=partition_cols or None,
            partitioning_flavor="hive" if partition_cols else None,
            compression=compression,
            max_rows_per_file=max_rows_per_file,
            row_group_size=row_group_size,
        )
        inserted_files = [m.path for m in write_res.files]
        inserted_meta = [
            MergeFileMetadata(
                path=m.path,
                row_count=m.row_count,
                operation="inserted",
                size_bytes=m.size_bytes,
            )
            for m in write_res.files
        ]
        files_meta = [
            MergeFileMetadata(path=f, row_count=0, operation="preserved")
            for f in preserved_files
        ] + inserted_meta
        return MergeResult(
            strategy="insert",
            source_count=source_table.num_rows,
            target_count_before=target_count_before,
            target_count_after=target_count_before + insert_table.num_rows,
            inserted=insert_table.num_rows,
            updated=0,
            deleted=0,
            files=files_meta,
            rewritten_files=[],
            inserted_files=inserted_files,
            preserved_files=preserved_files,
        )

    file_manager = IncrementalFileManager()
    import uuid

    staging_dir = file_manager.create_staging_directory(
        path, filesystem=self._filesystem
    )
    rewritten_files: list[str] = []
    rewritten_meta: list[MergeFileMetadata] = []
    preserved_files = [f for f in target_files if f not in affected_files]

    try:
        for file_path in affected_files:
            monitor.start_op("file_processing")
            file_matched = matched_keys_by_file.get(file_path)
            if not file_matched:
                preserved_files.append(file_path)
                monitor.end_op()
                continue

            # Load only source rows relevant to this file
            source_for_file = _select_rows_by_keys(source_table, file_matched)

            if partition_cols:
                file_schema = pq.read_schema(file_path, filesystem=self._filesystem)
                available_columns = set(file_schema.names)
                validation_columns = list(
                    dict.fromkeys(
                        list(key_cols)
                        + [
                            col
                            for col in partition_cols
                            if col in available_columns
                        ]
                    )
                )
                target_validation = pq.read_table(
                    file_path,
                    columns=validation_columns,
                    filesystem=self._filesystem,
                )
                missing_partition_cols = [
                    col
                    for col in partition_cols
                    if col not in target_validation.column_names
                ]
                if missing_partition_cols:
                    partition_values = parse_hive_partition_path(
                        file_path, partition_columns=partition_cols
                    )
                    for col in missing_partition_cols:
                        if col not in partition_values:
                            raise ValueError(
                                f"Partition column '{col}' must be present in "
                                "hive partition path for merge validation"
                            )
                        target_validation = target_validation.append_column(
                            col,
                            pa_mod.array(
                                [partition_values[col]] * target_validation.num_rows
                            ),
                        )
                validate_partition_column_immutability(
                    source_for_file,
                    target_validation,
                    key_cols,
                    partition_cols,
                )

            staging_file = f"{staging_dir}/{uuid.uuid4().hex[:16]}.parquet"

            if enable_streaming_merge:
                # Use streaming merge for this file
                monitor.start_op("streaming_merge")
                existing_dataset = ds.dataset(
                    file_path, filesystem=self._filesystem
                )

                with pq.ParquetWriter(
                    staging_file,
                    existing_dataset.schema,
                    filesystem=self._filesystem,
                    compression=compression or "snappy",
                ) as writer:
                    from fsspeckit.datasets.pyarrow.dataset import (
                        merge_upsert_pyarrow,
                        merge_update_pyarrow,
                    )

                    if strategy == "upsert":
                        merge_upsert_pyarrow(
                            existing_dataset,
                            source_for_file,
                            key_cols,
                            chunk_size=merge_chunk_size_rows,
                            max_memory_mb=merge_max_memory_mb,
                            memory_monitor=monitor._memory_monitor,
                            writer=writer,
                        )
                    else:  # strategy == "update"
                        merge_update_pyarrow(
                            existing_dataset,
                            source_for_file,
                            key_cols,
                            chunk_size=merge_chunk_size_rows,
                            max_memory_mb=merge_max_memory_mb,
                            memory_monitor=monitor._memory_monitor,
                            writer=writer,
                        )
                monitor.end_op()
                updated_row_count = pq.read_metadata(
                    staging_file, filesystem=self._filesystem
                ).num_rows
            else:
                # Classic in-memory merge
                monitor.start_op("in_memory_merge")
                target_table = pq.read_table(file_path, filesystem=self._filesystem)

                from fsspeckit.datasets.pyarrow.dataset import (
                    merge_upsert_pyarrow,
                    merge_update_pyarrow,
                )

                if strategy == "upsert":
                    updated_table = merge_upsert_pyarrow(
                        target_table,
                        source_for_file,
                        key_cols,
                        max_memory_mb=merge_max_memory_mb,
                        memory_monitor=monitor._memory_monitor,
                    )
                else:
                    updated_table = merge_update_pyarrow(
                        target_table,
                        source_for_file,
                        key_cols,
                        max_memory_mb=merge_max_memory_mb,
                        memory_monitor=monitor._memory_monitor,
                    )

                pq.write_table(
                    updated_table,
                    staging_file,
                    filesystem=self._filesystem,
                    compression=compression,
                    row_group_size=row_group_size,
                )
                updated_row_count = updated_table.num_rows
                monitor.end_op()

            size_bytes = None
            try:
                size_bytes = int(self._filesystem.size(staging_file))
            except Exception:
                size_bytes = None

            file_manager.atomic_replace_files(
                [staging_file], [file_path], filesystem=self._filesystem
            )
            rewritten_files.append(file_path)
            rewritten_meta.append(
                MergeFileMetadata(
                    path=file_path,
                    row_count=updated_row_count,
                    operation="rewritten",
                    size_bytes=size_bytes,
                )
            )
            monitor.files_processed += 1
            monitor.end_op()
    finally:
        file_manager.cleanup_staging_files(filesystem=self._filesystem)
        monitor.end_op()

    inserted_files: list[str] = []
    inserted_meta: list[MergeFileMetadata] = []
    inserted_rows = 0

    if (
        strategy == "upsert"
        and inserted_key_tracker.get_metrics()["unique_keys_estimate"] > 0
    ):
        insert_table = _select_rows_by_keys(source_table, inserted_key_tracker)
        inserted_rows = insert_table.num_rows
        write_res = self.write_dataset(
            insert_table,
            path,
            mode="append",
            partition_by=partition_cols or None,
            partitioning_flavor="hive" if partition_cols else None,
            compression=compression,
            max_rows_per_file=max_rows_per_file,
            row_group_size=row_group_size,
        )
        inserted_files = [m.path for m in write_res.files]
        inserted_meta = [
            MergeFileMetadata(
                path=m.path,
                row_count=m.row_count,
                operation="inserted",
                size_bytes=m.size_bytes,
            )
            for m in write_res.files
        ]

    updated_rows = matched_keys.get_metrics()["unique_keys_estimate"]
    files_meta = (
        rewritten_meta
        + inserted_meta
        + [
            MergeFileMetadata(path=f, row_count=0, operation="preserved")
            for f in preserved_files
        ]
    )

    monitor.track_memory()
    metrics = monitor.get_metrics(
        total_rows_before=target_count_before,
        total_rows_after=target_count_before + inserted_rows,
        total_bytes=sum(
            f.size_bytes for f in files_meta if f.size_bytes is not None
        ),
    )
    # Add tracker metrics
    metrics["key_tracker"] = matched_keys.get_metrics()
    logger.info("Merge operation completed", strategy=strategy, metrics=metrics)

    return MergeResult(
        strategy=strategy,
        source_count=source_table.num_rows,
        target_count_before=target_count_before,
        target_count_after=target_count_before + inserted_rows,
        inserted=inserted_rows,
        updated=updated_rows if strategy != "insert" else 0,
        deleted=0,
        files=files_meta,
        rewritten_files=rewritten_files,
        inserted_files=inserted_files,
        preserved_files=preserved_files,
        metrics=metrics,
    )

fsspeckit.datasets.PyarrowDatasetIO.read_parquet

read_parquet(
    path: str,
    columns: list[str] | None = None,
    filters: Any | None = None,
    use_threads: bool = True,
) -> Table

Read parquet file(s) using PyArrow.

Parameters:

Name	Type	Description	Default
path	str	Path to parquet file or directory	required
columns	list[str] | None	Optional list of columns to read	None
filters	Any | None	Optional row filter expression. Accepts PyArrow compute expressions, DNF tuples, or SQL-like strings (converted to PyArrow expressions). Note: DuckDB backend only accepts SQL WHERE clause strings.. Accepts PyArrow compute expressions, DNF tuples, or SQL-like strings (converted to PyArrow expressions). Note: DuckDB backend only accepts SQL WHERE clause strings.	None
use_threads	bool	Whether to use parallel reading (default: True)	True

Returns:

Type	Description
Table	PyArrow table containing the data

Example

io = PyarrowDatasetIO()
table = io.read_parquet("/path/to/file.parquet")

# With column selection
table = io.read_parquet(
    "/path/to/data/",
    columns=["id", "name", "value"]
)

Source code in src/fsspeckit/datasets/pyarrow/io.py

def read_parquet(
    self,
    path: str,
    columns: list[str] | None = None,
    filters: Any | None = None,
    use_threads: bool = True,
) -> pa.Table:
    """Read parquet file(s) using PyArrow.

    Args:
        path: Path to parquet file or directory
        columns: Optional list of columns to read
        filters: Optional row filter expression. Accepts PyArrow compute expressions,
            DNF tuples, or SQL-like strings (converted to PyArrow expressions).
            Note: DuckDB backend only accepts SQL WHERE clause strings.. Accepts PyArrow compute expressions,
            DNF tuples, or SQL-like strings (converted to PyArrow expressions).
            Note: DuckDB backend only accepts SQL WHERE clause strings.
        use_threads: Whether to use parallel reading (default: True)

    Returns:
        PyArrow table containing the data

    Example:
        ```python
        io = PyarrowDatasetIO()
        table = io.read_parquet("/path/to/file.parquet")

        # With column selection
        table = io.read_parquet(
            "/path/to/data/",
            columns=["id", "name", "value"]
        )
        ```
    """
    _import_pyarrow()
    import pyarrow.dataset as ds
    import pyarrow.parquet as pq

    path = self._normalize_path(path, operation="read")
    filters = self._normalize_filters(filters, path)

    # Check if path is a single file or directory
    if self._filesystem.isfile(path):
        return pq.read_table(
            path,
            filesystem=self._filesystem,
            columns=columns,
            filters=filters,
            use_threads=use_threads,
        )
    else:
        # Dataset directory
        dataset = ds.dataset(
            path,
            filesystem=self._filesystem,
            format="parquet",
        )
        return dataset.to_table(
            columns=columns,
            filter=filters,
        )

fsspeckit.datasets.PyarrowDatasetIO.write_dataset

write_dataset(
    data: Table | list[Table],
    path: str,
    *,
    mode: Literal["append", "overwrite"] = "append",
    basename_template: str | None = None,
    schema: Schema | None = None,
    partition_by: str | list[str] | None = None,
    partitioning_flavor: Literal["hive", "directory"]
    | None = None,
    compression: str | None = "snappy",
    max_rows_per_file: int | None = 5000000,
    row_group_size: int | None = 500000,
) -> "WriteDatasetResult"

Write a parquet dataset and return per-file metadata.

Parameters:

Name	Type	Description	Default
data	Table | list[Table]	PyArrow table or list of tables to write	required
path	str	Target dataset path	required
mode	Literal['append', 'overwrite']	Write mode - "append" or "overwrite"	'append'
basename_template	str | None	Template for output filenames	None
schema	Schema | None	Optional schema to enforce on the data	None
partition_by	str | list[str] | None	Column(s) to partition by	None
partitioning_flavor	Literal['hive', 'directory'] | None	Partitioning style - "hive" for Hive-style (col=val) or "directory" for simple directory partitioning. Default is hive when partition_by is provided.	None
compression	str | None	Compression codec (default: snappy)	'snappy'
max_rows_per_file	int | None	Maximum rows per output file	5000000
row_group_size	int | None	Rows per row group in parquet files	500000

Returns:

Type	Description
'WriteDatasetResult'	WriteDatasetResult with file metadata and statistics

Source code in src/fsspeckit/datasets/pyarrow/io.py

def write_dataset(
    self,
    data: pa.Table | list[pa.Table],
    path: str,
    *,
    mode: Literal["append", "overwrite"] = "append",
    basename_template: str | None = None,
    schema: pa.Schema | None = None,
    partition_by: str | list[str] | None = None,
    partitioning_flavor: Literal["hive", "directory"] | None = None,
    compression: str | None = "snappy",
    max_rows_per_file: int | None = 5_000_000,
    row_group_size: int | None = 500_000,
) -> "WriteDatasetResult":
    """Write a parquet dataset and return per-file metadata.

    Args:
        data: PyArrow table or list of tables to write
        path: Target dataset path
        mode: Write mode - "append" or "overwrite"
        basename_template: Template for output filenames
        schema: Optional schema to enforce on the data
        partition_by: Column(s) to partition by
        partitioning_flavor: Partitioning style - "hive" for Hive-style (col=val)
            or "directory" for simple directory partitioning. Default is hive
            when partition_by is provided.
        compression: Compression codec (default: snappy)
        max_rows_per_file: Maximum rows per output file
        row_group_size: Rows per row group in parquet files

    Returns:
        WriteDatasetResult with file metadata and statistics
    """
    import uuid

    import pyarrow as pa
    import pyarrow.dataset as pds
    import pyarrow.parquet as pq

    from fsspeckit.common.security import validate_compression_codec
    from fsspeckit.datasets.write_result import (
        FileWriteMetadata,
        WriteDatasetResult,
    )

    path = self._normalize_path(path, operation="write")
    validate_compression_codec(compression)

    self._validate_write_mode(mode)
    row_group_size = self._validate_write_parameters(
        max_rows_per_file,
        row_group_size,
    )

    table = self._combine_tables(data)

    if schema is not None:
        from fsspeckit.datasets.schema import cast_schema

        table = cast_schema(table, schema)

    # Ensure dataset directory exists.
    self._filesystem.mkdirs(path, exist_ok=True)

    if mode == "overwrite":
        self._clear_dataset_parquet_only(path)

    if mode == "append" and (
        basename_template is None or basename_template == "part-{i}.parquet"
    ):
        unique_id = uuid.uuid4().hex[:16]
        basename_template = f"part-{unique_id}-{{i}}.parquet"

    if basename_template is None:
        basename_template = "part-{i}.parquet"

    written: list[pds.WrittenFile] = []
    file_options = pds.ParquetFileFormat().make_write_options(
        compression=compression
    )

    write_options: dict[str, object] = {
        "basename_template": basename_template,
        "max_rows_per_file": max_rows_per_file,
        "max_rows_per_group": row_group_size,
        "existing_data_behavior": "overwrite_or_ignore",
    }
    if partition_by is not None:
        # Handle partitioning flavor
        effective_flavor = partitioning_flavor or "hive"
        if effective_flavor == "hive":
            # Create Hive partitioning with proper schema
            partition_cols = (
                [partition_by] if isinstance(partition_by, str) else partition_by
            )
            partition_schema = pa.schema(
                [
                    (col, table.schema.field(col).type)
                    for col in partition_cols
                    if col in table.schema.names
                ]
            )
            write_options["partitioning"] = pds.partitioning(
                partition_schema, flavor="hive"
            )
        else:
            # Simple directory partitioning
            write_options["partitioning"] = partition_by

    pds.write_dataset(
        table,
        base_dir=path,
        filesystem=self._filesystem,
        format="parquet",
        file_options=file_options,
        file_visitor=written.append,
        **write_options,
    )

    files: list[FileWriteMetadata] = []
    for wf in written:
        row_count = 0
        if wf.metadata is not None:
            row_count = int(wf.metadata.num_rows)
        else:
            try:
                row_count = int(
                    pq.read_metadata(wf.path, filesystem=self._filesystem).num_rows
                )
            except (IOError, RuntimeError, ValueError) as e:
                logger.warning(
                    "failed_to_read_metadata",
                    path=wf.path,
                    error=str(e),
                    operation="write_dataset",
                )
                row_count = 0

        size_bytes = None
        if wf.size is not None:
            size_bytes = int(wf.size)
        else:
            try:
                size_bytes = int(self._filesystem.size(wf.path))
            except (IOError, RuntimeError) as e:
                logger.warning(
                    "failed_to_get_file_size",
                    path=wf.path,
                    error=str(e),
                    operation="write_dataset",
                )
                size_bytes = None

        files.append(
            FileWriteMetadata(
                path=wf.path,
                row_count=row_count,
                size_bytes=size_bytes,
            )
        )

    return WriteDatasetResult(
        files=files,
        total_rows=sum(f.row_count for f in files),
        mode=mode,
        backend="pyarrow",
    )

fsspeckit.datasets.PyarrowDatasetIO.write_parquet

write_parquet(
    data: Table | list[Table],
    path: str,
    compression: str | None = "snappy",
    row_group_size: int | None = None,
    use_threads: bool = False,
) -> None

Write parquet file using PyArrow.

Parameters:

Name	Type	Description	Default
data	Table | list[Table]	PyArrow table or list of tables to write	required
path	str	Output file path	required
compression	str | None	Compression codec to use (default: snappy)	'snappy'
row_group_size	int | None	Rows per row group	None
use_threads	bool	Whether to use parallel writing (ignored by PyArrow)	False

Example

import pyarrow as pa

table = pa.table({'a': [1, 2, 3], 'b': ['x', 'y', 'z']})
io = PyarrowDatasetIO()
io.write_parquet(table, "/tmp/data.parquet")

Source code in src/fsspeckit/datasets/pyarrow/io.py

def write_parquet(
    self,
    data: pa.Table | list[pa.Table],
    path: str,
    compression: str | None = "snappy",
    row_group_size: int | None = None,
    use_threads: bool = False,
) -> None:
    """Write parquet file using PyArrow.

    Args:
        data: PyArrow table or list of tables to write
        path: Output file path
        compression: Compression codec to use (default: snappy)
        row_group_size: Rows per row group
        use_threads: Whether to use parallel writing (ignored by PyArrow)

    Example:
        ```python
        import pyarrow as pa

        table = pa.table({'a': [1, 2, 3], 'b': ['x', 'y', 'z']})
        io = PyarrowDatasetIO()
        io.write_parquet(table, "/tmp/data.parquet")
        ```
    """
    import pyarrow.parquet as pq

    from fsspeckit.common.security import validate_compression_codec

    path = self._normalize_path(path, operation="write")
    validate_compression_codec(compression)
    _ = use_threads

    data = self._combine_tables(data)

    pq.write_table(
        data,
        path,
        filesystem=self._filesystem,
        compression=compression,
        row_group_size=row_group_size,
    )

Functions

fsspeckit.datasets.normalize_path

normalize_path(
    path: str, filesystem: AbstractFileSystem
) -> str

Normalize path based on filesystem type.

This function now delegates to core/filesystem/paths.normalize_path to provide unified path normalization across the codebase. The delegation preserves all existing behavior while eliminating duplicate normalization logic.

Parameters:

Name	Type	Description	Default
path	str	The path to normalize.	required
filesystem	AbstractFileSystem	The filesystem instance.	required

Returns:

Type	Description
str	The normalized path.

Source code in src/fsspeckit/datasets/path_utils.py

def normalize_path(path: str, filesystem: AbstractFileSystem) -> str:
    """Normalize path based on filesystem type.

    This function now delegates to core/filesystem/paths.normalize_path to provide
    unified path normalization across the codebase. The delegation preserves all
    existing behavior while eliminating duplicate normalization logic.

    Args:
        path: The path to normalize.
        filesystem: The filesystem instance.

    Returns:
        The normalized path.
    """
    # Delegate to core normalize_path with filesystem
    return core_normalize_path(path, filesystem=filesystem)

fsspeckit.datasets.validate_dataset_path

validate_dataset_path(
    path: str,
    filesystem: AbstractFileSystem,
    operation: str,
) -> None

Comprehensive path validation for dataset operations.

Parameters:

Name	Type	Description	Default
path	str	The path to validate.	required
filesystem	AbstractFileSystem	The filesystem instance.	required
operation	str	The operation being performed ('read', 'write', 'merge', etc.)	required

Raises:

Type	Description
DatasetPathError	If the path is invalid for the given operation.

Source code in src/fsspeckit/datasets/path_utils.py

def validate_dataset_path(
    path: str, filesystem: AbstractFileSystem, operation: str
) -> None:
    """Comprehensive path validation for dataset operations.

    Args:
        path: The path to validate.
        filesystem: The filesystem instance.
        operation: The operation being performed ('read', 'write', 'merge', etc.)

    Raises:
        DatasetPathError: If the path is invalid for the given operation.
    """
    logger.debug("validating_path", path=path, operation=operation)

    # Basic security validation
    try:
        security_validate_path(path)
    except ValueError as e:
        raise DatasetPathError(
            str(e), operation=operation, details={"path": path}
        ) from e

    # Check path exists for read operations
    if operation in ["read", "merge"]:
        if not filesystem.exists(path):
            raise DatasetPathError(
                f"Dataset path does not exist: {path}",
                operation=operation,
                details={"path": path},
            )

    # Check parent directory exists for write operations
    if operation in ["write", "merge"]:
        try:
            # fsspec's AbstractFileSystem has _parent in recent versions
            parent = filesystem._parent(path)
        except (AttributeError, TypeError, ValueError):
            parent = None

        if (
            parent
            and parent not in ["", "/", "."]
            and not filesystem.exists(parent)
        ):
            if operation == "write":
                created = False
                for method_name, kwargs in (
                    ("mkdirs", {"exist_ok": True}),
                    ("makedirs", {"exist_ok": True}),
                    ("mkdir", {"create_parents": True}),
                ):
                    method = getattr(filesystem, method_name, None)
                    if method is None:
                        continue
                    try:
                        method(parent, **kwargs)
                        created = True
                        break
                    except TypeError:
                        try:
                            method(parent)
                            created = True
                            break
                        except Exception:
                            continue
                    except Exception:
                        continue

                if created and filesystem.exists(parent):
                    pass
                else:
                    raise DatasetPathError(
                        f"Parent directory does not exist: {parent}",
                        operation=operation,
                        details={"path": path, "parent": parent},
                    )
            else:
                raise DatasetPathError(
                    f"Parent directory does not exist: {parent}",
                    operation=operation,
                    details={"path": path, "parent": parent},
                )

    # Validate path format
    if "://" in path:
        # Remote path - validate protocol
        protocol = path.split("://")[0].lower()
        if protocol not in SUPPORTED_PROTOCOLS:
            raise DatasetPathError(
                f"Unsupported protocol: {protocol}",
                operation=operation,
                details={
                    "path": path,
                    "protocol": protocol,
                    "supported_protocols": SUPPORTED_PROTOCOLS,
                },
            )

fsspeckit.datasets.collect_dataset_stats_pyarrow

collect_dataset_stats_pyarrow(
    path: str,
    filesystem: AbstractFileSystem | None = None,
    partition_filter: list[str] | None = None,
) -> dict[str, Any]

Collect file-level statistics for a parquet dataset using shared core logic.

This function delegates to the shared fsspeckit.core.maintenance.collect_dataset_stats function, ensuring consistent dataset discovery and statistics across both DuckDB and PyArrow backends.

The helper walks the given dataset directory on the provided filesystem, discovers parquet files (recursively), and returns basic statistics:

• Per-file path, size in bytes, and number of rows
• Aggregated total bytes and total rows

The function is intentionally streaming/metadata-driven and never materializes the full dataset as a single :class:pyarrow.Table.

Parameters:

Name	Type	Description	Default
path	str	Root directory of the parquet dataset.	required
filesystem	AbstractFileSystem | None	Optional fsspec filesystem. If omitted, a local "file" filesystem is used.	None
partition_filter	list[str] | None	Optional list of partition prefix filters (e.g. ["date=2025-11-04"]). Only files whose path relative to path starts with one of these prefixes are included.	None

Returns:

Type	Description
dict[str, Any]	Dict with keys:
dict[str, Any]	files: list of {"path", "size_bytes", "num_rows"} dicts
dict[str, Any]	total_bytes: sum of file sizes
dict[str, Any]	total_rows: sum of row counts

Raises:

Type	Description
FileNotFoundError	If the path does not exist or no parquet files match the optional partition filter.

Note

This is a thin wrapper around the shared core function. See :func:fsspeckit.core.maintenance.collect_dataset_stats for the authoritative implementation.

Source code in src/fsspeckit/datasets/pyarrow/dataset.py

def collect_dataset_stats_pyarrow(
    path: str,
    filesystem: AbstractFileSystem | None = None,
    partition_filter: list[str] | None = None,
) -> dict[str, Any]:
    """Collect file-level statistics for a parquet dataset using shared core logic.

    This function delegates to the shared ``fsspeckit.core.maintenance.collect_dataset_stats``
    function, ensuring consistent dataset discovery and statistics across both DuckDB
    and PyArrow backends.

    The helper walks the given dataset directory on the provided filesystem,
    discovers parquet files (recursively), and returns basic statistics:

    - Per-file path, size in bytes, and number of rows
    - Aggregated total bytes and total rows

    The function is intentionally streaming/metadata-driven and never
    materializes the full dataset as a single :class:`pyarrow.Table`.

    Args:
        path: Root directory of the parquet dataset.
        filesystem: Optional fsspec filesystem. If omitted, a local "file"
            filesystem is used.
        partition_filter: Optional list of partition prefix filters
            (e.g. ["date=2025-11-04"]). Only files whose path relative to
            ``path`` starts with one of these prefixes are included.

    Returns:
        Dict with keys:

        - ``files``: list of ``{"path", "size_bytes", "num_rows"}`` dicts
        - ``total_bytes``: sum of file sizes
        - ``total_rows``: sum of row counts

    Raises:
        FileNotFoundError: If the path does not exist or no parquet files
            match the optional partition filter.

    Note:
        This is a thin wrapper around the shared core function. See
        :func:`fsspeckit.core.maintenance.collect_dataset_stats` for the
        authoritative implementation.
    """
    from fsspeckit.core.maintenance import collect_dataset_stats

    return collect_dataset_stats(
        path=path,
        filesystem=filesystem,
        partition_filter=partition_filter,
    )

fsspeckit.datasets.cast_schema

cast_schema(table: Table, schema: Schema) -> Table

Cast a PyArrow table to a given schema, updating the schema to match the table's columns.

Parameters:

Name	Type	Description	Default
table	Table	The PyArrow table to cast.	required
schema	Schema	The target schema to cast table to.	required

Returns:

Type	Description
Table	pa.Table: A new PyArrow table with the specified schema.

Source code in src/fsspeckit/datasets/schema.py

def cast_schema(table: pa.Table, schema: pa.Schema) -> pa.Table:
    """
    Cast a PyArrow table to a given schema, updating the schema to match the table's columns.

    Args:
        table (pa.Table): The PyArrow table to cast.
        schema (pa.Schema): The target schema to cast table to.

    Returns:
        pa.Table: A new PyArrow table with the specified schema.
    """
    # Append missing columns as nulls
    table_columns = set(table.schema.names)
    for field in schema:
        if field.name not in table_columns:
            table = table.append_column(
                field.name, pa.nulls(table.num_rows, type=field.type)
            )

    # Identify extra columns in table that are not in schema
    extra_fields = []
    schema_names = set(schema.names)
    for field in table.schema:
        if field.name not in schema_names:
            extra_fields.append(field)

    # Create target schema: requested schema + extra columns
    if extra_fields:
        target_schema = pa.schema(list(schema) + extra_fields, metadata=schema.metadata)
    else:
        target_schema = schema

    # Reorder table columns to match target schema
    table = table.select(target_schema.names)

    return table.cast(target_schema)

fsspeckit.datasets.convert_large_types_to_normal

convert_large_types_to_normal(schema: Schema) -> Schema

Convert large types in a PyArrow schema to their standard types.

Parameters:

Name	Type	Description	Default
schema	Schema	The PyArrow schema to convert.	required

Returns:

Type	Description
Schema	pa.Schema: A new PyArrow schema with large types converted to standard types.

Source code in src/fsspeckit/datasets/schema.py

def convert_large_types_to_normal(schema: pa.Schema) -> pa.Schema:
    """
    Convert large types in a PyArrow schema to their standard types.

    Args:
        schema (pa.Schema): The PyArrow schema to convert.

    Returns:
        pa.Schema: A new PyArrow schema with large types converted to standard types.
    """
    # Define mapping of large types to standard types
    type_mapping = {
        pa.large_string(): pa.string(),
        pa.large_binary(): pa.binary(),
        pa.large_utf8(): pa.utf8(),
        pa.large_list(pa.null()): pa.list_(pa.null()),
        pa.large_list_view(pa.null()): pa.list_view(pa.null()),
    }
    # Convert fields
    new_fields = []
    for field in schema:
        field_type = field.type
        # Check if type exists in mapping
        if field_type in type_mapping:
            new_field = pa.field(
                name=field.name,
                type=type_mapping[field_type],
                nullable=field.nullable,
                metadata=field.metadata,
            )
            new_fields.append(new_field)
        # Handle large lists with nested types
        elif isinstance(field_type, pa.LargeListType):
            new_field = pa.field(
                name=field.name,
                type=pa.list_(
                    type_mapping[field_type.value_type]
                    if field_type.value_type in type_mapping
                    else field_type.value_type
                ),
                nullable=field.nullable,
                metadata=field.metadata,
            )
            new_fields.append(new_field)
        # Handle dictionary with large_string, large_utf8, or large_binary values
        elif isinstance(field_type, pa.DictionaryType):
            new_field = pa.field(
                name=field.name,
                type=pa.dictionary(
                    field_type.index_type,
                    type_mapping[field_type.value_type]
                    if field_type.value_type in type_mapping
                    else field_type.value_type,
                    field_type.ordered,
                ),
                metadata=field.metadata,
            )
            new_fields.append(new_field)
        else:
            new_fields.append(field)

    return pa.schema(new_fields)

fsspeckit.datasets.opt_dtype_pa

opt_dtype_pa(
    table: Table,
    strict: bool = False,
    columns: list[str] | None = None,
) -> Table

Optimize dtypes in a PyArrow table based on data analysis.

This function analyzes the data in each column and attempts to downcast to more appropriate types (e.g., int64 -> int32, float64 -> float32, string -> int/bool where applicable).

Parameters:

Name	Type	Description	Default
table	Table	PyArrow table to optimize	required
strict	bool	Whether to use strict type checking	False
columns	list[str] | None	List of columns to optimize (None for all)	None

Returns:

Type	Description
Table	Table with optimized dtypes

Example

import pyarrow as pa

table = pa.table(
    {
        "a": pa.array([1, 2, 3], type=pa.int64()),
        "b": pa.array([1.0, 2.0, 3.0], type=pa.float64()),
    },
)
optimized = opt_dtype(table)
print(optimized.column(0).type)  # DataType(int32)
print(optimized.column(1).type)  # DataType(float32)

Source code in src/fsspeckit/datasets/schema.py

def opt_dtype(
    table: pa.Table,
    strict: bool = False,
    columns: list[str] | None = None,
) -> pa.Table:
    """Optimize dtypes in a PyArrow table based on data analysis.

    This function analyzes the data in each column and attempts to downcast
    to more appropriate types (e.g., int64 -> int32, float64 -> float32,
    string -> int/bool where applicable).

    Args:
        table: PyArrow table to optimize
        strict: Whether to use strict type checking
        columns: List of columns to optimize (None for all)

    Returns:
        Table with optimized dtypes

    Example:
        ```python
        import pyarrow as pa

        table = pa.table(
            {
                "a": pa.array([1, 2, 3], type=pa.int64()),
                "b": pa.array([1.0, 2.0, 3.0], type=pa.float64()),
            },
        )
        optimized = opt_dtype(table)
        print(optimized.column(0).type)  # DataType(int32)
        print(optimized.column(1).type)  # DataType(float32)
        ```
    """
    from fsspeckit.common.parallel import run_parallel

    if columns is None:
        columns = table.column_names

    # Process columns in parallel
    results = run_parallel(
        _process_column_for_opt_dtype,
        [(table, col, strict) for col in columns],
        backend="threading",
        n_jobs=-1,
    )

    # Build new table with optimized columns
    new_columns = {}
    for col_name, optimized_array in results:
        new_columns[col_name] = optimized_array

    # Keep non-optimized columns as-is
    for col_name in table.column_names:
        if col_name not in new_columns:
            new_columns[col_name] = table.column(col_name)

    return pa.table(new_columns)

fsspeckit.datasets.unify_schemas_pa

unify_schemas_pa(
    schemas: list[Schema],
    use_large_dtypes: bool = False,
    timezone: str | None = None,
    standardize_timezones: bool = True,
    verbose: bool = False,
    remove_conflicting_columns: bool = False,
) -> Schema

Unify a list of PyArrow schemas into a single schema using intelligent conflict resolution.

Parameters:

Name	Type	Description	Default
schemas	list[Schema]	List of PyArrow schemas to unify.	required
use_large_dtypes	bool	If True, keep large types like large_string.	False
timezone	str | None	If specified, standardize all timestamp columns to this timezone. If "auto", use the most frequent timezone across schemas. If None, remove timezone from all timestamp columns.	None
standardize_timezones	bool	If True, standardize all timestamp columns to most frequent timezone.	True
verbose	bool	If True, print conflict resolution details for debugging.	False
remove_conflicting_columns	bool	If True, allows removal of columns with type conflicts as a fallback strategy instead of converting them (e.g. to string). Defaults to False.	False

Returns:

Type	Description
Schema	pa.Schema: A unified PyArrow schema.

Raises:

Type	Description
ValueError	If no schemas are provided.

Source code in src/fsspeckit/datasets/schema.py

def unify_schemas(
    schemas: list[pa.Schema],
    use_large_dtypes: bool = False,
    timezone: str | None = None,
    standardize_timezones: bool = True,
    verbose: bool = False,
    remove_conflicting_columns: bool = False,
) -> pa.Schema:
    """
    Unify a list of PyArrow schemas into a single schema using intelligent conflict resolution.

    Args:
        schemas (list[pa.Schema]): List of PyArrow schemas to unify.
        use_large_dtypes (bool): If True, keep large types like large_string.
        timezone (str | None): If specified, standardize all timestamp columns to this timezone.
            If "auto", use the most frequent timezone across schemas.
            If None, remove timezone from all timestamp columns.
        standardize_timezones (bool): If True, standardize all timestamp columns to most frequent timezone.
        verbose (bool): If True, print conflict resolution details for debugging.
        remove_conflicting_columns (bool): If True, allows removal of columns with type conflicts as a fallback
            strategy instead of converting them (e.g. to string). Defaults to False.

    Returns:
        pa.Schema: A unified PyArrow schema.

    Raises:
        ValueError: If no schemas are provided.
    """
    if not schemas:
        raise ValueError("At least one schema must be provided for unification")

    # Early exit for single schema
    unique_schemas = _unique_schemas(schemas)
    # Keep a copy of original unique schemas for timezone detection if needed
    original_unique_schemas = unique_schemas

    if len(unique_schemas) == 1:
        result_schema = unique_schemas[0]
        if standardize_timezones:
            target_tz = timezone if timezone is not None else "auto"
            result_schema = standardize_schema_timezones([result_schema], target_tz)[0]
        return (
            result_schema
            if use_large_dtypes
            else convert_large_types_to_normal(result_schema)
        )

    # Step 1: Find and resolve conflicts first
    conflicts = _find_conflicting_fields(unique_schemas)
    if conflicts and verbose:
        _log_conflict_summary(conflicts, verbose)

    if conflicts:
        # Pre-process incompatible columns if removal is requested
        if remove_conflicting_columns:
             # Determine compatibility status first
             _normalize_schema_types(unique_schemas, conflicts, handle_incompatible=False)

             incompatible_fields = {
                 name for name, info in conflicts.items() 
                 if not info["compatible"]
             }

             if incompatible_fields:
                 # Remove conflicting columns from all schemas
                 unique_schemas = [
                     pa.schema([f for f in s if f.name not in incompatible_fields], metadata=s.metadata)
                     for s in unique_schemas
                 ]
                 # Re-evaluate conflicts? No, we removed them.
                 # Update conflicts dict?
                 for field in incompatible_fields:
                     del conflicts[field]

        # Normalize schemas using intelligent promotion rules
        unique_schemas = _normalize_schema_types(
            unique_schemas, conflicts, handle_incompatible=not remove_conflicting_columns
        )

    # Step 2: Attempt unification with conflict-resolved schemas
    try:
        unified_schema = pa.unify_schemas(unique_schemas, promote_options="permissive")

        # Step 3: Apply timezone standardization to the unified result
        if standardize_timezones:
            target_tz = timezone if timezone is not None else "auto"

            if target_tz == "auto":
                # Calculate majority from ORIGINAL input schemas, not the unified result
                # (since normalization might have shifted them)
                dom = dominant_timezone_per_column(original_unique_schemas)

                # Apply derived timezones to unified schema
                fields = []
                for field in unified_schema:
                    if pa.types.is_timestamp(field.type) and field.name in dom:
                        unit, tz = dom[field.name]
                        # Use unit from unified field to preserve precision promotion
                        current_unit = field.type.unit
                        fields.append(field.with_type(pa.timestamp(current_unit, tz)))
                    else:
                        fields.append(field)
                unified_schema = pa.schema(fields, unified_schema.metadata)
            else:
                # Explicit timezone
                unified_schema = standardize_schema_timezones([unified_schema], target_tz)[0]

        return (
            unified_schema
            if use_large_dtypes
            else convert_large_types_to_normal(unified_schema)
        )

    except (pa.ArrowInvalid, pa.ArrowTypeError) as e:
        # Step 4: Intelligent fallback strategies
        if verbose:
            logger.debug("Primary unification failed: %s", e)
            logger.debug("Attempting fallback strategies...")

        # Fallback 1: Try aggressive string conversion for remaining conflicts
        try:
            fallback_schema = _aggressive_fallback_unification(unique_schemas)
            if standardize_timezones:
                fallback_schema = standardize_schema_timezones(
                    [fallback_schema], timezone
                )[0]
            if verbose:
                logger.debug("✓ Aggressive fallback succeeded")
            return (
                fallback_schema
                if use_large_dtypes
                else convert_large_types_to_normal(fallback_schema)
            )

        except (pa.ArrowInvalid, pa.ArrowTypeError, ValueError) as e:
            if verbose:
                logger.debug("✗ Aggressive fallback failed: %s", str(e), exc_info=True)

        # Fallback 2: Remove conflicting fields (if enabled)
        if remove_conflicting_columns:
            try:
                non_conflicting_schema = _remove_conflicting_fields(unique_schemas)
                if standardize_timezones:
                    non_conflicting_schema = standardize_schema_timezones(
                        [non_conflicting_schema], timezone
                    )[0]
                if verbose:
                    logger.debug("✓ Remove conflicting fields fallback succeeded")
                return (
                    non_conflicting_schema
                    if use_large_dtypes
                    else convert_large_types_to_normal(non_conflicting_schema)
                )

            except (pa.ArrowInvalid, pa.ArrowTypeError, ValueError) as e:
                if verbose:
                    logger.debug(
                        "✗ Remove conflicting fields fallback failed: %s",
                        str(e),
                        exc_info=True,
                    )

        # Fallback 3: Remove problematic fields that can't be unified
        try:
            minimal_schema = _remove_problematic_fields(unique_schemas)
            if standardize_timezones:
                minimal_schema = standardize_schema_timezones(
                    [minimal_schema], timezone
                )[0]
            if verbose:
                logger.debug("✓ Minimal schema (removed problematic fields) succeeded")
            return (
                minimal_schema
                if use_large_dtypes
                else convert_large_types_to_normal(minimal_schema)
            )

        except (pa.ArrowInvalid, pa.ArrowTypeError, ValueError) as e:
            if verbose:
                logger.debug(
                    "✗ Minimal schema fallback failed: %s", str(e), exc_info=True
                )

        # Fallback 4: Return first schema as last resort
        if verbose:
            logger.debug("✗ All fallback strategies failed, returning first schema")

        first_schema = unique_schemas[0]
        if standardize_timezones:
            first_schema = standardize_schema_timezones([first_schema], timezone)[0]
        return (
            first_schema
            if use_large_dtypes
            else convert_large_types_to_normal(first_schema)
        )

fsspeckit.datasets.opt_dtype_pl

opt_dtype_pl(
    df: DataFrame,
    include: str | list[str] | None = None,
    exclude: str | list[str] | None = None,
    time_zone: str | None = None,
    shrink_numerics: bool = False,
    allow_unsigned: bool = True,
    allow_null: bool = True,
    sample_size: int | None = 1024,
    sample_method: SampleMethod = "first",
    strict: bool = False,
    *,
    force_timezone: str | None = None,
) -> DataFrame

Optimize data types of a Polars DataFrame for performance and memory efficiency.

This function analyzes each column and converts it to the most appropriate data type based on content, handling string-to-type conversions and numeric type downcasting.

Parameters:

Name	Type	Description	Default
df	DataFrame	The Polars DataFrame to optimize.	required
include	str | list[str] | None	Column(s) to include in optimization (default: all columns).	None
exclude	str | list[str] | None	Column(s) to exclude from optimization.	None
time_zone	str | None	Optional time zone hint during datetime parsing.	None
shrink_numerics	bool	Whether to downcast numeric types when possible.	False
allow_unsigned	bool	Whether to allow unsigned integer types.	True
allow_null	bool	Whether to allow columns with all null values to be cast to Null type.	True
sample_size	int | None	Maximum number of cleaned values to inspect for regex-based inference. Use None to inspect the entire column.	1024
sample_method	SampleMethod	Which subset to inspect ("first" or "random").	'first'
strict	bool	If True, will raise an error if any column cannot be optimized.	False
force_timezone	str | None	If set, ensure all parsed datetime columns end up with this timezone.	None

Returns:

Type	Description
DataFrame	DataFrame with optimized data types.

Source code in src/fsspeckit/datasets/polars.py

def opt_dtype(
    df: pl.DataFrame,
    include: str | list[str] | None = None,
    exclude: str | list[str] | None = None,
    time_zone: str | None = None,
    shrink_numerics: bool = False,
    allow_unsigned: bool = True,
    allow_null: bool = True,
    sample_size: int | None = 1024,
    sample_method: SampleMethod = "first",
    strict: bool = False,
    *,
    force_timezone: str | None = None,
) -> pl.DataFrame:
    """
    Optimize data types of a Polars DataFrame for performance and memory efficiency.

    This function analyzes each column and converts it to the most appropriate
    data type based on content, handling string-to-type conversions and
    numeric type downcasting.

    Args:
        df: The Polars DataFrame to optimize.
        include: Column(s) to include in optimization (default: all columns).
        exclude: Column(s) to exclude from optimization.
        time_zone: Optional time zone hint during datetime parsing.
        shrink_numerics: Whether to downcast numeric types when possible.
        allow_unsigned: Whether to allow unsigned integer types.
        allow_null: Whether to allow columns with all null values to be cast to Null type.
        sample_size: Maximum number of cleaned values to inspect for regex-based inference. Use None to inspect the entire column.
        sample_method: Which subset to inspect (`"first"` or `"random"`).
        strict: If True, will raise an error if any column cannot be optimized.
        force_timezone: If set, ensure all parsed datetime columns end up with this timezone.

    Returns:
        DataFrame with optimized data types.
    """
    if sample_method not in ("first", "random"):
        raise ValueError("sample_method must be 'first' or 'random'")

    if isinstance(df, pl.LazyFrame):
        return opt_dtype(
            df.collect(),
            include=include,
            exclude=exclude,
            time_zone=time_zone,
            shrink_numerics=shrink_numerics,
            allow_unsigned=allow_unsigned,
            allow_null=allow_null,
            sample_size=sample_size,
            sample_method=sample_method,
            strict=strict,
            force_timezone=force_timezone,
        ).lazy()

    # Normalize include/exclude parameters
    if isinstance(include, str):
        include = [include]
    if isinstance(exclude, str):
        exclude = [exclude]

    # Determine columns to process
    cols_to_process = df.columns
    if include:
        cols_to_process = [col for col in include if col in df.columns]
    if exclude:
        cols_to_process = [col for col in cols_to_process if col not in exclude]

    # Generate optimization expressions for all columns
    expressions = []
    for col_name in cols_to_process:
        try:
            expressions.append(
                _get_column_expr(
                    df,
                    col_name,
                    shrink_numerics,
                    allow_unsigned,
                    allow_null,
                    time_zone,
                    force_timezone,
                    sample_size,
                    sample_method,
                    strict,
                )
            )
        except Exception as e:
            if strict:
                raise e
            # If strict mode is off, just keep the original column
            continue

    # Apply all transformations at once if any exist
    return df if not expressions else df.with_columns(expressions)