fsspeckit.datasets.pyarrow API Reference

pyarrow

PyArrow dataset integration for fsspeckit.

This package contains focused submodules for PyArrow functionality: - dataset: Dataset merge and maintenance operations - io: PyarrowDatasetIO class for dataset operations

All public APIs are re-exported here for convenient access.

Classes

fsspeckit.datasets.pyarrow.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.pyarrow.PyarrowDatasetIO.filesystem property

filesystem: AbstractFileSystem

Return the filesystem instance.

Functions

fsspeckit.datasets.pyarrow.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.pyarrow.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.pyarrow.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.pyarrow.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.pyarrow.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.pyarrow.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,
    )

fsspeckit.datasets.pyarrow.MemoryMonitor

MemoryMonitor(
    max_pyarrow_mb: int = 2048,
    max_process_memory_mb: int | None = None,
    min_system_available_mb: int = 512,
)

Monitors both PyArrow and system memory usage.

This class provides tiered memory pressure monitoring to help prevent out-of-memory errors during large dataset operations. It combines PyArrow-specific allocation tracking with general system memory monitoring.

Initialize the memory monitor.

Parameters:

Name	Type	Description	Default
max_pyarrow_mb	int	Maximum allowed PyArrow-allocated memory in MB.	2048
max_process_memory_mb	int | None	Optional maximum total process memory (RSS) in MB.	None
min_system_available_mb	int	Minimum required system available memory in MB.	512

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

def __init__(
    self,
    max_pyarrow_mb: int = 2048,
    max_process_memory_mb: int | None = None,
    min_system_available_mb: int = 512,
):
    """Initialize the memory monitor.

    Args:
        max_pyarrow_mb: Maximum allowed PyArrow-allocated memory in MB.
        max_process_memory_mb: Optional maximum total process memory (RSS) in MB.
        min_system_available_mb: Minimum required system available memory in MB.
    """
    self.max_pyarrow_mb = max_pyarrow_mb
    self.max_process_memory_mb = max_process_memory_mb
    self.min_system_available_mb = min_system_available_mb

    if psutil is not None:
        self._process = psutil.Process(os.getpid())
    else:
        self._process = None
        if max_process_memory_mb or min_system_available_mb:
            logger.debug(
                "psutil not available; system memory thresholds will be ignored."
            )

Functions

fsspeckit.datasets.pyarrow.MemoryMonitor.check_memory_pressure

check_memory_pressure() -> MemoryPressureLevel

Evaluate current memory usage against thresholds.

Tiered thresholds: - NORMAL: < 70% of limits - WARNING: 70-90% of limits - CRITICAL: > 90% of limits - EMERGENCY: Exceeds absolute limits

Returns:

Type	Description
MemoryPressureLevel	Current MemoryPressureLevel

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

def check_memory_pressure(self) -> MemoryPressureLevel:
    """Evaluate current memory usage against thresholds.

    Tiered thresholds:
    - NORMAL: < 70% of limits
    - WARNING: 70-90% of limits
    - CRITICAL: > 90% of limits
    - EMERGENCY: Exceeds absolute limits

    Returns:
        Current MemoryPressureLevel
    """
    status = self.get_memory_status()
    pa_used = status["pyarrow_allocated_mb"]
    pa_ratio = pa_used / self.max_pyarrow_mb

    pressure = MemoryPressureLevel.NORMAL

    # Check PyArrow limits
    if pa_ratio >= 1.0:
        pressure = MemoryPressureLevel.EMERGENCY
    elif pa_ratio > 0.9:
        pressure = MemoryPressureLevel.CRITICAL
    elif pa_ratio > 0.7:
        pressure = MemoryPressureLevel.WARNING

    # Check Process limits if configured and available
    if (
        psutil is not None
        and self.max_process_memory_mb
        and "process_rss_mb" in status
    ):
        rss_used = status["process_rss_mb"]
        rss_ratio = rss_used / self.max_process_memory_mb

        if rss_ratio >= 1.0:
            pressure = max_pressure(pressure, MemoryPressureLevel.EMERGENCY)
        elif rss_ratio > 0.9:
            pressure = max_pressure(pressure, MemoryPressureLevel.CRITICAL)
        elif rss_ratio > 0.7:
            pressure = max_pressure(pressure, MemoryPressureLevel.WARNING)

    # Check System availability if available
    if psutil is not None and "system_available_mb" in status:
        sys_avail = status["system_available_mb"]
        # For system availability, we check how close we are to the minimum
        if sys_avail < self.min_system_available_mb:
            pressure = max_pressure(pressure, MemoryPressureLevel.EMERGENCY)
        elif sys_avail < self.min_system_available_mb * 1.5:
            pressure = max_pressure(pressure, MemoryPressureLevel.CRITICAL)
        elif sys_avail < self.min_system_available_mb * 2.0:
            pressure = max_pressure(pressure, MemoryPressureLevel.WARNING)

    if pressure != MemoryPressureLevel.NORMAL:
        logger.warning(
            f"Memory pressure detected: {pressure.value}. {self.get_detailed_status()}"
        )

    return pressure

