datachain 0.2.0__tar.gz → 0.2.1__tar.gz

{datachain-0.2.0 → datachain-0.2.1}/.gitignore

@@ -136,9 +136,6 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
-# setuptools-scm
-/src/datachain/_version.py
-
 .idea/
 .vscode/
 .datachain/

{datachain-0.2.0/src/datachain.egg-info → datachain-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.2.0
+Version: 0.2.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -56,7 +56,7 @@ Requires-Dist: lz4; extra == "remote"
 Requires-Dist: msgpack<2,>=1.0.4; extra == "remote"
 Requires-Dist: requests>=2.22.0; extra == "remote"
 Provides-Extra: vector
-Requires-Dist: scipy; extra == "vector"
+Requires-Dist: usearch; extra == "vector"
 Provides-Extra: tests
 Requires-Dist: datachain[cv,pandas,remote,vector]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"

{datachain-0.2.0 → datachain-0.2.1}/pyproject.toml

@@ -66,7 +66,7 @@ remote = [
   "requests>=2.22.0"
 ]
 vector = [
-  "scipy"
+  "usearch"
 ]
 tests = [
   "datachain[cv,pandas,remote,vector]",
@@ -107,7 +107,6 @@ where = ["src"]
 namespaces = false
 
 [tool.setuptools_scm]
-write_to = "src/datachain/_version.py"
 
 [tool.pytest.ini_options]
 addopts = "-rfEs -m 'not benchmark'"

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/cli.py

@@ -5,13 +5,14 @@ import sys
 import traceback
 from argparse import SUPPRESS, Action, ArgumentParser, ArgumentTypeError, Namespace
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from importlib.metadata import PackageNotFoundError, version
 from itertools import chain
 from multiprocessing import freeze_support
 from typing import TYPE_CHECKING, Optional, Union
 
 import shtab
 
-from datachain import __version__, utils
+from datachain import utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
 from datachain.utils import DataChainDir
 
@@ -96,6 +97,12 @@ def add_show_args(parser: ArgumentParser) -> None:
 
 
 def get_parser() -> ArgumentParser:  # noqa: PLR0915
+    try:
+        __version__ = version("datachain")
+    except PackageNotFoundError:
+        # package is not installed
+        __version__ = "unknown"
+
     parser = ArgumentParser(
         description="DataChain: Wrangle unstructured AI data at scale", prog="datachain"
     )

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/data_storage/schema.py

@@ -31,7 +31,7 @@ def dedup_columns(columns: Iterable[sa.Column]) -> list[sa.Column]:
     """
     c_set: dict[str, sa.Column] = {}
     for c in columns:
-        if ec := c_set.get(c.name, None):
+        if (ec := c_set.get(c.name, None)) is not None:
             if str(ec.type) != str(c.type):
                 raise ValueError(
                     f"conflicting types for column {c.name}:{c.type!s} and {ec.type!s}"
@@ -171,8 +171,8 @@ class DataTable:
     ):
         # copy columns, since re-using the same objects from another table
         # may raise an error
-        columns = [cls.copy_column(c) for c in columns if c.name != "id"]
-        columns = [sa.Column("id", Int, primary_key=True), *columns]
+        columns = cls.sys_columns() + [cls.copy_column(c) for c in columns]
+        columns = dedup_columns(columns)
 
         if metadata is None:
             metadata = sa.MetaData()
@@ -230,11 +230,17 @@ class DataTable:
     def delete(self):
         return self.apply_conditions(self.table.delete())
 
+    @staticmethod
+    def sys_columns():
+        return [
+            sa.Column("id", Int, primary_key=True),
+            sa.Column("random", Int64, nullable=False, default=f.random()),
+        ]
+
     @classmethod
     def file_columns(cls) -> list[sa.Column]:
         return [
-            sa.Column("id", Int, primary_key=True),
-            sa.Column("random", Int64, nullable=False),
+            *cls.sys_columns(),
             sa.Column("vtype", String, nullable=False, index=True),
             sa.Column("dir_type", Int, index=True),
             sa.Column("parent", String, index=True),

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/data_storage/sqlite.py

@@ -33,6 +33,7 @@ from datachain.data_storage.schema import (
 from datachain.dataset import DatasetRecord
 from datachain.error import DataChainError
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
+from datachain.sql.sqlite.base import load_usearch_extension
 from datachain.sql.types import SQLType
 from datachain.storage import StorageURI
 from datachain.utils import DataChainDir
@@ -114,6 +115,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
                 db.set_trace_callback(print)
 
+            load_usearch_extension(db)
+
             return cls(engine, MetaData(), db, db_file)
         except RuntimeError:
             raise DataChainError("Can't connect to SQLite DB") from None

datachain-0.2.1/src/datachain/lib/cached_stream.py

@@ -0,0 +1,38 @@
+from abc import ABC
+from contextlib import AbstractContextManager
+
+from datachain.cache import UniqueId
+
+
+class AbstractCachedStream(AbstractContextManager, ABC):
+    def __init__(self, catalog, uid: UniqueId):
+        self.catalog = catalog
+        self.uid = uid
+        self.mode = "rb"
+
+    def set_mode(self, mode):
+        self.mode = mode
+
+
+class PreCachedStream(AbstractCachedStream):
+    def __init__(self, catalog, uid: UniqueId):
+        super().__init__(catalog, uid)
+        self.client = self.catalog.get_client(self.uid.storage)
+        self.cached_file = None
+
+    def get_path_in_cache(self):
+        return self.catalog.cache.path_from_checksum(self.uid.get_hash())
+
+    def __enter__(self):
+        self.client.download(self.uid)
+        self.cached_file = open(self.get_path_in_cache(), self.mode)
+        return self.cached_file
+
+    def __exit__(self, *args):
+        self.cached_file.close()
+
+
+class PreDownloadStream(PreCachedStream):
+    def __exit__(self, *args):
+        super().__exit__(*args)
+        self.catalog.cache.remove(self.uid)

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/lib/dc.py

@@ -39,6 +39,8 @@ if TYPE_CHECKING:
     import pandas as pd
     from typing_extensions import Self
 
+    from datachain.catalog import Catalog
+
 C = Column
 
 
@@ -200,10 +202,12 @@ class DataChain(DatasetQuery):
     def from_storage(
         cls,
         path,
+        *,
         type: Literal["binary", "text", "image"] = "binary",
+        catalog: Optional["Catalog"] = None,
         recursive: Optional[bool] = True,
         anon: bool = False,
-    ) -> "DataChain":
+    ) -> "Self":
         """Get data from a storage as a list of file with all file attributes. It
         returns the chain itself as usual.
 
@@ -220,7 +224,7 @@ class DataChain(DatasetQuery):
             ```
         """
         func = get_file(type)
-        return DataChain(path, recursive=recursive, anon=anon).map(file=func)
+        return cls(path, catalog=catalog, recursive=recursive, anon=anon).map(file=func)
 
     @classmethod
     def from_dataset(cls, name: str, version: Optional[int] = None) -> "DataChain":

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/lib/feature.py

@@ -7,6 +7,7 @@ from datetime import datetime
 from functools import lru_cache
 from types import GenericAlias
 from typing import (
+    TYPE_CHECKING,
     Any,
     ClassVar,
     Literal,
@@ -39,6 +40,9 @@ from datachain.sql.types import (
     String,
 )
 
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
 FeatureStandardType = Union[
     type[int],
     type[str],
@@ -158,7 +162,7 @@ class Feature(BaseModel):
         s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
         return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
 
-    def _set_stream(self, catalog, stream=None, caching_enabled: bool = False) -> None:
+    def _set_stream(self, catalog: "Catalog", caching_enabled: bool = False) -> None:
         pass
 
     @classmethod

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/lib/feature_registry.py

@@ -1,6 +1,7 @@
+import logging
 from typing import Any, ClassVar, Optional
 
-from datachain.cli import logger
+logger = logging.getLogger(__name__)
 
 
 class Registry:
@@ -16,7 +17,7 @@ class Registry:
         version = fr._version  # type: ignore[attr-defined]
         if version in cls.reg[name]:
             full_name = f"{name}@{version}"
-            logger.warning(f"Feature {full_name} is already registered")
+            logger.warning("Feature %s is already registered", full_name)
         cls.reg[name][version] = fr
 
     @classmethod

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/lib/file.py

@@ -2,11 +2,10 @@ import json
 from abc import ABC, abstractmethod
 from datetime import datetime
 from pathlib import Path
-from typing import Any, ClassVar, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union
 from urllib.parse import unquote, urlparse
 from urllib.request import url2pathname
 
-from fsspec import Callback
 from fsspec.implementations.local import LocalFileSystem
 from pydantic import Field, field_validator
 
@@ -18,6 +17,9 @@ from datachain.lib.utils import DataChainError
 from datachain.sql.types import JSON, Int, String
 from datachain.utils import TIME_ZERO
 
+if TYPE_CHECKING:
+    from datachain.catalog import Catalog
+
 
 class FileFeature(Feature):
     _is_file = True
@@ -182,26 +184,17 @@ class File(FileFeature):
 
     def open(self):
         if self._stream is None:
-            if self._catalog is None:
-                raise FileError(self, "stream is not set")
-            self._stream = self._open_stream()
+            raise FileError(self, "stream is not set")
 
         if self.location:
             return VFileRegistry.resolve(self, self.location)
 
         return self._stream
 
-    def _set_stream(
-        self, catalog=None, stream=None, caching_enabled: bool = False
-    ) -> None:
-        if self._catalog is None and catalog is None:
-            raise DataChainError(f"Cannot set file '{stream}' without catalog")
-
-        if catalog:
-            self._catalog = catalog
-
+    def _set_stream(self, catalog: "Catalog", caching_enabled: bool = False) -> None:
+        self._catalog = catalog
         stream_class = PreCachedStream if caching_enabled else PreDownloadStream
-        self._stream = stream_class(stream, self.size, self._catalog, self.get_uid())
+        self._stream = stream_class(self._catalog, self.get_uid())
         self._caching_enabled = caching_enabled
 
     def get_uid(self) -> UniqueId:
@@ -232,11 +225,6 @@ class File(FileFeature):
     def get_uri(self):
         return f"{self.source}/{self.get_full_name()}"
 
-    def _open_stream(self, cache: bool = False, cb: Optional[Callback] = None):
-        client = self._catalog.get_client(self.source)
-        uid = self.get_uid()
-        return client.open_object(uid, use_cache=cache, cb=cb)
-
     def get_path(self) -> str:
         path = unquote(self.get_uri())
         fs = self.get_fs()
@@ -258,10 +246,8 @@ class TextFile(File):
         super().__init__(**kwargs)
         self._stream = None
 
-    def _set_stream(
-        self, catalog=None, stream=None, caching_enabled: bool = False
-    ) -> None:
-        super()._set_stream(catalog, stream, caching_enabled)
+    def _set_stream(self, catalog: "Catalog", caching_enabled: bool = False) -> None:
+        super()._set_stream(catalog, caching_enabled)
         self._stream.set_mode("r")
 
 

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/lib/udf.py

@@ -6,10 +6,10 @@ from typing import TYPE_CHECKING, Callable, Optional
 from datachain.lib.feature import Feature
 from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.utils import DataChainError, DataChainParamsError
-from datachain.query import Stream, udf
+from datachain.query import udf
 
 if TYPE_CHECKING:
-    from dvxc.query.udf import UDFWrapper
+    from datachain.query.udf import UDFWrapper
 
 
 class UdfError(DataChainParamsError):
@@ -34,11 +34,6 @@ class UDFBase:
 
         params_spec = params.to_udf_spec()
         self.params_spec = list(params_spec.keys())
-        self._contains_stream = False
-        if params.contains_file():
-            self.params_spec.insert(0, Stream())  # type: ignore[arg-type]
-            self._contains_stream = True
-
         self.output_spec = output.to_udf_spec()
 
         self._catalog = None
@@ -122,18 +117,10 @@ class UDFBase:
             rows = [rows]
         objs = []
         for row in rows:
-            if self._contains_stream:
-                stream, *row = row
-            else:
-                stream = None
-
             obj_row = self.params.row_to_objs(row)
-
-            if self._contains_stream:
-                for obj in obj_row:
-                    if isinstance(obj, Feature):
-                        obj._set_stream(self._catalog, stream, True)
-
+            for obj in obj_row:
+                if isinstance(obj, Feature):
+                    obj._set_stream(self._catalog, caching_enabled=True)
             objs.append(obj_row)
         return objs
 
@@ -150,13 +137,7 @@ class UDFBase:
             output_map[name] = []
 
         for flat_obj in group:
-            if self._contains_stream:
-                position = 1
-                stream = flat_obj[0]
-            else:
-                position = 0
-                stream = None
-
+            position = 0
             for signal, (cls, length) in spec_map.items():
                 slice = flat_obj[position : position + length]
                 position += length
@@ -167,7 +148,7 @@ class UDFBase:
                     obj = slice[0]
 
                 if isinstance(obj, Feature):
-                    obj._set_stream(self._catalog, stream)
+                    obj._set_stream(self._catalog)
                 output_map[signal].append(obj)
 
         return list(output_map.values())

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/query/dataset.py

@@ -1737,22 +1737,16 @@ class DatasetQuery:
 
             # Exclude the id column and let the db create it to avoid unique
             # constraint violations.
-            cols = [col.name for col in dr.get_table().c if col.name != "id"]
-            assert cols
             q = query.exclude(("id",))
-
             if q._order_by_clauses:
                 # ensuring we have id sorted by order by clause if it exists in a query
                 q = q.add_columns(
                     f.row_number().over(order_by=q._order_by_clauses).label("id")
                 )
-                cols.append("id")
-
-            self.catalog.warehouse.db.execute(
-                sqlalchemy.insert(dr.get_table()).from_select(cols, q),
-                **kwargs,
-            )
 
+            cols = tuple(c.name for c in q.columns)
+            insert_q = sqlalchemy.insert(dr.get_table()).from_select(cols, q)
+            self.catalog.warehouse.db.execute(insert_q, **kwargs)
             self.catalog.metastore.update_dataset_status(
                 dataset, DatasetStatus.COMPLETE, version=version
             )

{datachain-0.2.0 → datachain-0.2.1}/src/datachain/sql/sqlite/base.py

@@ -71,8 +71,6 @@ def setup():
     compiles(sql_path.name, "sqlite")(compile_path_name)
     compiles(sql_path.file_stem, "sqlite")(compile_path_file_stem)
     compiles(sql_path.file_ext, "sqlite")(compile_path_file_ext)
-    compiles(array.cosine_distance, "sqlite")(compile_cosine_distance)
-    compiles(array.euclidean_distance, "sqlite")(compile_euclidean_distance)
     compiles(array.length, "sqlite")(compile_array_length)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
@@ -81,6 +79,13 @@ def setup():
     compiles(Values, "sqlite")(compile_values)
     compiles(random.rand, "sqlite")(compile_rand)
 
+    if load_usearch_extension(sqlite3.connect(":memory:")):
+        compiles(array.cosine_distance, "sqlite")(compile_cosine_distance_ext)
+        compiles(array.euclidean_distance, "sqlite")(compile_euclidean_distance_ext)
+    else:
+        compiles(array.cosine_distance, "sqlite")(compile_cosine_distance)
+        compiles(array.euclidean_distance, "sqlite")(compile_euclidean_distance)
+
     register_user_defined_sql_functions()
     setup_is_complete = True
 
@@ -246,11 +251,23 @@ def compile_path_file_ext(element, compiler, **kwargs):
     return compiler.process(path_file_ext(*element.clauses.clauses), **kwargs)
 
 
+def compile_cosine_distance_ext(element, compiler, **kwargs):
+    run_compiler_hook("cosine_distance")
+    return f"distance_cosine_f32({compiler.process(element.clauses, **kwargs)})"
+
+
 def compile_cosine_distance(element, compiler, **kwargs):
     run_compiler_hook("cosine_distance")
     return f"cosine_distance({compiler.process(element.clauses, **kwargs)})"
 
 
+def compile_euclidean_distance_ext(element, compiler, **kwargs):
+    run_compiler_hook("euclidean_distance")
+    return (
+        f"sqrt(distance_sqeuclidean_f32({compiler.process(element.clauses, **kwargs)}))"
+    )
+
+
 def compile_euclidean_distance(element, compiler, **kwargs):
     run_compiler_hook("euclidean_distance")
     return f"euclidean_distance({compiler.process(element.clauses, **kwargs)})"
@@ -330,3 +347,18 @@ def compile_values(element, compiler, **kwargs):
 
 def compile_rand(element, compiler, **kwargs):
     return compiler.process(func.random(), **kwargs)
+
+
+def load_usearch_extension(conn) -> bool:
+    try:
+        # usearch is part of the vector optional dependencies
+        # we use the extension's cosine and euclidean distance functions
+        from usearch import sqlite_path
+
+        conn.enable_load_extension(True)
+        conn.load_extension(sqlite_path())
+        conn.enable_load_extension(False)
+        return True
+
+    except Exception:  # noqa: BLE001
+        return False

datachain-0.2.1/src/datachain/sql/sqlite/vector.py

@@ -0,0 +1,23 @@
+import math
+
+import numpy as np
+
+
+def euclidean_distance(a: str, b: str):
+    a_np = np.fromstring(a[1:-1], sep=",")
+    b_np = np.fromstring(b[1:-1], sep=",")
+
+    return np.linalg.norm(b_np - a_np)
+
+
+def cosine_distance(a: str, b: str):
+    u = np.fromstring(a[1:-1], sep=",")
+    v = np.fromstring(b[1:-1], sep=",")
+
+    uv = np.inner(u, v)
+    uu = np.inner(u, u)
+    vv = np.inner(v, v)
+
+    dist = 1.0 - uv / math.sqrt(uu * vv)
+
+    return max(0, min(dist, 2.0))

{datachain-0.2.0 → datachain-0.2.1/src/datachain.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.2.0
+Version: 0.2.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -56,7 +56,7 @@ Requires-Dist: lz4; extra == "remote"
 Requires-Dist: msgpack<2,>=1.0.4; extra == "remote"
 Requires-Dist: requests>=2.22.0; extra == "remote"
 Provides-Extra: vector
-Requires-Dist: scipy; extra == "vector"
+Requires-Dist: usearch; extra == "vector"
 Provides-Extra: tests
 Requires-Dist: datachain[cv,pandas,remote,vector]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"

{datachain-0.2.0 → datachain-0.2.1}/src/datachain.egg-info/SOURCES.txt

@@ -84,7 +84,6 @@ examples/zalando/zalando_splits_and_classes_ds.py
 examples/zalando/zalando_splits_and_classes_output.py
 src/datachain/__init__.py
 src/datachain/__main__.py
-src/datachain/_version.py
 src/datachain/asyn.py
 src/datachain/cache.py
 src/datachain/cli.py
@@ -233,7 +232,6 @@ tests/unit/test_utils.py
 tests/unit/test_warehouse.py
 tests/unit/lib/__init__.py
 tests/unit/lib/test_arrow.py
-tests/unit/lib/test_cached_stream.py
 tests/unit/lib/test_datachain.py
 tests/unit/lib/test_datachain_merge.py
 tests/unit/lib/test_feature.py

{datachain-0.2.0 → datachain-0.2.1}/src/datachain.egg-info/requires.txt

@@ -70,4 +70,4 @@ aiotools>=1.7.0
 requests-mock
 
 [vector]
-scipy
+usearch

datachain-0.2.1/tests/func/test_datachain.py

@@ -0,0 +1,40 @@
+import pytest
+
+from datachain.lib.dc import DataChain
+from datachain.lib.file import File
+
+
+@pytest.mark.parametrize("anon", [True, False])
+def test_catalog_anon(catalog, anon):
+    chain = (
+        DataChain.from_storage("gs://dvcx-datalakes/dogs-and-cats/", anon=anon)
+        .limit(5)
+        .save("test_catalog_anon")
+    )
+    assert chain.catalog.client_config.get("anon", False) is anon
+
+
+def test_from_storage(cloud_test_catalog):
+    ctc = cloud_test_catalog
+    dc = DataChain.from_storage(ctc.src_uri, catalog=ctc.catalog)
+    assert dc.count() == 7
+
+
+def test_map_file(cloud_test_catalog):
+    ctc = cloud_test_catalog
+
+    def new_signal(file: File) -> str:
+        with file.open() as f:
+            return file.name + " -> " + f.read().decode("utf-8")
+
+    dc = DataChain.from_storage(ctc.src_uri, catalog=ctc.catalog).map(signal=new_signal)
+    expected = {
+        "description -> Cats and Dogs",
+        "cat1 -> meow",
+        "cat2 -> mrow",
+        "dog1 -> woof",
+        "dog2 -> arf",
+        "dog3 -> bark",
+        "dog4 -> ruff",
+    }
+    assert set(dc.collect_one("signal")) == expected

{datachain-0.2.0 → datachain-0.2.1}/tests/func/test_dataset_query.py

@@ -1,5 +1,6 @@
 import io
 import json
+import math
 import os
 import pickle
 import random
@@ -56,7 +57,6 @@ from tests.utils import (
     SIMPLE_DS_QUERY_RECORDS,
     TARRED_TREE,
     WEBFORMAT_TREE,
-    adjusted_float_diff,
     assert_row_names,
     create_tar_dataset,
     dataset_dependency_asdict,
@@ -1140,6 +1140,8 @@ def test_udf_distributed_cancel(cloud_test_catalog_tmpfile, capfd, datachain_job
             workers=2,
             team_id=metastore.team_id,
             created_at=datetime.now(timezone.utc),
+            params="{}",
+            metrics="{}",
         ),
     )
 
@@ -2851,13 +2853,12 @@ def test_similarity_search(cloud_test_catalog):
         ("dogs", "dog3", 0.7695916496857775, 1.8344983482620636),
         ("dogs/others", "dog4", 0.9789704524691446, 2.0531542018152322),
     ]
-    expected_diffs = [(p, n, 0.0, 0.0) for p, n, _, _ in expected]
 
-    diffs = [
-        (p2, n2, adjusted_float_diff(c1, c2, 1e-8), adjusted_float_diff(e1, e2, 1e-5))
-        for (p1, n1, c1, e1), (p2, n2, c2, e2) in zip(result, expected)
-    ]
-    assert diffs == expected_diffs
+    for (p1, n1, c1, e1), (p2, n2, c2, e2) in zip(result, expected):
+        assert p1.endswith(p2)
+        assert n1 == n2
+        assert math.isclose(c1, c2, abs_tol=1e-5)
+        assert math.isclose(e1, e2, abs_tol=1e-5)
 
 
 @pytest.mark.parametrize(

{datachain-0.2.0 → datachain-0.2.1}/tests/unit/lib/test_arrow.py

@@ -22,11 +22,10 @@ def test_arrow_generator(tmp_path, catalog):
     pq_path = tmp_path / name
     df.to_parquet(pq_path)
     stream = File(name=name, parent=tmp_path.as_posix(), source="file:///")
-    with open(pq_path) as fd:
-        stream._set_stream(catalog, fd, caching_enabled=False)
+    stream._set_stream(catalog, caching_enabled=False)
 
-        func = ArrowGenerator()
-        objs = list(func(stream))
+    func = ArrowGenerator()
+    objs = list(func(stream))
 
     assert len(objs) == len(ids)
     for index, (o, id, text) in enumerate(zip(objs, ids, texts)):