|
18 | 18 |
|
19 | 19 | # pypaimon.api implementation based on Java code & py4j lib |
20 | 20 |
|
21 | | -import duckdb |
22 | 21 | import pandas as pd |
23 | 22 | import pyarrow as pa |
24 | | -import ray |
25 | 23 |
|
26 | | -from duckdb.duckdb import DuckDBPyConnection |
27 | 24 | from pypaimon.py4j.java_gateway import get_gateway |
28 | 25 | from pypaimon.py4j.util import java_utils, constants |
29 | 26 | from pypaimon.api import \ |
30 | 27 | (catalog, table, read_builder, table_scan, split, |
31 | 28 | table_read, write_builder, table_write, commit_message, |
32 | 29 | table_commit, Schema, predicate) |
33 | | -from typing import List, Iterator, Optional, Any |
| 30 | +from typing import List, Iterator, Optional, Any, TYPE_CHECKING |
| 31 | + |
| 32 | +if TYPE_CHECKING: |
| 33 | + import ray |
| 34 | + from duckdb.duckdb import DuckDBPyConnection |
34 | 35 |
|
35 | 36 |
|
36 | 37 | class Catalog(catalog.Catalog): |
@@ -171,12 +172,16 @@ def to_duckdb( |
171 | 172 | self, |
172 | 173 | splits: List[Split], |
173 | 174 | table_name: str, |
174 | | - connection: Optional[DuckDBPyConnection] = None) -> DuckDBPyConnection: |
| 175 | + connection: Optional["DuckDBPyConnection"] = None) -> "DuckDBPyConnection": |
| 176 | + import duckdb |
| 177 | + |
175 | 178 | con = connection or duckdb.connect(database=":memory:") |
176 | 179 | con.register(table_name, self.to_arrow(splits)) |
177 | 180 | return con |
178 | 181 |
|
179 | | - def to_ray(self, splits: List[Split]) -> ray.data.dataset.Dataset: |
| 182 | + def to_ray(self, splits: List[Split]) -> "ray.data.dataset.Dataset": |
| 183 | + import ray |
| 184 | + |
180 | 185 | return ray.data.from_arrow(self.to_arrow(splits)) |
181 | 186 |
|
182 | 187 | def _init(self): |
|
0 commit comments