fsspeckit.datasets.pyarrow.MemoryMonitor.get_detailed_status

get_detailed_status() -> str

Get detailed status string for logging and error messages.

Returns:

Type	Description
str	Formatted string with memory metrics and limits.

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

def get_detailed_status(self) -> str:
    """Get detailed status string for logging and error messages.

    Returns:
        Formatted string with memory metrics and limits.
    """
    status = self.get_memory_status()
    details = [
        f"PyArrow: {status['pyarrow_allocated_mb']:.1f}/{self.max_pyarrow_mb} MB"
    ]

    if "process_rss_mb" in status and self.max_process_memory_mb:
        details.append(
            f"Process RSS: {status['process_rss_mb']:.1f}/{self.max_process_memory_mb} MB"
        )
    elif "process_rss_mb" in status:
        details.append(f"Process RSS: {status['process_rss_mb']:.1f} MB")

    if "system_available_mb" in status:
        details.append(
            f"System Available: {status['system_available_mb']:.1f} MB (min: {self.min_system_available_mb} MB)"
        )

    return " | ".join(details)

fsspeckit.datasets.pyarrow.MemoryMonitor.get_memory_status

get_memory_status() -> dict[str, float]

Return current memory metrics in MB.

Returns:

Type	Description
dict[str, float]	Dictionary with keys: - pyarrow_allocated_mb: Bytes allocated by PyArrow - process_rss_mb: Resident Set Size of the process (if psutil available) - system_available_mb: Total system available memory (if psutil available)

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

def get_memory_status(self) -> dict[str, float]:
    """Return current memory metrics in MB.

    Returns:
        Dictionary with keys:
            - pyarrow_allocated_mb: Bytes allocated by PyArrow
            - process_rss_mb: Resident Set Size of the process (if psutil available)
            - system_available_mb: Total system available memory (if psutil available)
    """
    status = {"pyarrow_allocated_mb": pa.total_allocated_bytes() / BYTES_IN_MB}

    if psutil is not None and self._process:
        try:
            status["process_rss_mb"] = self._process.memory_info().rss / BYTES_IN_MB
            status["system_available_mb"] = (
                psutil.virtual_memory().available / BYTES_IN_MB
            )
        except Exception as e:
            # Catching generic Exception because psutil can raise various errors
            # depending on OS/permissions
            logger.warning(f"Could not access system memory info via psutil: {e}")

    return status

fsspeckit.datasets.pyarrow.MemoryMonitor.should_check_memory

should_check_memory(
    chunks_processed: int, check_interval: int = 10
) -> bool

Determine if memory should be checked based on process progress.

Parameters:

Name	Type	Description	Default
chunks_processed	int	Number of chunks processed so far.	required
check_interval	int	How often (in chunks) to perform a memory check.	10

Returns:

Type	Description
bool	True if memory should be checked now.

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

