|
6 | 6 |
|
7 | 7 | from collections import defaultdict
|
8 | 8 | from enum import Enum, auto, unique
|
9 |
| -from typing import Any, DefaultDict, NamedTuple |
| 9 | +from typing import Any, Callable, DefaultDict, NamedTuple |
| 10 | + |
| 11 | +from .options import installed_pandas |
| 12 | +from .options import pyarrow as pa |
| 13 | + |
| 14 | +if installed_pandas: |
| 15 | + DataType = pa.DataType |
| 16 | +else: |
| 17 | + DataType = None |
| 18 | + |
10 | 19 |
|
11 | 20 | DBAPI_TYPE_STRING = 0
|
12 | 21 | DBAPI_TYPE_BINARY = 1
|
|
17 | 26 | class FieldType(NamedTuple):
|
18 | 27 | name: str
|
19 | 28 | dbapi_type: list[int]
|
20 |
| - |
21 |
| - |
22 |
| -FIELD_TYPES: list[FieldType] = [ |
23 |
| - FieldType(name="FIXED", dbapi_type=[DBAPI_TYPE_NUMBER]), |
24 |
| - FieldType(name="REAL", dbapi_type=[DBAPI_TYPE_NUMBER]), |
25 |
| - FieldType(name="TEXT", dbapi_type=[DBAPI_TYPE_STRING]), |
26 |
| - FieldType(name="DATE", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
27 |
| - FieldType(name="TIMESTAMP", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
28 |
| - FieldType(name="VARIANT", dbapi_type=[DBAPI_TYPE_BINARY]), |
29 |
| - FieldType(name="TIMESTAMP_LTZ", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
30 |
| - FieldType(name="TIMESTAMP_TZ", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
31 |
| - FieldType(name="TIMESTAMP_NTZ", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
32 |
| - FieldType(name="OBJECT", dbapi_type=[DBAPI_TYPE_BINARY]), |
33 |
| - FieldType(name="ARRAY", dbapi_type=[DBAPI_TYPE_BINARY]), |
34 |
| - FieldType(name="BINARY", dbapi_type=[DBAPI_TYPE_BINARY]), |
35 |
| - FieldType(name="TIME", dbapi_type=[DBAPI_TYPE_TIMESTAMP]), |
36 |
| - FieldType(name="BOOLEAN", dbapi_type=[]), |
37 |
| - FieldType(name="GEOGRAPHY", dbapi_type=[DBAPI_TYPE_STRING]), |
38 |
| -] |
| 29 | + pa_type: Callable[[], DataType] |
| 30 | + |
| 31 | + |
| 32 | +# This type mapping holds column type definitions. |
| 33 | +# Be careful to not change the ordering as the index is what Snowflake |
| 34 | +# gives to as schema |
| 35 | +FIELD_TYPES: tuple[FieldType] = ( |
| 36 | + FieldType(name="FIXED", dbapi_type=[DBAPI_TYPE_NUMBER], pa_type=lambda: pa.int64()), |
| 37 | + FieldType( |
| 38 | + name="REAL", dbapi_type=[DBAPI_TYPE_NUMBER], pa_type=lambda: pa.float64() |
| 39 | + ), |
| 40 | + FieldType(name="TEXT", dbapi_type=[DBAPI_TYPE_STRING], pa_type=lambda: pa.string()), |
| 41 | + FieldType( |
| 42 | + name="DATE", dbapi_type=[DBAPI_TYPE_TIMESTAMP], pa_type=lambda: pa.date64() |
| 43 | + ), |
| 44 | + FieldType( |
| 45 | + name="TIMESTAMP", |
| 46 | + dbapi_type=[DBAPI_TYPE_TIMESTAMP], |
| 47 | + pa_type=lambda: pa.time64("ns"), |
| 48 | + ), |
| 49 | + FieldType( |
| 50 | + name="VARIANT", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda: pa.string() |
| 51 | + ), |
| 52 | + FieldType( |
| 53 | + name="TIMESTAMP_LTZ", |
| 54 | + dbapi_type=[DBAPI_TYPE_TIMESTAMP], |
| 55 | + pa_type=lambda: pa.timestamp("ns"), |
| 56 | + ), |
| 57 | + FieldType( |
| 58 | + name="TIMESTAMP_TZ", |
| 59 | + dbapi_type=[DBAPI_TYPE_TIMESTAMP], |
| 60 | + pa_type=lambda: pa.timestamp("ns"), |
| 61 | + ), |
| 62 | + FieldType( |
| 63 | + name="TIMESTAMP_NTZ", |
| 64 | + dbapi_type=[DBAPI_TYPE_TIMESTAMP], |
| 65 | + pa_type=lambda: pa.timestamp("ns"), |
| 66 | + ), |
| 67 | + FieldType( |
| 68 | + name="OBJECT", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda: pa.string() |
| 69 | + ), |
| 70 | + FieldType( |
| 71 | + name="ARRAY", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda: pa.string() |
| 72 | + ), |
| 73 | + FieldType( |
| 74 | + name="BINARY", dbapi_type=[DBAPI_TYPE_BINARY], pa_type=lambda: pa.binary() |
| 75 | + ), |
| 76 | + FieldType( |
| 77 | + name="TIME", dbapi_type=[DBAPI_TYPE_TIMESTAMP], pa_type=lambda: pa.time64("ns") |
| 78 | + ), |
| 79 | + FieldType(name="BOOLEAN", dbapi_type=[], pa_type=lambda: pa.bool_()), |
| 80 | + FieldType( |
| 81 | + name="GEOGRAPHY", dbapi_type=[DBAPI_TYPE_STRING], pa_type=lambda: pa.string() |
| 82 | + ), |
| 83 | +) |
39 | 84 |
|
40 | 85 | FIELD_NAME_TO_ID: DefaultDict[Any, int] = defaultdict(int)
|
41 | 86 | FIELD_ID_TO_NAME: DefaultDict[int, str] = defaultdict(str)
|
|
0 commit comments