18
18
if TYPE_CHECKING :
19
19
from collections .abc import Iterable , Mapping , Sequence , Set
20
20
21
- from typing_extensions import Any , Self
21
+ from typing_extensions import Any , CapsuleType , Self
22
22
23
23
from cudf_polars .typing import ColumnOptions , DataFrameHeader , Slice
24
24
@@ -40,6 +40,20 @@ def _create_polars_column_metadata(
40
40
return plc .interop .ColumnMetadata (name = name , children_meta = children_meta )
41
41
42
42
43
+ # This is also defined in pylibcudf.interop
44
+ class _ObjectWithArrowMetadata :
45
+ def __init__ (
46
+ self , obj : plc .Table , metadata : list [plc .interop .ColumnMetadata ]
47
+ ) -> None :
48
+ self .obj = obj
49
+ self .metadata = metadata
50
+
51
+ def __arrow_c_array__ (
52
+ self , requested_schema : None = None
53
+ ) -> tuple [CapsuleType , CapsuleType ]:
54
+ return self .obj ._to_schema (self .metadata ), self .obj ._to_host_array ()
55
+
56
+
43
57
# Pacify the type checker. DataFrame init asserts that all the columns
44
58
# have a string name, so let's narrow the type.
45
59
class NamedColumn (Column ):
@@ -82,8 +96,8 @@ def to_polars(self) -> pl.DataFrame:
82
96
)
83
97
for name , col in zip (name_map , self .columns , strict = True )
84
98
]
85
- table = plc . interop . to_arrow (self .table , metadata = metadata )
86
- df : pl . DataFrame = pl .from_arrow ( table )
99
+ table_with_metadata = _ObjectWithArrowMetadata (self .table , metadata )
100
+ df = pl .DataFrame ( table_with_metadata )
87
101
return df .rename (name_map ).with_columns (
88
102
pl .col (c .name ).set_sorted (descending = c .order == plc .types .Order .DESCENDING )
89
103
if c .is_sorted
0 commit comments