def should_check_memory(
    self, chunks_processed: int, check_interval: int = 10
) -> bool:
    """Determine if memory should be checked based on process progress.

    Args:
        chunks_processed: Number of chunks processed so far.
        check_interval: How often (in chunks) to perform a memory check.

    Returns:
        True if memory should be checked now.
    """
    return chunks_processed % check_interval == 0

fsspeckit.datasets.pyarrow.MemoryPressureLevel

Bases: Enum

Memory pressure levels for graceful degradation.

fsspeckit.datasets.pyarrow.AdaptiveKeyTracker

AdaptiveKeyTracker(
    max_exact_keys: int = 1000000,
    max_lru_keys: int = 10000000,
    false_positive_rate: float = 0.001,
)

Tracks unique keys with adaptive memory usage.

Tiers: 1. EXACT (Set): Guaranteed accuracy, low memory (up to max_exact_keys). 2. LRU (OrderedDict): Bounded memory, misses possible if evicted (up to max_lru_keys). 3. BLOOM (Probabilistic): Fixed/Scalable memory, false positives possible, no false negatives.

Automatically transitions between tiers as cardinality grows or memory pressure increases.

Initialize the tracker.

Parameters:

Name	Type	Description	Default
max_exact_keys	int	Maximum number of keys to track exactly in a set.	1000000
max_lru_keys	int	Maximum number of keys to track in LRU cache before switching to Bloom.	10000000
false_positive_rate	float	Desired false positive rate for Bloom filter.	0.001

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

def __init__(
    self,
    max_exact_keys: int = 1_000_000,
    max_lru_keys: int = 10_000_000,
    false_positive_rate: float = 0.001,
):
    """
    Initialize the tracker.

    Args:
        max_exact_keys: Maximum number of keys to track exactly in a set.
        max_lru_keys: Maximum number of keys to track in LRU cache before switching to Bloom.
        false_positive_rate: Desired false positive rate for Bloom filter.
    """
    self.max_exact_keys = max_exact_keys
    self.max_lru_keys = max_lru_keys
    self.false_positive_rate = false_positive_rate

    self._lock = threading.Lock()
    self._tier = "EXACT"
    self._exact_keys: Optional[Set[Any]] = set()
    self._lru_keys: Optional[OrderedDict] = None
    self._bloom_filter: Any = None

    self._keys_added_count = 0
    self._unique_keys_seen = 0
    self._transitions = 0

    # Performance tracking
    self._peak_estimated_mem = 0

Functions

fsspeckit.datasets.pyarrow.AdaptiveKeyTracker.__contains__

__contains__(key: Any) -> bool

Check if a key has been seen before.

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

def __contains__(self, key: Any) -> bool:
    """
    Check if a key has been seen before.
    """
    if isinstance(key, list):
        key = tuple(key)

    with self._lock:
        if self._tier == "EXACT":
            return (
                key in self._exact_keys if self._exact_keys is not None else False
            )
        elif self._tier == "LRU":
            if self._lru_keys is not None and key in self._lru_keys:
                self._lru_keys.move_to_end(key)
                return True
            return False
        elif self._tier == "BLOOM":
            if self._bloom_filter:
                # Bloom filter might have false positives but no false negatives
                return key in self._bloom_filter
            return False
        return False

fsspeckit.datasets.pyarrow.AdaptiveKeyTracker.add

add(key: Any) -> None

Add a key to the tracker.

Parameters:

Name	Type	Description	Default
key	Any	The key to track. Can be a single value or a tuple for multi-column keys.	required

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

