Skip to content

Commit ec5c119

Browse files
Feature: added properties to VortexScanner (#5500)
added property to VortexScanner fixes #5454 ``` import vortex as vx vx.io.write(vx.array([{"col1": "a string"}]), 'foo.vortex') x = vx.open('foo.vortex').to_dataset().scanner().projected_schema ``` Signed-off-by: Siddharth Kumar <[email protected]> Co-authored-by: Connor Tsui <[email protected]>
1 parent fdafb28 commit ec5c119

File tree

2 files changed

+26
-0
lines changed

2 files changed

+26
-0
lines changed

vortex-python/python/vortex/dataset.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -791,6 +791,22 @@ def __init__(
791791
def schema(self):
792792
return self._dataset.schema
793793

794+
@property
795+
@override
796+
def dataset_schema(self) -> pyarrow.Schema:
797+
return self._dataset.schema
798+
799+
@property
800+
@override
801+
def projected_schema(self) -> pyarrow.Schema:
802+
if self._columns:
803+
fields: list[pa.Field[pa.DataType]] = [
804+
self._dataset.schema.field(c) # pyright: ignore[reportUnknownMemberType]
805+
for c in self._columns
806+
]
807+
return pyarrow.schema(fields)
808+
return self._dataset.schema
809+
794810
@override
795811
def count_rows(self):
796812
return self._dataset.count_rows(

vortex-python/test/test_scan.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,3 +64,13 @@ def test_scan_with_cast(vxfile: vx.VortexFile):
6464
[{"index": 1, "string": pa.scalar("1", pa.string_view()), "bool": False, "float": math.sqrt(1)}]
6565
)
6666
assert str(actual.to_arrow_array()) == str(expected)
67+
68+
69+
def test_scanner_property_projected(vxfile: vx.VortexFile):
70+
assert vxfile.to_dataset().scanner(columns=["bool"]).projected_schema == pa.schema([("bool", pa.bool_())])
71+
72+
73+
def test_scanner_property_dataset_schema(vxfile: vx.VortexFile):
74+
assert vxfile.to_dataset().scanner().dataset_schema == pa.schema(
75+
[("bool", pa.bool_()), ("float", pa.float64()), ("index", pa.int64()), ("string", pa.string_view())]
76+
)

0 commit comments

Comments
 (0)