kumoai 2.13.0.dev202511211730__cp313-cp313-win_amd64.whl → 2.13.0.dev202512040649__cp313-cp313-win_amd64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kumoai/_version.py +1 -1
- kumoai/connector/utils.py +23 -2
- kumoai/experimental/rfm/__init__.py +20 -45
- kumoai/experimental/rfm/backend/__init__.py +0 -0
- kumoai/experimental/rfm/backend/local/__init__.py +38 -0
- kumoai/experimental/rfm/backend/local/table.py +109 -0
- kumoai/experimental/rfm/backend/snow/__init__.py +35 -0
- kumoai/experimental/rfm/backend/snow/table.py +115 -0
- kumoai/experimental/rfm/backend/sqlite/__init__.py +30 -0
- kumoai/experimental/rfm/backend/sqlite/table.py +101 -0
- kumoai/experimental/rfm/base/__init__.py +10 -0
- kumoai/experimental/rfm/base/column.py +66 -0
- kumoai/experimental/rfm/base/source.py +18 -0
- kumoai/experimental/rfm/{local_table.py → base/table.py} +134 -139
- kumoai/experimental/rfm/{local_graph.py → graph.py} +287 -62
- kumoai/experimental/rfm/infer/__init__.py +6 -0
- kumoai/experimental/rfm/infer/dtype.py +79 -0
- kumoai/experimental/rfm/infer/pkey.py +126 -0
- kumoai/experimental/rfm/infer/time_col.py +62 -0
- kumoai/experimental/rfm/local_graph_sampler.py +42 -1
- kumoai/experimental/rfm/local_graph_store.py +13 -27
- kumoai/experimental/rfm/rfm.py +16 -17
- kumoai/experimental/rfm/sagemaker.py +11 -3
- kumoai/kumolib.cp313-win_amd64.pyd +0 -0
- kumoai/testing/decorators.py +1 -1
- {kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/METADATA +9 -8
- {kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/RECORD +30 -18
- kumoai/experimental/rfm/utils.py +0 -344
- {kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/WHEEL +0 -0
- {kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/licenses/LICENSE +0 -0
- {kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/top_level.txt +0 -0
kumoai/_version.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
__version__ = '2.13.0.
|
|
1
|
+
__version__ = '2.13.0.dev202512040649'
|
kumoai/connector/utils.py
CHANGED
|
@@ -381,8 +381,29 @@ def _handle_duplicate_names(names: List[str]) -> List[str]:
|
|
|
381
381
|
|
|
382
382
|
|
|
383
383
|
def _sanitize_columns(names: List[str]) -> Tuple[List[str], bool]:
|
|
384
|
-
|
|
384
|
+
"""Normalize column names in a CSV or Parquet file.
|
|
385
|
+
|
|
386
|
+
Rules:
|
|
387
|
+
- Replace any non-alphanumeric character with "_"
|
|
388
|
+
- Strip leading/trailing underscores
|
|
389
|
+
- Ensure uniqueness by appending suffixes: _1, _2, ...
|
|
390
|
+
- Auto-name empty columns as auto_named_<n>
|
|
391
|
+
|
|
392
|
+
Returns:
|
|
393
|
+
(new_column_names, changed)
|
|
394
|
+
"""
|
|
395
|
+
_SAN_RE = re.compile(r"[^0-9A-Za-z,\t]")
|
|
396
|
+
# 1) Replace non-alphanumeric sequences with underscore
|
|
385
397
|
new = [_SAN_RE.sub("_", n).strip("_") for n in names]
|
|
398
|
+
|
|
399
|
+
# 2) Auto-name any empty column names to match UI behavior
|
|
400
|
+
unnamed_counter = 0
|
|
401
|
+
for i, n in enumerate(new):
|
|
402
|
+
if not n:
|
|
403
|
+
new[i] = f"auto_named_{unnamed_counter}"
|
|
404
|
+
unnamed_counter += 1
|
|
405
|
+
|
|
406
|
+
# 3) Ensure uniqueness (append suffixes where needed)
|
|
386
407
|
new = _handle_duplicate_names(new)
|
|
387
408
|
return new, new != names
|
|
388
409
|
|
|
@@ -1168,7 +1189,7 @@ def _detect_and_validate_csv(head_bytes: bytes) -> str:
|
|
|
1168
1189
|
- Re-serializes those rows and validates with pandas (small nrows) to catch
|
|
1169
1190
|
malformed inputs.
|
|
1170
1191
|
- Raises ValueError on empty input or if parsing fails with the chosen
|
|
1171
|
-
|
|
1192
|
+
delimiter.
|
|
1172
1193
|
"""
|
|
1173
1194
|
if not head_bytes:
|
|
1174
1195
|
raise ValueError("Could not auto-detect a delimiter: file is empty.")
|
|
@@ -1,54 +1,26 @@
|
|
|
1
|
-
try:
|
|
2
|
-
import kumoai.kumolib # noqa: F401
|
|
3
|
-
except Exception as e:
|
|
4
|
-
import platform
|
|
5
|
-
|
|
6
|
-
_msg = f"""RFM is not supported in your environment.
|
|
7
|
-
|
|
8
|
-
💻 Your Environment:
|
|
9
|
-
Python version: {platform.python_version()}
|
|
10
|
-
Operating system: {platform.system()}
|
|
11
|
-
CPU architecture: {platform.machine()}
|
|
12
|
-
glibc version: {platform.libc_ver()[1]}
|
|
13
|
-
|
|
14
|
-
✅ Supported Environments:
|
|
15
|
-
* Python versions: 3.10, 3.11, 3.12, 3.13
|
|
16
|
-
* Operating systems and CPU architectures:
|
|
17
|
-
* Linux (x86_64)
|
|
18
|
-
* macOS (arm64)
|
|
19
|
-
* Windows (x86_64)
|
|
20
|
-
* glibc versions: >=2.28
|
|
21
|
-
|
|
22
|
-
❌ Unsupported Environments:
|
|
23
|
-
* Python versions: 3.8, 3.9, 3.14
|
|
24
|
-
* Operating systems and CPU architectures:
|
|
25
|
-
* Linux (arm64)
|
|
26
|
-
* macOS (x86_64)
|
|
27
|
-
* Windows (arm64)
|
|
28
|
-
* glibc versions: <2.28
|
|
29
|
-
|
|
30
|
-
Please create a feature request at 'https://github.com/kumo-ai/kumo-rfm'."""
|
|
31
|
-
|
|
32
|
-
raise RuntimeError(_msg) from e
|
|
33
|
-
|
|
34
|
-
from dataclasses import dataclass
|
|
35
|
-
from enum import Enum
|
|
36
1
|
import ipaddress
|
|
37
2
|
import logging
|
|
3
|
+
import os
|
|
38
4
|
import re
|
|
39
5
|
import socket
|
|
40
6
|
import threading
|
|
41
|
-
from
|
|
42
|
-
import
|
|
7
|
+
from dataclasses import dataclass
|
|
8
|
+
from enum import Enum
|
|
9
|
+
from typing import Dict, Optional, Tuple
|
|
43
10
|
from urllib.parse import urlparse
|
|
11
|
+
|
|
44
12
|
import kumoai
|
|
45
13
|
from kumoai.client.client import KumoClient
|
|
46
|
-
|
|
47
|
-
KumoClient_SageMakerProxy_Local)
|
|
48
|
-
from .local_table import LocalTable
|
|
49
|
-
from .local_graph import LocalGraph
|
|
50
|
-
from .rfm import ExplainConfig, Explanation, KumoRFM
|
|
14
|
+
|
|
51
15
|
from .authenticate import authenticate
|
|
16
|
+
from .sagemaker import (
|
|
17
|
+
KumoClient_SageMakerAdapter,
|
|
18
|
+
KumoClient_SageMakerProxy_Local,
|
|
19
|
+
)
|
|
20
|
+
from .base import Table
|
|
21
|
+
from .backend.local import LocalTable
|
|
22
|
+
from .graph import Graph
|
|
23
|
+
from .rfm import ExplainConfig, Explanation, KumoRFM
|
|
52
24
|
|
|
53
25
|
logger = logging.getLogger('kumoai_rfm')
|
|
54
26
|
|
|
@@ -197,12 +169,15 @@ def init(
|
|
|
197
169
|
url)
|
|
198
170
|
|
|
199
171
|
|
|
172
|
+
LocalGraph = Graph # NOTE Backward compatibility - do not use anymore.
|
|
173
|
+
|
|
200
174
|
__all__ = [
|
|
175
|
+
'authenticate',
|
|
176
|
+
'init',
|
|
177
|
+
'Table',
|
|
201
178
|
'LocalTable',
|
|
202
|
-
'
|
|
179
|
+
'Graph',
|
|
203
180
|
'KumoRFM',
|
|
204
181
|
'ExplainConfig',
|
|
205
182
|
'Explanation',
|
|
206
|
-
'authenticate',
|
|
207
|
-
'init',
|
|
208
183
|
]
|
|
File without changes
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
try:
|
|
2
|
+
import kumoai.kumolib # noqa: F401
|
|
3
|
+
except Exception as e:
|
|
4
|
+
import platform
|
|
5
|
+
|
|
6
|
+
_msg = f"""RFM is not supported in your environment.
|
|
7
|
+
|
|
8
|
+
💻 Your Environment:
|
|
9
|
+
Python version: {platform.python_version()}
|
|
10
|
+
Operating system: {platform.system()}
|
|
11
|
+
CPU architecture: {platform.machine()}
|
|
12
|
+
glibc version: {platform.libc_ver()[1]}
|
|
13
|
+
|
|
14
|
+
✅ Supported Environments:
|
|
15
|
+
* Python versions: 3.10, 3.11, 3.12, 3.13
|
|
16
|
+
* Operating systems and CPU architectures:
|
|
17
|
+
* Linux (x86_64)
|
|
18
|
+
* macOS (arm64)
|
|
19
|
+
* Windows (x86_64)
|
|
20
|
+
* glibc versions: >=2.28
|
|
21
|
+
|
|
22
|
+
❌ Unsupported Environments:
|
|
23
|
+
* Python versions: 3.8, 3.9, 3.14
|
|
24
|
+
* Operating systems and CPU architectures:
|
|
25
|
+
* Linux (arm64)
|
|
26
|
+
* macOS (x86_64)
|
|
27
|
+
* Windows (arm64)
|
|
28
|
+
* glibc versions: <2.28
|
|
29
|
+
|
|
30
|
+
Please create a feature request at 'https://github.com/kumo-ai/kumo-rfm'."""
|
|
31
|
+
|
|
32
|
+
raise RuntimeError(_msg) from e
|
|
33
|
+
|
|
34
|
+
from .table import LocalTable
|
|
35
|
+
|
|
36
|
+
__all__ = [
|
|
37
|
+
'LocalTable',
|
|
38
|
+
]
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
import warnings
|
|
2
|
+
from typing import List, Optional
|
|
3
|
+
|
|
4
|
+
import pandas as pd
|
|
5
|
+
|
|
6
|
+
from kumoai.experimental.rfm.base import SourceColumn, SourceForeignKey, Table
|
|
7
|
+
from kumoai.experimental.rfm.infer import infer_dtype
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class LocalTable(Table):
|
|
11
|
+
r"""A table backed by a :class:`pandas.DataFrame`.
|
|
12
|
+
|
|
13
|
+
A :class:`LocalTable` fully specifies the relevant metadata, *i.e.*
|
|
14
|
+
selected columns, column semantic types, primary keys and time columns.
|
|
15
|
+
:class:`LocalTable` is used to create a :class:`Graph`.
|
|
16
|
+
|
|
17
|
+
.. code-block:: python
|
|
18
|
+
|
|
19
|
+
import pandas as pd
|
|
20
|
+
import kumoai.experimental.rfm as rfm
|
|
21
|
+
|
|
22
|
+
# Load data from a CSV file:
|
|
23
|
+
df = pd.read_csv("data.csv")
|
|
24
|
+
|
|
25
|
+
# Create a table from a `pandas.DataFrame` and infer its metadata ...
|
|
26
|
+
table = rfm.LocalTable(df, name="my_table").infer_metadata()
|
|
27
|
+
|
|
28
|
+
# ... or create a table explicitly:
|
|
29
|
+
table = rfm.LocalTable(
|
|
30
|
+
df=df,
|
|
31
|
+
name="my_table",
|
|
32
|
+
primary_key="id",
|
|
33
|
+
time_column="time",
|
|
34
|
+
end_time_column=None,
|
|
35
|
+
)
|
|
36
|
+
|
|
37
|
+
# Verify metadata:
|
|
38
|
+
table.print_metadata()
|
|
39
|
+
|
|
40
|
+
# Change the semantic type of a column:
|
|
41
|
+
table[column].stype = "text"
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
df: The data frame to create this table from.
|
|
45
|
+
name: The name of this table.
|
|
46
|
+
primary_key: The name of the primary key of this table, if it exists.
|
|
47
|
+
time_column: The name of the time column of this table, if it exists.
|
|
48
|
+
end_time_column: The name of the end time column of this table, if it
|
|
49
|
+
exists.
|
|
50
|
+
"""
|
|
51
|
+
def __init__(
|
|
52
|
+
self,
|
|
53
|
+
df: pd.DataFrame,
|
|
54
|
+
name: str,
|
|
55
|
+
primary_key: Optional[str] = None,
|
|
56
|
+
time_column: Optional[str] = None,
|
|
57
|
+
end_time_column: Optional[str] = None,
|
|
58
|
+
) -> None:
|
|
59
|
+
|
|
60
|
+
if df.empty:
|
|
61
|
+
raise ValueError("Data frame is empty")
|
|
62
|
+
if isinstance(df.columns, pd.MultiIndex):
|
|
63
|
+
raise ValueError("Data frame must not have a multi-index")
|
|
64
|
+
if not df.columns.is_unique:
|
|
65
|
+
raise ValueError("Data frame must have unique column names")
|
|
66
|
+
if any(col == '' for col in df.columns):
|
|
67
|
+
raise ValueError("Data frame must have non-empty column names")
|
|
68
|
+
|
|
69
|
+
self._data = df.copy(deep=False)
|
|
70
|
+
|
|
71
|
+
super().__init__(
|
|
72
|
+
name=name,
|
|
73
|
+
columns=list(df.columns),
|
|
74
|
+
primary_key=primary_key,
|
|
75
|
+
time_column=time_column,
|
|
76
|
+
end_time_column=end_time_column,
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
def _get_source_columns(self) -> List[SourceColumn]:
|
|
80
|
+
source_columns: List[SourceColumn] = []
|
|
81
|
+
for column in self._data.columns:
|
|
82
|
+
ser = self._data[column]
|
|
83
|
+
try:
|
|
84
|
+
dtype = infer_dtype(ser)
|
|
85
|
+
except Exception:
|
|
86
|
+
warnings.warn(f"Data type inference for column '{column}' in "
|
|
87
|
+
f"table '{self.name}' failed. Consider changing "
|
|
88
|
+
f"the data type of the column to use it within "
|
|
89
|
+
f"this table.")
|
|
90
|
+
continue
|
|
91
|
+
|
|
92
|
+
source_column = SourceColumn(
|
|
93
|
+
name=column,
|
|
94
|
+
dtype=dtype,
|
|
95
|
+
is_primary_key=False,
|
|
96
|
+
is_unique_key=False,
|
|
97
|
+
)
|
|
98
|
+
source_columns.append(source_column)
|
|
99
|
+
|
|
100
|
+
return source_columns
|
|
101
|
+
|
|
102
|
+
def _get_source_foreign_keys(self) -> List[SourceForeignKey]:
|
|
103
|
+
return []
|
|
104
|
+
|
|
105
|
+
def _get_sample_df(self) -> pd.DataFrame:
|
|
106
|
+
return self._data
|
|
107
|
+
|
|
108
|
+
def _get_num_rows(self) -> Optional[int]:
|
|
109
|
+
return len(self._data)
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
from typing import Any, TypeAlias
|
|
2
|
+
|
|
3
|
+
try:
|
|
4
|
+
import snowflake.connector
|
|
5
|
+
except ImportError:
|
|
6
|
+
raise ImportError("No module named 'snowflake'. Please install Kumo SDK "
|
|
7
|
+
"with the 'snowflake' extension via "
|
|
8
|
+
"`pip install kumoai[snowflake]`.")
|
|
9
|
+
|
|
10
|
+
Connection: TypeAlias = snowflake.connector.SnowflakeConnection
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def connect(**kwargs: Any) -> Connection:
|
|
14
|
+
r"""Opens a connection to a :class:`snowflake` database.
|
|
15
|
+
|
|
16
|
+
If available, will return a connection to the active session.
|
|
17
|
+
|
|
18
|
+
kwargs: Connection arguments, following the :class:`snowflake` protocol.
|
|
19
|
+
"""
|
|
20
|
+
try:
|
|
21
|
+
from snowflake.snowpark.context import get_active_session
|
|
22
|
+
return get_active_session().connection
|
|
23
|
+
except Exception:
|
|
24
|
+
pass
|
|
25
|
+
|
|
26
|
+
return snowflake.connector.connect(**kwargs)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
from .table import SnowTable # noqa: E402
|
|
30
|
+
|
|
31
|
+
__all__ = [
|
|
32
|
+
'connect',
|
|
33
|
+
'Connection',
|
|
34
|
+
'SnowTable',
|
|
35
|
+
]
|
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
import re
|
|
2
|
+
from typing import List, Optional, Sequence
|
|
3
|
+
|
|
4
|
+
import pandas as pd
|
|
5
|
+
from kumoapi.typing import Dtype
|
|
6
|
+
|
|
7
|
+
from kumoai.experimental.rfm.backend.sqlite import Connection
|
|
8
|
+
from kumoai.experimental.rfm.base import SourceColumn, SourceForeignKey, Table
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class SnowTable(Table):
|
|
12
|
+
r"""A table backed by a :class:`sqlite` database.
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
connection: The connection to a :class:`snowflake` database.
|
|
16
|
+
name: The name of this table.
|
|
17
|
+
columns: The selected columns of this table.
|
|
18
|
+
primary_key: The name of the primary key of this table, if it exists.
|
|
19
|
+
time_column: The name of the time column of this table, if it exists.
|
|
20
|
+
end_time_column: The name of the end time column of this table, if it
|
|
21
|
+
exists.
|
|
22
|
+
"""
|
|
23
|
+
def __init__(
|
|
24
|
+
self,
|
|
25
|
+
connection: Connection,
|
|
26
|
+
name: str,
|
|
27
|
+
database: str | None = None,
|
|
28
|
+
schema: str | None = None,
|
|
29
|
+
columns: Optional[Sequence[str]] = None,
|
|
30
|
+
primary_key: Optional[str] = None,
|
|
31
|
+
time_column: Optional[str] = None,
|
|
32
|
+
end_time_column: Optional[str] = None,
|
|
33
|
+
) -> None:
|
|
34
|
+
|
|
35
|
+
if database is not None and schema is None:
|
|
36
|
+
raise ValueError(f"Missing 'schema' for table '{name}' in "
|
|
37
|
+
f"database '{database}'")
|
|
38
|
+
|
|
39
|
+
self._connection = connection
|
|
40
|
+
self._database = database
|
|
41
|
+
self._schema = schema
|
|
42
|
+
|
|
43
|
+
super().__init__(
|
|
44
|
+
name=name,
|
|
45
|
+
columns=columns,
|
|
46
|
+
primary_key=primary_key,
|
|
47
|
+
time_column=time_column,
|
|
48
|
+
end_time_column=end_time_column,
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
@property
|
|
52
|
+
def fqn_name(self) -> str:
|
|
53
|
+
names: List[str] = []
|
|
54
|
+
if self._database is not None:
|
|
55
|
+
assert self._schema is not None
|
|
56
|
+
names.extend([self._database, self._schema])
|
|
57
|
+
elif self._schema is not None:
|
|
58
|
+
names.append(self._schema)
|
|
59
|
+
names.append(self._name)
|
|
60
|
+
return '.'.join(names)
|
|
61
|
+
|
|
62
|
+
def _get_source_columns(self) -> List[SourceColumn]:
|
|
63
|
+
source_columns: List[SourceColumn] = []
|
|
64
|
+
with self._connection.cursor() as cursor:
|
|
65
|
+
try:
|
|
66
|
+
cursor.execute(f"DESCRIBE TABLE {self.fqn_name}")
|
|
67
|
+
except Exception as e:
|
|
68
|
+
raise ValueError(
|
|
69
|
+
f"Table '{self.fqn_name}' does not exist") from e
|
|
70
|
+
|
|
71
|
+
for row in cursor.fetchall():
|
|
72
|
+
column, type, _, _, _, is_pkey, is_unique = row[:7]
|
|
73
|
+
|
|
74
|
+
type = type.strip().upper()
|
|
75
|
+
if type.startswith('NUMBER'):
|
|
76
|
+
dtype = Dtype.int
|
|
77
|
+
elif type.startswith('VARCHAR'):
|
|
78
|
+
dtype = Dtype.string
|
|
79
|
+
elif type == 'FLOAT':
|
|
80
|
+
dtype = Dtype.float
|
|
81
|
+
elif type == 'BOOLEAN':
|
|
82
|
+
dtype = Dtype.bool
|
|
83
|
+
elif re.search('DATE|TIMESTAMP', type):
|
|
84
|
+
dtype = Dtype.date
|
|
85
|
+
else:
|
|
86
|
+
continue
|
|
87
|
+
|
|
88
|
+
source_column = SourceColumn(
|
|
89
|
+
name=column,
|
|
90
|
+
dtype=dtype,
|
|
91
|
+
is_primary_key=is_pkey.strip().upper() == 'Y',
|
|
92
|
+
is_unique_key=is_unique.strip().upper() == 'Y',
|
|
93
|
+
)
|
|
94
|
+
source_columns.append(source_column)
|
|
95
|
+
|
|
96
|
+
return source_columns
|
|
97
|
+
|
|
98
|
+
def _get_source_foreign_keys(self) -> List[SourceForeignKey]:
|
|
99
|
+
source_fkeys: List[SourceForeignKey] = []
|
|
100
|
+
with self._connection.cursor() as cursor:
|
|
101
|
+
cursor.execute(f"SHOW IMPORTED KEYS IN TABLE {self.fqn_name}")
|
|
102
|
+
for row in cursor.fetchall():
|
|
103
|
+
_, _, _, dst_table, pkey, _, _, _, fkey = row[:9]
|
|
104
|
+
source_fkeys.append(SourceForeignKey(fkey, dst_table, pkey))
|
|
105
|
+
return source_fkeys
|
|
106
|
+
|
|
107
|
+
def _get_sample_df(self) -> pd.DataFrame:
|
|
108
|
+
with self._connection.cursor() as cursor:
|
|
109
|
+
columns = ', '.join(self._source_column_dict.keys())
|
|
110
|
+
cursor.execute(f"SELECT {columns} FROM {self.fqn_name} LIMIT 1000")
|
|
111
|
+
table = cursor.fetch_arrow_all()
|
|
112
|
+
return table.to_pandas(types_mapper=pd.ArrowDtype)
|
|
113
|
+
|
|
114
|
+
def _get_num_rows(self) -> Optional[int]:
|
|
115
|
+
return None
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
from pathlib import Path
|
|
2
|
+
from typing import Any, TypeAlias, Union
|
|
3
|
+
|
|
4
|
+
try:
|
|
5
|
+
import adbc_driver_sqlite.dbapi as adbc
|
|
6
|
+
except ImportError:
|
|
7
|
+
raise ImportError("No module named 'adbc_driver_sqlite'. Please install "
|
|
8
|
+
"Kumo SDK with the 'sqlite' extension via "
|
|
9
|
+
"`pip install kumoai[sqlite]`.")
|
|
10
|
+
|
|
11
|
+
Connection: TypeAlias = adbc.AdbcSqliteConnection
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def connect(uri: Union[str, Path, None] = None, **kwargs: Any) -> Connection:
|
|
15
|
+
r"""Opens a connection to a :class:`sqlite` database.
|
|
16
|
+
|
|
17
|
+
uri: The path to the database file to be opened.
|
|
18
|
+
kwargs: Additional connection arguments, following the
|
|
19
|
+
:class:`adbc_driver_sqlite` protocol.
|
|
20
|
+
"""
|
|
21
|
+
return adbc.connect(uri, **kwargs)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
from .table import SQLiteTable # noqa: E402
|
|
25
|
+
|
|
26
|
+
__all__ = [
|
|
27
|
+
'connect',
|
|
28
|
+
'Connection',
|
|
29
|
+
'SQLiteTable',
|
|
30
|
+
]
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
import re
|
|
2
|
+
import warnings
|
|
3
|
+
from typing import List, Optional, Sequence
|
|
4
|
+
|
|
5
|
+
import pandas as pd
|
|
6
|
+
from kumoapi.typing import Dtype
|
|
7
|
+
|
|
8
|
+
from kumoai.experimental.rfm.backend.sqlite import Connection
|
|
9
|
+
from kumoai.experimental.rfm.base import SourceColumn, SourceForeignKey, Table
|
|
10
|
+
from kumoai.experimental.rfm.infer import infer_dtype
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class SQLiteTable(Table):
|
|
14
|
+
r"""A table backed by a :class:`sqlite` database.
|
|
15
|
+
|
|
16
|
+
Args:
|
|
17
|
+
connection: The connection to a :class:`sqlite` database.
|
|
18
|
+
name: The name of this table.
|
|
19
|
+
columns: The selected columns of this table.
|
|
20
|
+
primary_key: The name of the primary key of this table, if it exists.
|
|
21
|
+
time_column: The name of the time column of this table, if it exists.
|
|
22
|
+
end_time_column: The name of the end time column of this table, if it
|
|
23
|
+
exists.
|
|
24
|
+
"""
|
|
25
|
+
def __init__(
|
|
26
|
+
self,
|
|
27
|
+
connection: Connection,
|
|
28
|
+
name: str,
|
|
29
|
+
columns: Optional[Sequence[str]] = None,
|
|
30
|
+
primary_key: Optional[str] = None,
|
|
31
|
+
time_column: Optional[str] = None,
|
|
32
|
+
end_time_column: Optional[str] = None,
|
|
33
|
+
) -> None:
|
|
34
|
+
|
|
35
|
+
self._connection = connection
|
|
36
|
+
|
|
37
|
+
super().__init__(
|
|
38
|
+
name=name,
|
|
39
|
+
columns=columns,
|
|
40
|
+
primary_key=primary_key,
|
|
41
|
+
time_column=time_column,
|
|
42
|
+
end_time_column=end_time_column,
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
def _get_source_columns(self) -> List[SourceColumn]:
|
|
46
|
+
source_columns: List[SourceColumn] = []
|
|
47
|
+
with self._connection.cursor() as cursor:
|
|
48
|
+
cursor.execute(f"PRAGMA table_info({self.name})")
|
|
49
|
+
rows = cursor.fetchall()
|
|
50
|
+
|
|
51
|
+
if len(rows) == 0:
|
|
52
|
+
raise ValueError(f"Table '{self.name}' does not exist")
|
|
53
|
+
|
|
54
|
+
for _, column, type, _, _, is_pkey in rows:
|
|
55
|
+
# Determine column affinity:
|
|
56
|
+
type = type.strip().upper()
|
|
57
|
+
if re.search('INT', type):
|
|
58
|
+
dtype = Dtype.int
|
|
59
|
+
elif re.search('TEXT|CHAR|CLOB', type):
|
|
60
|
+
dtype = Dtype.string
|
|
61
|
+
elif re.search('REAL|FLOA|DOUB', type):
|
|
62
|
+
dtype = Dtype.float
|
|
63
|
+
else: # NUMERIC affinity.
|
|
64
|
+
ser = self._sample_df[column]
|
|
65
|
+
try:
|
|
66
|
+
dtype = infer_dtype(ser)
|
|
67
|
+
except Exception:
|
|
68
|
+
warnings.warn(
|
|
69
|
+
f"Data type inference for column '{column}' in "
|
|
70
|
+
f"table '{self.name}' failed. Consider changing "
|
|
71
|
+
f"the data type of the column to use it within "
|
|
72
|
+
f"this table.")
|
|
73
|
+
continue
|
|
74
|
+
|
|
75
|
+
source_column = SourceColumn(
|
|
76
|
+
name=column,
|
|
77
|
+
dtype=dtype,
|
|
78
|
+
is_primary_key=bool(is_pkey),
|
|
79
|
+
is_unique_key=False,
|
|
80
|
+
)
|
|
81
|
+
source_columns.append(source_column)
|
|
82
|
+
|
|
83
|
+
return source_columns
|
|
84
|
+
|
|
85
|
+
def _get_source_foreign_keys(self) -> List[SourceForeignKey]:
|
|
86
|
+
source_fkeys: List[SourceForeignKey] = []
|
|
87
|
+
with self._connection.cursor() as cursor:
|
|
88
|
+
cursor.execute(f"PRAGMA foreign_key_list({self.name})")
|
|
89
|
+
for _, _, dst_table, fkey, pkey, _, _, _ in cursor.fetchall():
|
|
90
|
+
source_fkeys.append(SourceForeignKey(fkey, dst_table, pkey))
|
|
91
|
+
return source_fkeys
|
|
92
|
+
|
|
93
|
+
def _get_sample_df(self) -> pd.DataFrame:
|
|
94
|
+
with self._connection.cursor() as cursor:
|
|
95
|
+
cursor.execute(f"SELECT * FROM {self.name} "
|
|
96
|
+
f"ORDER BY rowid LIMIT 1000")
|
|
97
|
+
table = cursor.fetch_arrow_table()
|
|
98
|
+
return table.to_pandas(types_mapper=pd.ArrowDtype)
|
|
99
|
+
|
|
100
|
+
def _get_num_rows(self) -> Optional[int]:
|
|
101
|
+
return None
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from typing import Any
|
|
3
|
+
|
|
4
|
+
from kumoapi.typing import Dtype, Stype
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@dataclass(init=False, repr=False, eq=False)
|
|
8
|
+
class Column:
|
|
9
|
+
stype: Stype
|
|
10
|
+
|
|
11
|
+
def __init__(
|
|
12
|
+
self,
|
|
13
|
+
name: str,
|
|
14
|
+
dtype: Dtype,
|
|
15
|
+
stype: Stype,
|
|
16
|
+
is_primary_key: bool = False,
|
|
17
|
+
is_time_column: bool = False,
|
|
18
|
+
is_end_time_column: bool = False,
|
|
19
|
+
) -> None:
|
|
20
|
+
self._name = name
|
|
21
|
+
self._dtype = Dtype(dtype)
|
|
22
|
+
self._is_primary_key = is_primary_key
|
|
23
|
+
self._is_time_column = is_time_column
|
|
24
|
+
self._is_end_time_column = is_end_time_column
|
|
25
|
+
self.stype = Stype(stype)
|
|
26
|
+
|
|
27
|
+
@property
|
|
28
|
+
def name(self) -> str:
|
|
29
|
+
return self._name
|
|
30
|
+
|
|
31
|
+
@property
|
|
32
|
+
def dtype(self) -> Dtype:
|
|
33
|
+
return self._dtype
|
|
34
|
+
|
|
35
|
+
def __setattr__(self, key: str, val: Any) -> None:
|
|
36
|
+
if key == 'stype':
|
|
37
|
+
if isinstance(val, str):
|
|
38
|
+
val = Stype(val)
|
|
39
|
+
assert isinstance(val, Stype)
|
|
40
|
+
if not val.supports_dtype(self.dtype):
|
|
41
|
+
raise ValueError(f"Column '{self.name}' received an "
|
|
42
|
+
f"incompatible semantic type (got "
|
|
43
|
+
f"dtype='{self.dtype}' and stype='{val}')")
|
|
44
|
+
if self._is_primary_key and val != Stype.ID:
|
|
45
|
+
raise ValueError(f"Primary key '{self.name}' must have 'ID' "
|
|
46
|
+
f"semantic type (got '{val}')")
|
|
47
|
+
if self._is_time_column and val != Stype.timestamp:
|
|
48
|
+
raise ValueError(f"Time column '{self.name}' must have "
|
|
49
|
+
f"'timestamp' semantic type (got '{val}')")
|
|
50
|
+
if self._is_end_time_column and val != Stype.timestamp:
|
|
51
|
+
raise ValueError(f"End time column '{self.name}' must have "
|
|
52
|
+
f"'timestamp' semantic type (got '{val}')")
|
|
53
|
+
|
|
54
|
+
super().__setattr__(key, val)
|
|
55
|
+
|
|
56
|
+
def __hash__(self) -> int:
|
|
57
|
+
return hash((self.name, self.stype, self.dtype))
|
|
58
|
+
|
|
59
|
+
def __eq__(self, other: Any) -> bool:
|
|
60
|
+
if not isinstance(other, Column):
|
|
61
|
+
return False
|
|
62
|
+
return hash(self) == hash(other)
|
|
63
|
+
|
|
64
|
+
def __repr__(self) -> str:
|
|
65
|
+
return (f'{self.__class__.__name__}(name={self.name}, '
|
|
66
|
+
f'stype={self.stype}, dtype={self.dtype})')
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
|
|
3
|
+
from kumoapi.typing import Dtype
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass
|
|
7
|
+
class SourceColumn:
|
|
8
|
+
name: str
|
|
9
|
+
dtype: Dtype
|
|
10
|
+
is_primary_key: bool
|
|
11
|
+
is_unique_key: bool
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
@dataclass
|
|
15
|
+
class SourceForeignKey:
|
|
16
|
+
name: str
|
|
17
|
+
dst_table: str
|
|
18
|
+
primary_key: str
|