def add(self, key: Any) -> None:
    """
    Add a key to the tracker.

    Args:
        key: The key to track. Can be a single value or a tuple for multi-column keys.
    """
    # Ensure key is hashable (lists are common in some arrow contexts)
    if isinstance(key, list):
        key = tuple(key)

    with self._lock:
        self._keys_added_count += 1

        # Periodically update memory peak (every 1000 additions to reduce overhead)
        if self._keys_added_count % 1000 == 0:
            self._update_mem_peak()

        # Use a loop to handle potential tier transitions during addition
        while True:
            if self._tier == "EXACT":
                if self._exact_keys is None:
                    # Defensive: tier mismatch, try to recover
                    self._tier = "LRU"
                    continue
                if key in self._exact_keys:
                    return
                if len(self._exact_keys) < self.max_exact_keys:
                    self._exact_keys.add(key)
                    self._unique_keys_seen += 1
                    return
                else:
                    self._transition_to_lru()
                    # Continue to next iteration to add to the new tier
            elif self._tier == "LRU":
                if self._lru_keys is None:
                    # Defensive: tier mismatch, try to recover
                    if (
                        HAS_BLOOM
                        and ScalableBloomFilter is not None
                        and self._bloom_filter
                    ):
                        self._tier = "BLOOM"
                    else:
                        # Re-initialize LRU if lost
                        self._lru_keys = OrderedDict()
                    continue
                if key in self._lru_keys:
                    # Refresh LRU position
                    self._lru_keys.move_to_end(key)
                    return
                if len(self._lru_keys) < self.max_lru_keys:
                    self._add_to_lru(key)
                    return
                else:
                    # Try transitioning to Bloom if available
                    if HAS_BLOOM and ScalableBloomFilter is not None:
                        self._transition_to_bloom()
                        # Continue to next iteration to add to the new tier
                    else:
                        # Fallback: stay in LRU and evict oldest
                        self._add_to_lru(key)
                        return
            elif self._tier == "BLOOM":
                if self._bloom_filter is None:
                    # Defensive: try to recover by initializing Bloom
                    if HAS_BLOOM and ScalableBloomFilter is not None:
                        self._bloom_filter = ScalableBloomFilter(
                            initial_capacity=self.max_lru_keys,
                            error_rate=self.false_positive_rate,
                        )
                    else:
                        self._tier = "LRU"
                        self._lru_keys = OrderedDict()
                        continue
                self._add_to_bloom(key)
                return
            else:
                # Should not be reached
                break

fsspeckit.datasets.pyarrow.AdaptiveKeyTracker.get_estimated_memory_usage

get_estimated_memory_usage() -> int

Estimate the current memory usage of the tracker in bytes.

Returns:

Type	Description
int	Estimated memory usage in bytes.

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

def get_estimated_memory_usage(self) -> int:
    """
    Estimate the current memory usage of the tracker in bytes.

    Returns:
        Estimated memory usage in bytes.
    """
    # Note: sys.getsizeof is shallow, so we add estimates for container contents
    if self._tier == "EXACT":
        assert self._exact_keys is not None
        # Approx 24 bytes per entry in set + set overhead
        return sys.getsizeof(self._exact_keys) + (len(self._exact_keys) * 64)
    elif self._tier == "LRU":
        assert self._lru_keys is not None
        # Approx 100 bytes per entry in OrderedDict + overhead
        return sys.getsizeof(self._lru_keys) + (len(self._lru_keys) * 128)
    elif self._tier == "BLOOM":
        if self._bloom_filter:
            # Bloom filters are usually bitarrays, we try to get size if possible
            if hasattr(self._bloom_filter, "bitarray"):
                return sys.getsizeof(self._bloom_filter.bitarray)
            # Fallback to a rough estimate based on capacity and error rate
            # m = - (n * ln(p)) / (ln(2)^2)
            if hasattr(self._bloom_filter, "capacity"):
                import math

                n = self._bloom_filter.capacity
                p = self.false_positive_rate
                m = -(n * math.log(p)) / (math.log(2) ** 2)
                return int(m / 8)  # bits to bytes
        return 0
    return 0

fsspeckit.datasets.pyarrow.AdaptiveKeyTracker.get_metrics

get_metrics() -> Dict[str, Any]

Return quality and performance metrics.

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

