fixes

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-python/python/vortex/file.py

@@ -61,9 +61,7 @@ def open(
     See also: :class:`vortex.dataset.VortexDataset`
     """
 
-    return VortexFile(
-        _file.open(path, store=store, without_segment_cache=without_segment_cache, session=session)
-    )
+    return VortexFile(_file.open(path, store=store, without_segment_cache=without_segment_cache, session=session))
 
 
 @final

vortex-python/python/vortex/ray/datasource.py

@@ -18,6 +18,7 @@
 from .. import open as vx_open
 from ..arrow.expression import ensure_vortex_expression
 from ..expr import Expr as VortexExpr
+from ..session import Session
 from ..type_aliases import IntoProjection
 
 if TYPE_CHECKING:
@@ -51,12 +52,14 @@ def __init__(
         self,
         *,
         url: str,
+        session: Session,
         columns: IntoProjection = None,
         filter: pc.Expression | VortexExpr | None = None,
         batch_size: int | None = None,
         meta_provider: BaseFileMetadataProvider = DefaultFileMetadataProvider(),  # pyright: ignore[reportCallInDefaultInitializer]
     ):
         super().__init__()
+        self._session = session
         self._columns = columns
         self._filter = filter
 
@@ -101,6 +104,7 @@ def get_read_tasks(
                 self._columns,
                 self._filter,
                 per_task_row_limit if per_task_row_limit is not None else self._batch_size,
+                self._session,
             )
             for paths in partition(parallelism, self._paths)
             if len(paths) > 0
@@ -118,11 +122,12 @@ def _read_task(
     columns: IntoProjection,
     filter: pc.Expression | VortexExpr | None,
     batch_size: int | None,
+    session: Session,
 ) -> ReadTask:
     if not paths:
         raise ValueError("no paths specified")
 
-    files = [vx_open(path) for path in paths]
+    files = [vx_open(path, session=session) for path in paths]
     schemas = [f.dtype.to_arrow_schema() for f in files]
     schema = schemas[0]
     assert all(s == schema for s in schemas[1:])
@@ -140,8 +145,11 @@ def read() -> Iterable[pandas.DataFrame]:
         # If we could serialize a PyVortexFile and a PyExpr, we could set those up earlier.
 
         vx_filter = ensure_vortex_expression(filter, schema=schema)
+        # Ray read functions may execute in worker processes, so use a worker-local session
+        # instead of closing over the driver session.
+        session = Session()
         for path in paths:
-            f = vx_open(path)
+            f = vx_open(path, session=session)
             for rb in f.to_arrow(columns, expr=vx_filter, batch_size=batch_size):
                 # We would prefer to generate Arrow, but we run into this issue: https://github.com/apache/arrow/issues/47279
                 #

vortex-python/test/conftest.py

@@ -3,4 +3,13 @@
 
 import logging
 
+import pytest
+
+import vortex as vx
+
 logging.basicConfig(level=logging.DEBUG)
+
+
+@pytest.fixture
+def session() -> vx.Session:
+    return vx.Session()

vortex-python/test/test_array.py

@@ -7,39 +7,39 @@
 import vortex
 
 
-def test_primitive_array_round_trip():
+def test_primitive_array_round_trip(session: vortex.Session):
     a = pa.array([0, 1, 2, 3])
     arr = vortex.array(a)
-    assert arr.to_arrow_array() == a
+    assert arr.to_arrow_array(session=session) == a
 
 
-def test_array_with_nulls():
+def test_array_with_nulls(session: vortex.Session):
     a = pa.array([b"123", None], type=pa.string_view())
     arr = vortex.array(a)
-    assert arr.to_arrow_array() == a
+    assert arr.to_arrow_array(session=session) == a
 
 
-def test_varbin_array_round_trip():
+def test_varbin_array_round_trip(session: vortex.Session):
     a = pa.array(["a", "b", "c"], type=pa.string_view())
     arr = vortex.array(a)
-    assert arr.to_arrow_array() == a
+    assert arr.to_arrow_array(session=session) == a
 
 
-def test_varbin_array_take():
+def test_varbin_array_take(session: vortex.Session):
     a = vortex.array(pa.array(["a", "b", "c", "d"], type=pa.string_view()))
-    assert a.take(vortex.array(pa.array([0, 2]))).to_arrow_array() == pa.array(
+    assert a.take(vortex.array(pa.array([0, 2]))).to_arrow_array(session=session) == pa.array(
         ["a", "c"],
         type=pa.string_view(),
     )
 
 
-def test_empty_array():
+def test_empty_array(session: vortex.Session):
     a = pa.array([], type=pa.uint8())
     primitive = vortex.array(a)
-    assert primitive.to_arrow_array().type == pa.uint8()
+    assert primitive.to_arrow_array(session=session).type == pa.uint8()
 
 
 @pytest.mark.xfail(raises=IndexError)
-def test_scalar_at_out_of_bounds():
+def test_scalar_at_out_of_bounds(session: vortex.Session):
     a = vortex.array([10, 42, 999, 1992])
-    _s = a.scalar_at(10)
+    _s = a.scalar_at(10, session=session)

vortex-python/test/test_compress.py

@@ -12,42 +12,42 @@
 import vortex
 
 
-def test_primitive_compress():
+def test_primitive_compress(session: vortex.Session):
     a = pa.array([0, 0, 0, 0, 9, 9, 9, 9, 1, 5])
-    arr_compressed = vortex.compress(vortex.array(a))
+    arr_compressed = vortex.compress(vortex.array(a), session=session)
     assert not isinstance(arr_compressed, vortex.PrimitiveArray)
     assert arr_compressed.nbytes < a.nbytes
 
 
-def test_for_compress():
+def test_for_compress(session: vortex.Session):
     a = pa.array(np.arange(10_000) + 10_000_000)
-    arr_compressed = vortex.compress(vortex.array(a))
+    arr_compressed = vortex.compress(vortex.array(a), session=session)
     assert not isinstance(arr_compressed, vortex.PrimitiveArray)
 
 
-def test_arrange_encode():
+def test_arrange_encode(session: vortex.Session):
     a = vortex.array(pa.array(np.arange(10_000), type=pa.uint32()))
-    compressed = vortex.compress(a)
+    compressed = vortex.compress(a, session=session)
     assert isinstance(compressed, vortex.FastLanesDeltaArray | vortex.SequenceArray)
     assert compressed.nbytes < a.nbytes
 
 
-def test_zigzag_encode():
+def test_zigzag_encode(session: vortex.Session):
     a = vortex.array(pa.array([-1, -1, 0, -1, 1, -1]))
-    zarr = vortex.ZigZagArray.encode(a)
+    zarr = vortex.ZigZagArray.encode(a, session=session)
     assert isinstance(zarr, vortex.ZigZagArray)
     # TODO(ngates): support decoding once we have decompressor.
 
 
-def test_chunked_encode():
+def test_chunked_encode(session: vortex.Session):
     chunked = pa.chunked_array([pa.array([0, 1, 2]), pa.array([3, 4, 5])])
     encoded = vortex.array(chunked)
-    arrow = encoded.to_arrow_array()
+    arrow = encoded.to_arrow_array(session=session)
     assert isinstance(arrow, pa.ChunkedArray)
     assert arrow.combine_chunks() == pa.array([0, 1, 2, 3, 4, 5])
 
 
-def test_table_encode():
+def test_table_encode(session: vortex.Session):
     table = pa.table(  # pyright: ignore[reportCallIssue, reportUnknownVariableType]
         {  # pyright: ignore[reportArgumentType]
             "number": pa.chunked_array([pa.array([0, 1, 2]), pa.array([3, 4, 5])]),
@@ -59,7 +59,7 @@ def test_table_encode():
     assert isinstance(table, pa.Table)
 
     encoded = vortex.array(table)
-    arrow = encoded.to_arrow_array()
+    arrow = encoded.to_arrow_array(session=session)
     assert isinstance(arrow, pa.ChunkedArray)
     assert arrow.combine_chunks() == pa.StructArray.from_arrays(  # pyright: ignore[reportUnknownMemberType]
         [pa.array([0, 1, 2, 3, 4, 5]), pa.array(["a", "b", "c", "d", "e", "f"], type=pa.string_view())],
@@ -68,11 +68,11 @@ def test_table_encode():
 
 
 @pytest.mark.skip(reason="We have no way to guarantee that the vortex-bench data has been downloaded.")
-def test_taxi():
+def test_taxi(session: vortex.Session):
     curdir = Path(os.path.dirname(__file__)).parent.parent
     table = pq.read_table(curdir / "vortex-bench/data/yellow-tripdata-2023-11.parquet")  # pyright: ignore[reportUnknownMemberType]
-    compressed = vortex.compress(vortex.array(table[:100]))
-    decompressed = compressed.to_arrow_array()
+    compressed = vortex.compress(vortex.array(table[:100]), session=session)
+    decompressed = compressed.to_arrow_array(session=session)
     assert len(decompressed) == 100
     # hard to test because of string_view
     # assert pc.equal(decompressed, table[:100].to_struct_array()), (decompressed, table[:100].to_struct_array())

vortex-python/test/test_dataset.py

@@ -29,9 +29,10 @@ def ds(tmpdir_factory) -> vx.dataset.VortexDataset:  # pyright: ignore[reportUnk
 
     assert not os.path.exists(fname)  # pyright: ignore[reportUnknownArgumentType]
 
+    session = vx.Session()
     a = pa.array([record(x) for x in range(1_000_000)])
-    vx.io.write(vx.array(a), str(fname))  # pyright: ignore[reportUnknownArgumentType]
-    return vx.dataset.VortexDataset.from_path(str(fname))  # pyright: ignore[reportUnknownArgumentType]
+    vx.io.write(vx.array(a), str(fname), session=session)  # pyright: ignore[reportUnknownArgumentType]
+    return vx.dataset.VortexDataset.from_path(str(fname), session=session)  # pyright: ignore[reportUnknownArgumentType]
 
 
 def test_schema(ds: pd.Dataset):
@@ -129,7 +130,7 @@ def test_filter(ds: vx.dataset.VortexDataset):
     assert len(tbl) == 20
 
 
-def test_filter_with_nested_null_dtype(tmp_path: Path):
+def test_filter_with_nested_null_dtype(tmp_path: Path, session: vx.Session):
     path = tmp_path / "test.vortex"
 
     batch = pa.RecordBatch.from_pylist(
@@ -140,9 +141,9 @@ def test_filter_with_nested_null_dtype(tmp_path: Path):
     )
 
     arr = vx.array(batch.to_struct_array())
-    vx.io.write(vx.ArrayIterator.from_iter(arr.dtype, iter([arr])), str(path))
+    vx.io.write(vx.ArrayIterator.from_iter(arr.dtype, iter([arr]), session=session), str(path), session=session)
 
-    dataset = vx.open(str(path)).to_dataset()
+    dataset = vx.open(str(path), session=session).to_dataset()
     actual = dataset.to_table(filter=pc.field("a") == 0)
 
     assert actual.to_pylist() == [{"a": 0, "b": {"x": None}}]

vortex-python/test/test_datasource.py

@@ -41,16 +41,16 @@ def test_partition():
     assert partition(3, list(range(11))) == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10]]
 
 
-def test_vortex_datasource(ray_init, tmpdir_factory):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType, reportUnusedParameter]
+def test_vortex_datasource(ray_init, tmpdir_factory, session: vx.Session):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType, reportUnusedParameter]
     folder = tmpdir_factory.mktemp("data")  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
 
     arr1 = vx.array([record(x) for x in range(5)])
-    vx.io.write(arr1, str(folder / "01.vortex"))  # pyright: ignore[reportUnknownArgumentType]
+    vx.io.write(arr1, str(folder / "01.vortex"), session=session)  # pyright: ignore[reportUnknownArgumentType]
 
     arr2 = vx.array([record(x) for x in range(5, 10)])
-    vx.io.write(arr2, str(folder / "02.vortex"))  # pyright: ignore[reportUnknownArgumentType]
+    vx.io.write(arr2, str(folder / "02.vortex"), session=session)  # pyright: ignore[reportUnknownArgumentType]
 
-    ds = read_datasource(VortexDatasource(url=str(folder)))  # pyright: ignore[reportUnknownArgumentType]
+    ds = read_datasource(VortexDatasource(url=str(folder), session=session))  # pyright: ignore[reportUnknownArgumentType]
 
     # Without an explicit sort, Ray may reorder rows *even within a single record batch*.
     ds = ds.sort("index")

vortex-python/test/test_duckdb.py

@@ -12,15 +12,15 @@
 import vortex as vx
 
 
-def test_duckdb_via_substrait(tmp_path: Path) -> None:
+def test_duckdb_via_substrait(tmp_path: Path, session: vx.Session) -> None:
     con = duckdb.connect()
 
     arr = pa.array([datetime(2024, 1, 1), datetime(2024, 6, 15), datetime(2024, 12, 31)])
     table = pa.table({"ts": arr})
     path = str(tmp_path / "test_timestamp.vortex")
-    vx.io.write(table, path)
+    vx.io.write(table, path, session=session)
 
-    ds = vx.open(path).to_dataset()  # noqa: F841  # pyright: ignore[reportUnusedVariable] - used by duckdb via SQL
+    ds = vx.open(path, session=session).to_dataset()  # noqa: F841  # pyright: ignore[reportUnusedVariable] - used by duckdb via SQL
     result = con.execute("SELECT * FROM ds WHERE ts > '2024-06-01'").fetchall()
     assert len(result) == 2
     print(result)

vortex-python/test/test_file.py

@@ -23,11 +23,12 @@ def record(x: int, columns: list[str] | set[str] | None = None) -> dict[str, int
 def vxf(tmpdir_factory) -> vx.VortexFile:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
     fname = tmpdir_factory.mktemp("data") / "foo.vortex"  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
 
+    session = vx.Session()
     if not os.path.exists(fname):  # pyright: ignore[reportUnknownArgumentType]
         a = pa.array([record(x) for x in range(1_000_000)])
-        arr = vx.compress(vx.array(a))
-        vx.io.write(arr, str(fname))  # pyright: ignore[reportUnknownArgumentType]
-    return vx.open(str(fname), without_segment_cache=True)  # pyright: ignore[reportUnknownArgumentType]
+        arr = vx.compress(vx.array(a), session=session)
+        vx.io.write(arr, str(fname), session=session)  # pyright: ignore[reportUnknownArgumentType]
+    return vx.open(str(fname), without_segment_cache=True, session=session)  # pyright: ignore[reportUnknownArgumentType]
 
 
 def test_dtype(vxf: VortexFile):
@@ -62,7 +63,7 @@ def test_to_arrow_columns(vxf: VortexFile):
     assert rbr.schema == pa.schema([("string", pa.string_view()), ("bool", pa.bool_())])
 
 
-def test_empty_file(tmpdir_factory):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
+def test_empty_file(tmpdir_factory, session: vx.Session):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
     # test for writing empty files with null columns
     # create an empty table with schema `empty: null`
     table = pa.Table.from_pydict({"empty": []})
@@ -75,10 +76,10 @@ def test_empty_file(tmpdir_factory):  # pyright: ignore[reportUnknownParameterTy
 
     # writing file should succeed
     empty_file = tmpdir_factory.mktemp("data") / "empty.vortex"  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-    vx.io.write(empty, str(empty_file))  # pyright: ignore[reportUnknownArgumentType]
+    vx.io.write(empty, str(empty_file), session=session)  # pyright: ignore[reportUnknownArgumentType]
 
 
-def test_stream_pyarrow(tmpdir_factory):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
+def test_stream_pyarrow(tmpdir_factory, session: vx.Session):  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
     import pyarrow.parquet as pq
 
     data_dir = tmpdir_factory.mktemp("data")  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
@@ -91,4 +92,4 @@ def test_stream_pyarrow(tmpdir_factory):  # pyright: ignore[reportUnknownParamet
     pq.write_table(table, str(data_dir / "names.parquet"))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
 
     df = pq.read_table(str(data_dir / "names.parquet"))  # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType]
-    vx.io.write(df, str(data_dir / "names.vortex"))  # pyright: ignore[reportUnknownArgumentType]
+    vx.io.write(df, str(data_dir / "names.vortex"), session=session)  # pyright: ignore[reportUnknownArgumentType]