diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 439d400..0bab98e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,6 +10,7 @@ repos: additional_dependencies: - 'daft>=0.7.0' - 'pylance>=6.0.0' + - 'pyarrow-stubs>=20.0.0' - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 @@ -26,6 +27,7 @@ repos: hooks: - id: pretty-format-toml args: [--autofix, --no-sort] + exclude: uv.lock - repo: https://github.com/astral-sh/ruff-pre-commit # ruff version. diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index ac624b7..595bdd5 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -142,13 +142,13 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: f"GroupFragmentMergeUDF expected {len(self.read_columns)} data columns, received {len(data_cols)}." ) - arrays: list[_pa.Array] = [] + arrays: list[_pa.Array[Any]] = [] for col_name, s in zip(self.read_columns, data_cols): pylist = s.to_pylist() if hasattr(s, "to_pylist") else list(s) if col_name == self.right_on: - key_arr: _pa.Array + key_arr: _pa.Array[Any] if self.right_on == "_rowaddr": key_arr = _pa.array(pylist, type=_pa.uint64()) else: diff --git a/daft_lance/lance_scan.py b/daft_lance/lance_scan.py index de3349f..e0ea8fb 100644 --- a/daft_lance/lance_scan.py +++ b/daft_lance/lance_scan.py @@ -143,7 +143,10 @@ def __init__( self._enable_strict_filter_pushdown = get_context().daft_planning_config.enable_strict_filter_pushdown base = self._ds.schema if self._include_fragment_id: - base = pa.schema([*base, pa.field("fragment_id", pa.int64())], metadata=base.metadata) + metadata: dict[bytes | str, bytes | str] | None = ( + {k: v for k, v in base.metadata.items()} if base.metadata else None + ) + base = pa.schema([*base, pa.field("fragment_id", pa.int64())], metadata=metadata) self._schema = convert_lance_schema(base) def name(self) -> str: diff --git a/daft_lance/utils.py b/daft_lance/utils.py index 0488262..0eb32c7 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -106,14 +106,14 @@ def construct_lance_dataset( } effective_kwargs.update(kwargs or {}) try: - ds._lance_open_kwargs = effective_kwargs + ds._lance_open_kwargs = effective_kwargs # type: ignore[attr-defined] except Exception: pass # Preserve the full user-provided defaults (including nearest) for Daft's planning # even if we stripped keys out before calling `lance.dataset`. try: - ds._daft_default_scan_options = original_default_scan_options + ds._daft_default_scan_options = original_default_scan_options # type: ignore[attr-defined] except Exception: pass @@ -174,4 +174,7 @@ def select_required_columns(schema: pa.Schema, required_columns: list[str] | Non if missing: raise KeyError(f"Required columns missing in schema: {missing}") fields = [schema.field(schema.get_field_index(c)) for c in required_columns] - return pa.schema(fields, metadata=schema.metadata) + metadata: dict[bytes | str, bytes | str] | None = ( + {k: v for k, v in schema.metadata.items()} if schema.metadata else None + ) + return pa.schema(fields, metadata=metadata)