def get_metrics(self) -> Dict[str, Any]:
    """
    Return quality and performance metrics.
    """
    with self._lock:
        metrics = {
            "tier": self._tier,
            "total_add_calls": self._keys_added_count,
            "unique_keys_estimate": self._unique_keys_seen,
            "transitions": self._transitions,
            "has_bloom_dependency": HAS_BLOOM,
            "estimated_memory_mb": self.get_estimated_memory_usage()
            / (1024 * 1024),
            "peak_estimated_memory_mb": self._peak_estimated_mem / (1024 * 1024),
        }

        if self._tier == "EXACT":
            assert self._exact_keys is not None
            metrics["current_count"] = len(self._exact_keys)
            metrics["accuracy_type"] = "exact"
        elif self._tier == "LRU":
            assert self._lru_keys is not None
            metrics["current_count"] = len(self._lru_keys)
            metrics["accuracy_type"] = "bounded_lru"
        elif self._tier == "BLOOM":
            metrics["accuracy_type"] = "probabilistic"
            metrics["false_positive_rate_target"] = self.false_positive_rate
            if self._bloom_filter and hasattr(self._bloom_filter, "capacity"):
                metrics["bloom_capacity"] = self._bloom_filter.capacity

        return metrics

Functions

fsspeckit.datasets.pyarrow.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.pyarrow.pyarrow_dataset

pyarrow_dataset(
    self: AbstractFileSystem,
    path: str,
    format: str = "parquet",
    schema: Schema | None = None,
    partitioning: str | list[str] | Partitioning = None,
    **kwargs: Any,
) -> Dataset

Create a PyArrow dataset from files in any supported format.

Creates a dataset that provides optimized reading and querying capabilities including: - Schema inference and enforcement - Partition discovery and pruning - Predicate pushdown - Column projection

Parameters:

Name	Type	Description	Default
path	str	Base path to dataset files	required
format	str	File format. Currently supports: - "parquet" (default) - "csv" - "json" (experimental)	'parquet'
schema	Schema | None	Optional schema to enforce. If None, inferred from data.	None
partitioning	str | list[str] | Partitioning	How the dataset is partitioned. Can be: - str: Single partition field - list[str]: Multiple partition fields - pds.Partitioning: Custom partitioning scheme	None
**kwargs	Any	Additional arguments for dataset creation	{}

Returns:

Type	Description
Dataset	pds.Dataset: PyArrow dataset instance

Example

fs = LocalFileSystem()

# Simple Parquet dataset
ds = fs.pyarrow_dataset("data/")
print(ds.schema)

# Partitioned dataset
ds = fs.pyarrow_dataset(
    "events/",
    partitioning=["year", "month"],
)
# Query with partition pruning
table = ds.to_table(filter=(ds.field("year") == 2024))

# CSV with schema
ds = fs.pyarrow_dataset(
    "logs/",
    format="csv",
    schema=pa.schema(
        [
            ("timestamp", pa.timestamp("s")),
            ("level", pa.string()),
            ("message", pa.string()),
        ],
    ),
)

Source code in src/fsspeckit/core/ext/dataset.py

def pyarrow_dataset(
    self: AbstractFileSystem,
    path: str,
    format: str = "parquet",
    schema: pa.Schema | None = None,
    partitioning: str | list[str] | pds.Partitioning = None,
    **kwargs: Any,
) -> pds.Dataset:
    """Create a PyArrow dataset from files in any supported format.

    Creates a dataset that provides optimized reading and querying capabilities
    including:
    - Schema inference and enforcement
    - Partition discovery and pruning
    - Predicate pushdown
    - Column projection

    Args:
        path: Base path to dataset files
        format: File format. Currently supports:
            - "parquet" (default)
            - "csv"
            - "json" (experimental)
        schema: Optional schema to enforce. If None, inferred from data.
        partitioning: How the dataset is partitioned. Can be:
            - str: Single partition field
            - list[str]: Multiple partition fields
            - pds.Partitioning: Custom partitioning scheme
        **kwargs: Additional arguments for dataset creation

    Returns:
        pds.Dataset: PyArrow dataset instance

    Example:
        ```python
        fs = LocalFileSystem()

        # Simple Parquet dataset
        ds = fs.pyarrow_dataset("data/")
        print(ds.schema)

        # Partitioned dataset
        ds = fs.pyarrow_dataset(
            "events/",
            partitioning=["year", "month"],
        )
        # Query with partition pruning
        table = ds.to_table(filter=(ds.field("year") == 2024))

        # CSV with schema
        ds = fs.pyarrow_dataset(
            "logs/",
            format="csv",
            schema=pa.schema(
                [
                    ("timestamp", pa.timestamp("s")),
                    ("level", pa.string()),
                    ("message", pa.string()),
                ],
            ),
        )
        ```
    """
    return pds.dataset(
        path,
        filesystem=self,
        partitioning=partitioning,
        schema=schema,
        format=format,
        **kwargs,
    )

