kumoai 2.12.0.dev202510261832__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.
Files changed (43) hide show
  1. kumoai/__init__.py +6 -9
  2. kumoai/_version.py +1 -1
  3. kumoai/client/client.py +9 -13
  4. kumoai/client/endpoints.py +1 -0
  5. kumoai/client/rfm.py +35 -7
  6. kumoai/connector/utils.py +23 -2
  7. kumoai/experimental/rfm/__init__.py +164 -46
  8. kumoai/experimental/rfm/backend/__init__.py +0 -0
  9. kumoai/experimental/rfm/backend/local/__init__.py +38 -0
  10. kumoai/experimental/rfm/backend/local/table.py +109 -0
  11. kumoai/experimental/rfm/backend/snow/__init__.py +35 -0
  12. kumoai/experimental/rfm/backend/snow/table.py +115 -0
  13. kumoai/experimental/rfm/backend/sqlite/__init__.py +30 -0
  14. kumoai/experimental/rfm/backend/sqlite/table.py +101 -0
  15. kumoai/experimental/rfm/base/__init__.py +10 -0
  16. kumoai/experimental/rfm/base/column.py +66 -0
  17. kumoai/experimental/rfm/base/source.py +18 -0
  18. kumoai/experimental/rfm/{local_table.py → base/table.py} +134 -139
  19. kumoai/experimental/rfm/{local_graph.py → graph.py} +287 -62
  20. kumoai/experimental/rfm/infer/__init__.py +6 -0
  21. kumoai/experimental/rfm/infer/dtype.py +79 -0
  22. kumoai/experimental/rfm/infer/pkey.py +126 -0
  23. kumoai/experimental/rfm/infer/time_col.py +62 -0
  24. kumoai/experimental/rfm/local_graph_sampler.py +42 -3
  25. kumoai/experimental/rfm/local_graph_store.py +13 -27
  26. kumoai/experimental/rfm/local_pquery_driver.py +221 -26
  27. kumoai/experimental/rfm/pquery/__init__.py +0 -4
  28. kumoai/experimental/rfm/pquery/pandas_executor.py +34 -8
  29. kumoai/experimental/rfm/rfm.py +152 -95
  30. kumoai/experimental/rfm/sagemaker.py +138 -0
  31. kumoai/kumolib.cp313-win_amd64.pyd +0 -0
  32. kumoai/spcs.py +1 -3
  33. kumoai/testing/decorators.py +1 -1
  34. kumoai/trainer/trainer.py +9 -10
  35. kumoai/utils/progress_logger.py +10 -4
  36. {kumoai-2.12.0.dev202510261832.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/METADATA +12 -2
  37. {kumoai-2.12.0.dev202510261832.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/RECORD +40 -29
  38. kumoai/experimental/rfm/pquery/backend.py +0 -136
  39. kumoai/experimental/rfm/pquery/pandas_backend.py +0 -478
  40. kumoai/experimental/rfm/utils.py +0 -344
  41. {kumoai-2.12.0.dev202510261832.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/WHEEL +0 -0
  42. {kumoai-2.12.0.dev202510261832.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/licenses/LICENSE +0 -0
  43. {kumoai-2.12.0.dev202510261832.dist-info → kumoai-2.13.0.dev202512040649.dist-info}/top_level.txt +0 -0
kumoai/__init__.py CHANGED
@@ -184,15 +184,12 @@ def init(
184
184
  snowflake_credentials
185
185
  ) if not api_key and snowflake_credentials else None
186
186
  client = KumoClient(url=url, api_key=api_key, spcs_token=spcs_token)
187
- if client.authenticate():
188
- global_state._url = client._url
189
- global_state._api_key = client._api_key
190
- global_state._snowflake_credentials = snowflake_credentials
191
- global_state._spcs_token = client._spcs_token
192
- global_state._snowpark_session = snowpark_session
193
- else:
194
- raise ValueError("Client authentication failed. Please check if you "
195
- "have a valid API key.")
187
+ client.authenticate()
188
+ global_state._url = client._url
189
+ global_state._api_key = client._api_key
190
+ global_state._snowflake_credentials = snowflake_credentials
191
+ global_state._spcs_token = client._spcs_token
192
+ global_state._snowpark_session = snowpark_session
196
193
 
197
194
  if not api_key and snowflake_credentials:
198
195
  # Refresh token every 10 minutes (expires in 1 hour):
kumoai/_version.py CHANGED
@@ -1 +1 @@
1
- __version__ = '2.12.0.dev202510261832'
1
+ __version__ = '2.13.0.dev202512040649'
kumoai/client/client.py CHANGED
@@ -20,7 +20,6 @@ if TYPE_CHECKING:
20
20
  )
21
21
  from kumoai.client.online import OnlineServingEndpointAPI
22
22
  from kumoai.client.pquery import PQueryAPI
23
- from kumoai.client.rfm import RFMAPI
24
23
  from kumoai.client.source_table import SourceTableAPI
25
24
  from kumoai.client.table import TableAPI
26
25
 
@@ -73,12 +72,15 @@ class KumoClient:
73
72
  self._session.headers.update(
74
73
  {'Authorization': f'Snowflake Token={self._spcs_token}'})
75
74
 
76
- def authenticate(self) -> bool:
77
- r"""Raises an exception if authentication fails. Succeeds if the
78
- client is properly formed.
79
- """
80
- return self._session.get(f"{self._url}/v1/connectors",
81
- verify=self._verify_ssl).ok
75
+ def authenticate(self) -> None:
76
+ """Raises an exception if authentication fails."""
77
+ try:
78
+ self._session.get(self._url + '/v1/connectors',
79
+ verify=self._verify_ssl).raise_for_status()
80
+ except Exception:
81
+ raise ValueError(
82
+ "Client authentication failed. Please check if you "
83
+ "have a valid API key/credentials.")
82
84
 
83
85
  def set_spcs_token(self, spcs_token: str) -> None:
84
86
  r"""Sets the SPCS token for the client and updates the session
@@ -163,12 +165,6 @@ class KumoClient:
163
165
  from kumoai.client.online import OnlineServingEndpointAPI
164
166
  return OnlineServingEndpointAPI(self)
165
167
 
166
- @property
167
- def rfm_api(self) -> 'RFMAPI':
168
- r"""Returns the typed RFM API."""
169
- from kumoai.client.rfm import RFMAPI
170
- return RFMAPI(self)
171
-
172
168
  def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
173
169
  r"""Send a HTTP request to the specified endpoint."""
174
170
  endpoint_str = endpoint.get_path()
@@ -147,3 +147,4 @@ class RFMEndpoints:
147
147
  explain = Endpoint(f"{BASE}/explain", HTTPMethod.POST)
148
148
  evaluate = Endpoint(f"{BASE}/evaluate", HTTPMethod.POST)
149
149
  validate_query = Endpoint(f"{BASE}/validate_query", HTTPMethod.POST)
150
+ parse_query = Endpoint(f"{BASE}/parse_query", HTTPMethod.POST)
kumoai/client/rfm.py CHANGED
@@ -1,7 +1,11 @@
1
+ from typing import Any
2
+
1
3
  from kumoapi.json_serde import to_json_dict
2
4
  from kumoapi.rfm import (
3
5
  RFMEvaluateResponse,
4
6
  RFMExplanationResponse,
7
+ RFMParseQueryRequest,
8
+ RFMParseQueryResponse,
5
9
  RFMPredictResponse,
6
10
  RFMValidateQueryRequest,
7
11
  RFMValidateQueryResponse,
@@ -26,25 +30,32 @@ class RFMAPI:
26
30
  Returns:
27
31
  RFMPredictResponse containing the predictions
28
32
  """
29
- # Send binary data to the predict endpoint
30
33
  response = self._client._request(
31
- RFMEndpoints.predict, data=request,
32
- headers={'Content-Type': 'application/x-protobuf'})
34
+ RFMEndpoints.predict,
35
+ data=request,
36
+ headers={'Content-Type': 'application/x-protobuf'},
37
+ )
33
38
  raise_on_error(response)
34
39
  return parse_response(RFMPredictResponse, response)
35
40
 
36
- def explain(self, request: bytes) -> RFMExplanationResponse:
41
+ def explain(
42
+ self,
43
+ request: bytes,
44
+ skip_summary: bool = False,
45
+ ) -> RFMExplanationResponse:
37
46
  """Explain the RFM model on the given context.
38
47
 
39
48
  Args:
40
49
  request: The predict request as serialized protobuf.
50
+ skip_summary: Whether to skip generating a human-readable summary
51
+ of the explanation.
41
52
 
42
53
  Returns:
43
54
  RFMPredictResponse containing the explanations
44
55
  """
45
- # Send binary data to the explain endpoint
56
+ params: dict[str, Any] = {'generate_summary': not skip_summary}
46
57
  response = self._client._request(
47
- RFMEndpoints.explain, data=request,
58
+ RFMEndpoints.explain, data=request, params=params,
48
59
  headers={'Content-Type': 'application/x-protobuf'})
49
60
  raise_on_error(response)
50
61
  return parse_response(RFMExplanationResponse, response)
@@ -58,7 +69,6 @@ class RFMAPI:
58
69
  Returns:
59
70
  RFMEvaluateResponse containing the computed metrics
60
71
  """
61
- # Send binary data to the evaluate endpoint
62
72
  response = self._client._request(
63
73
  RFMEndpoints.evaluate, data=request,
64
74
  headers={'Content-Type': 'application/x-protobuf'})
@@ -82,3 +92,21 @@ class RFMAPI:
82
92
  json=to_json_dict(request))
83
93
  raise_on_error(response)
84
94
  return parse_response(RFMValidateQueryResponse, response)
95
+
96
+ def parse_query(
97
+ self,
98
+ request: RFMParseQueryRequest,
99
+ ) -> RFMParseQueryResponse:
100
+ """Validate a predictive query against a graph.
101
+
102
+ Args:
103
+ request: The request object containing
104
+ the query and graph definition
105
+
106
+ Returns:
107
+ RFMParseQueryResponse containing the QueryDefinition
108
+ """
109
+ response = self._client._request(RFMEndpoints.parse_query,
110
+ json=to_json_dict(request))
111
+ raise_on_error(response)
112
+ return parse_response(RFMParseQueryResponse, response)
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
- _SAN_RE = re.compile(r"[^0-9A-Za-z]+")
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
- delimiter.
1192
+ delimiter.
1172
1193
  """
1173
1194
  if not head_bytes:
1174
1195
  raise ValueError("Could not auto-detect a delimiter: file is empty.")
@@ -1,43 +1,123 @@
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 typing import Optional, Dict
1
+ import ipaddress
2
+ import logging
35
3
  import os
4
+ import re
5
+ import socket
6
+ import threading
7
+ from dataclasses import dataclass
8
+ from enum import Enum
9
+ from typing import Dict, Optional, Tuple
10
+ from urllib.parse import urlparse
11
+
36
12
  import kumoai
37
- from .local_table import LocalTable
38
- from .local_graph import LocalGraph
39
- from .rfm import KumoRFM
13
+ from kumoai.client.client import KumoClient
14
+
40
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
24
+
25
+ logger = logging.getLogger('kumoai_rfm')
26
+
27
+
28
+ def _is_local_address(host: str | None) -> bool:
29
+ """Return True if the hostname/IP refers to the local machine."""
30
+ if not host:
31
+ return False
32
+ try:
33
+ infos = socket.getaddrinfo(host, None)
34
+ for _, _, _, _, sockaddr in infos:
35
+ ip = sockaddr[0]
36
+ ip_obj = ipaddress.ip_address(ip)
37
+ if ip_obj.is_loopback or ip_obj.is_unspecified:
38
+ return True
39
+ return False
40
+ except Exception:
41
+ return False
42
+
43
+
44
+ class InferenceBackend(str, Enum):
45
+ REST = "REST"
46
+ LOCAL_SAGEMAKER = "LOCAL_SAGEMAKER"
47
+ AWS_SAGEMAKER = "AWS_SAGEMAKER"
48
+ UNKNOWN = "UNKNOWN"
49
+
50
+
51
+ def _detect_backend(
52
+ url: str) -> Tuple[InferenceBackend, Optional[str], Optional[str]]:
53
+ parsed = urlparse(url)
54
+
55
+ # Remote SageMaker
56
+ if ("runtime.sagemaker" in parsed.netloc
57
+ and parsed.path.endswith("/invocations")):
58
+ # Example: https://runtime.sagemaker.us-west-2.amazonaws.com/
59
+ # endpoints/Name/invocations
60
+ match = re.search(r"runtime\.sagemaker\.([a-z0-9-]+)\.amazonaws\.com",
61
+ parsed.netloc)
62
+ region = match.group(1) if match else None
63
+ m = re.search(r"/endpoints/([^/]+)/invocations", parsed.path)
64
+ endpoint_name = m.group(1) if m else None
65
+ return InferenceBackend.AWS_SAGEMAKER, region, endpoint_name
66
+
67
+ # Local SageMaker
68
+ if parsed.port == 8080 and parsed.path.endswith(
69
+ "/invocations") and _is_local_address(parsed.hostname):
70
+ return InferenceBackend.LOCAL_SAGEMAKER, None, None
71
+
72
+ # Default: regular REST
73
+ return InferenceBackend.REST, None, None
74
+
75
+
76
+ @dataclass
77
+ class RfmGlobalState:
78
+ _url: str = '__url_not_provided__'
79
+ _backend: InferenceBackend = InferenceBackend.UNKNOWN
80
+ _region: Optional[str] = None
81
+ _endpoint_name: Optional[str] = None
82
+ _thread_local = threading.local()
83
+
84
+ # Thread-safe init-once.
85
+ _initialized: bool = False
86
+ _lock: threading.Lock = threading.Lock()
87
+
88
+ @property
89
+ def client(self) -> KumoClient:
90
+ if self._backend == InferenceBackend.REST:
91
+ return kumoai.global_state.client
92
+
93
+ if hasattr(self._thread_local, '_sagemaker'):
94
+ # Set the spcs token in the client to ensure it has the latest.
95
+ return self._thread_local._sagemaker
96
+
97
+ sagemaker_client: KumoClient
98
+ if self._backend == InferenceBackend.LOCAL_SAGEMAKER:
99
+ sagemaker_client = KumoClient_SageMakerProxy_Local(self._url)
100
+ else:
101
+ assert self._backend == InferenceBackend.AWS_SAGEMAKER
102
+ assert self._region
103
+ assert self._endpoint_name
104
+ sagemaker_client = KumoClient_SageMakerAdapter(
105
+ self._region, self._endpoint_name)
106
+
107
+ self._thread_local._sagemaker = sagemaker_client
108
+ return sagemaker_client
109
+
110
+ def reset(self) -> None: # For testing only.
111
+ with self._lock:
112
+ self._initialized = False
113
+ self._url = '__url_not_provided__'
114
+ self._backend = InferenceBackend.UNKNOWN
115
+ self._region = None
116
+ self._endpoint_name = None
117
+ self._thread_local = threading.local()
118
+
119
+
120
+ global_state = RfmGlobalState()
41
121
 
42
122
 
43
123
  def init(
@@ -47,19 +127,57 @@ def init(
47
127
  snowflake_application: Optional[str] = None,
48
128
  log_level: str = "INFO",
49
129
  ) -> None:
50
- if url is None:
51
- url = os.getenv("KUMO_API_URL", "https://kumorfm.ai/api")
130
+ with global_state._lock:
131
+ if global_state._initialized:
132
+ if url != global_state._url:
133
+ raise ValueError(
134
+ "Kumo RFM has already been initialized with a different "
135
+ "URL. Re-initialization with a different URL is not "
136
+ "supported.")
137
+ return
52
138
 
53
- kumoai.init(url=url, api_key=api_key,
54
- snowflake_credentials=snowflake_credentials,
55
- snowflake_application=snowflake_application,
56
- log_level=log_level)
139
+ if url is None:
140
+ url = os.getenv("RFM_API_URL", "https://kumorfm.ai/api")
57
141
 
142
+ backend, region, endpoint_name = _detect_backend(url)
143
+ if backend == InferenceBackend.REST:
144
+ # Initialize kumoai.global_state
145
+ if (kumoai.global_state.initialized
146
+ and kumoai.global_state._url != url):
147
+ raise ValueError(
148
+ "Kumo AI SDK has already been initialized with different "
149
+ "API URL. Please restart Python interpreter and "
150
+ "initialize via kumoai.rfm.init()")
151
+ kumoai.init(url=url, api_key=api_key,
152
+ snowflake_credentials=snowflake_credentials,
153
+ snowflake_application=snowflake_application,
154
+ log_level=log_level)
155
+ elif backend == InferenceBackend.AWS_SAGEMAKER:
156
+ assert region
157
+ assert endpoint_name
158
+ KumoClient_SageMakerAdapter(region, endpoint_name).authenticate()
159
+ else:
160
+ assert backend == InferenceBackend.LOCAL_SAGEMAKER
161
+ KumoClient_SageMakerProxy_Local(url).authenticate()
162
+
163
+ global_state._url = url
164
+ global_state._backend = backend
165
+ global_state._region = region
166
+ global_state._endpoint_name = endpoint_name
167
+ global_state._initialized = True
168
+ logger.info("Kumo RFM initialized with backend: %s, url: %s", backend,
169
+ url)
170
+
171
+
172
+ LocalGraph = Graph # NOTE Backward compatibility - do not use anymore.
58
173
 
59
174
  __all__ = [
60
- 'LocalTable',
61
- 'LocalGraph',
62
- 'KumoRFM',
63
175
  'authenticate',
64
176
  'init',
177
+ 'Table',
178
+ 'LocalTable',
179
+ 'Graph',
180
+ 'KumoRFM',
181
+ 'ExplainConfig',
182
+ 'Explanation',
65
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
+ ]