fsspeckit.datasets.pyarrow.pyarrow_parquet_dataset

pyarrow_parquet_dataset(
    self: AbstractFileSystem,
    path: str,
    schema: Schema | None = None,
    partitioning: str | list[str] | Partitioning = None,
    **kwargs: Any,
) -> Dataset

Create a PyArrow dataset optimized for Parquet files.

Creates a dataset specifically for Parquet data, automatically handling _metadata files for optimized reading.

This function is particularly useful for: - Datasets with existing _metadata files - Multi-file datasets that should be treated as one - Partitioned Parquet datasets

Parameters:

Name	Type	Description	Default
path	str	Path to dataset directory or _metadata file	required
schema	Schema | None	Optional schema to enforce. If None, inferred from data.	None
partitioning	str | list[str] | Partitioning	How the dataset is partitioned. Can be: - str: Single partition field - list[str]: Multiple partition fields - pds.Partitioning: Custom partitioning scheme	None
**kwargs	Any	Additional dataset arguments	{}

Returns:

Type	Description
Dataset	pds.Dataset: PyArrow dataset instance

Example

fs = LocalFileSystem()

# Dataset with _metadata
ds = fs.pyarrow_parquet_dataset("data/_metadata")
print(ds.files)  # Shows all data files

# Partitioned dataset directory
ds = fs.pyarrow_parquet_dataset(
    "sales/",
    partitioning=["year", "region"],
)
# Query with partition pruning
table = ds.to_table(
    filter=(
        (ds.field("year") == 2024)
        & (ds.field("region") == "EMEA")
    ),
)

Source code in src/fsspeckit/core/ext/dataset.py

def pyarrow_parquet_dataset(
    self: AbstractFileSystem,
    path: str,
    schema: pa.Schema | None = None,
    partitioning: str | list[str] | pds.Partitioning = None,
    **kwargs: Any,
) -> pds.Dataset:
    """Create a PyArrow dataset optimized for Parquet files.

    Creates a dataset specifically for Parquet data, automatically handling
    _metadata files for optimized reading.

    This function is particularly useful for:
    - Datasets with existing _metadata files
    - Multi-file datasets that should be treated as one
    - Partitioned Parquet datasets

    Args:
        path: Path to dataset directory or _metadata file
        schema: Optional schema to enforce. If None, inferred from data.
        partitioning: How the dataset is partitioned. Can be:
            - str: Single partition field
            - list[str]: Multiple partition fields
            - pds.Partitioning: Custom partitioning scheme
        **kwargs: Additional dataset arguments

    Returns:
        pds.Dataset: PyArrow dataset instance

    Example:
        ```python
        fs = LocalFileSystem()

        # Dataset with _metadata
        ds = fs.pyarrow_parquet_dataset("data/_metadata")
        print(ds.files)  # Shows all data files

        # Partitioned dataset directory
        ds = fs.pyarrow_parquet_dataset(
            "sales/",
            partitioning=["year", "region"],
        )
        # Query with partition pruning
        table = ds.to_table(
            filter=(
                (ds.field("year") == 2024)
                & (ds.field("region") == "EMEA")
            ),
        )
        ```
    """
    if not self.isfile(path):
        path = posixpath.join(path, "_metadata")
    return pds.parquet_dataset(
        path,
        filesystem=self,
        partitioning=partitioning,
        schema=schema,
        **kwargs,
    )