kumoai 2.12.0.dev202511071730__cp310-cp310-win_amd64.whl → 2.13.0.dev202512021731__cp310-cp310-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 (33) hide show
  1. kumoai/__init__.py +6 -9
  2. kumoai/_version.py +1 -1
  3. kumoai/client/client.py +9 -13
  4. kumoai/client/rfm.py +15 -7
  5. kumoai/connector/utils.py +23 -2
  6. kumoai/experimental/rfm/__init__.py +164 -46
  7. kumoai/experimental/rfm/backend/__init__.py +0 -0
  8. kumoai/experimental/rfm/backend/local/__init__.py +38 -0
  9. kumoai/experimental/rfm/backend/local/table.py +244 -0
  10. kumoai/experimental/rfm/backend/snow/__init__.py +32 -0
  11. kumoai/experimental/rfm/backend/sqlite/__init__.py +30 -0
  12. kumoai/experimental/rfm/backend/sqlite/table.py +124 -0
  13. kumoai/experimental/rfm/base/__init__.py +7 -0
  14. kumoai/experimental/rfm/base/column.py +66 -0
  15. kumoai/experimental/rfm/{local_table.py → base/table.py} +71 -139
  16. kumoai/experimental/rfm/{local_graph.py → graph.py} +144 -57
  17. kumoai/experimental/rfm/infer/__init__.py +2 -0
  18. kumoai/experimental/rfm/infer/stype.py +35 -0
  19. kumoai/experimental/rfm/local_graph_sampler.py +0 -2
  20. kumoai/experimental/rfm/local_graph_store.py +12 -11
  21. kumoai/experimental/rfm/local_pquery_driver.py +2 -2
  22. kumoai/experimental/rfm/rfm.py +83 -28
  23. kumoai/experimental/rfm/sagemaker.py +138 -0
  24. kumoai/experimental/rfm/utils.py +1 -120
  25. kumoai/kumolib.cp310-win_amd64.pyd +0 -0
  26. kumoai/spcs.py +1 -3
  27. kumoai/testing/decorators.py +1 -1
  28. kumoai/utils/progress_logger.py +10 -4
  29. {kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202512021731.dist-info}/METADATA +11 -2
  30. {kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202512021731.dist-info}/RECORD +33 -23
  31. {kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202512021731.dist-info}/WHEEL +0 -0
  32. {kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202512021731.dist-info}/licenses/LICENSE +0 -0
  33. {kumoai-2.12.0.dev202511071730.dist-info → kumoai-2.13.0.dev202512021731.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.dev202511071730'
1
+ __version__ = '2.13.0.dev202512021731'
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()
kumoai/client/rfm.py CHANGED
@@ -1,3 +1,5 @@
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,
@@ -28,25 +30,32 @@ class RFMAPI:
28
30
  Returns:
29
31
  RFMPredictResponse containing the predictions
30
32
  """
31
- # Send binary data to the predict endpoint
32
33
  response = self._client._request(
33
- RFMEndpoints.predict, data=request,
34
- headers={'Content-Type': 'application/x-protobuf'})
34
+ RFMEndpoints.predict,
35
+ data=request,
36
+ headers={'Content-Type': 'application/x-protobuf'},
37
+ )
35
38
  raise_on_error(response)
36
39
  return parse_response(RFMPredictResponse, response)
37
40
 
38
- def explain(self, request: bytes) -> RFMExplanationResponse:
41
+ def explain(
42
+ self,
43
+ request: bytes,
44
+ skip_summary: bool = False,
45
+ ) -> RFMExplanationResponse:
39
46
  """Explain the RFM model on the given context.
40
47
 
41
48
  Args:
42
49
  request: The predict request as serialized protobuf.
50
+ skip_summary: Whether to skip generating a human-readable summary
51
+ of the explanation.
43
52
 
44
53
  Returns:
45
54
  RFMPredictResponse containing the explanations
46
55
  """
47
- # Send binary data to the explain endpoint
56
+ params: dict[str, Any] = {'generate_summary': not skip_summary}
48
57
  response = self._client._request(
49
- RFMEndpoints.explain, data=request,
58
+ RFMEndpoints.explain, data=request, params=params,
50
59
  headers={'Content-Type': 'application/x-protobuf'})
51
60
  raise_on_error(response)
52
61
  return parse_response(RFMExplanationResponse, response)
@@ -60,7 +69,6 @@ class RFMAPI:
60
69
  Returns:
61
70
  RFMEvaluateResponse containing the computed metrics
62
71
  """
63
- # Send binary data to the evaluate endpoint
64
72
  response = self._client._request(
65
73
  RFMEndpoints.evaluate, data=request,
66
74
  headers={'Content-Type': 'application/x-protobuf'})
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,244 @@
1
+ from typing import Any, Dict, List, Optional, Tuple
2
+
3
+ import numpy as np
4
+ import pandas as pd
5
+ import pyarrow as pa
6
+ from kumoapi.typing import Dtype, Stype
7
+ from typing_extensions import Self
8
+
9
+ from kumoai.experimental.rfm import utils
10
+ from kumoai.experimental.rfm.base import Column, Table
11
+ from kumoai.experimental.rfm.infer import infer_stype
12
+
13
+
14
+ class LocalTable(Table):
15
+ r"""A table backed by a :class:`pandas.DataFrame`.
16
+
17
+ A :class:`LocalTable` fully specifies the relevant metadata, *i.e.*
18
+ selected columns, column semantic types, primary keys and time columns.
19
+ :class:`LocalTable` is used to create a :class:`Graph`.
20
+
21
+ .. code-block:: python
22
+
23
+ import pandas as pd
24
+ import kumoai.experimental.rfm as rfm
25
+
26
+ # Load data from a CSV file:
27
+ df = pd.read_csv("data.csv")
28
+
29
+ # Create a table from a `pandas.DataFrame` and infer its metadata ...
30
+ table = rfm.LocalTable(df, name="my_table").infer_metadata()
31
+
32
+ # ... or create a table explicitly:
33
+ table = rfm.LocalTable(
34
+ df=df,
35
+ name="my_table",
36
+ primary_key="id",
37
+ time_column="time",
38
+ end_time_column=None,
39
+ )
40
+
41
+ # Verify metadata:
42
+ table.print_metadata()
43
+
44
+ # Change the semantic type of a column:
45
+ table[column].stype = "text"
46
+
47
+ Args:
48
+ df: The data frame to create this table from.
49
+ name: The name of this table.
50
+ primary_key: The name of the primary key of this table, if it exists.
51
+ time_column: The name of the time column of this table, if it exists.
52
+ end_time_column: The name of the end time column of this table, if it
53
+ exists.
54
+ """
55
+ def __init__(
56
+ self,
57
+ df: pd.DataFrame,
58
+ name: str,
59
+ primary_key: Optional[str] = None,
60
+ time_column: Optional[str] = None,
61
+ end_time_column: Optional[str] = None,
62
+ ) -> None:
63
+
64
+ if df.empty:
65
+ raise ValueError("Data frame must have at least one row")
66
+ if isinstance(df.columns, pd.MultiIndex):
67
+ raise ValueError("Data frame must not have a multi-index")
68
+ if not df.columns.is_unique:
69
+ raise ValueError("Data frame must have unique column names")
70
+ if any(col == '' for col in df.columns):
71
+ raise ValueError("Data frame must have non-empty column names")
72
+
73
+ self._data = df.copy(deep=False)
74
+
75
+ super().__init__(
76
+ name=name,
77
+ columns=list(df.columns),
78
+ primary_key=primary_key,
79
+ time_column=time_column,
80
+ end_time_column=end_time_column,
81
+ )
82
+
83
+ def infer_metadata(self, verbose: bool = True) -> Self:
84
+ r"""Infers metadata, *i.e.*, primary keys and time columns, in the
85
+ table.
86
+
87
+ Args:
88
+ verbose: Whether to print verbose output.
89
+ """
90
+ logs = []
91
+
92
+ # Try to detect primary key if not set:
93
+ if not self.has_primary_key():
94
+
95
+ def is_candidate(column: Column) -> bool:
96
+ if column.stype == Stype.ID:
97
+ return True
98
+ if all(column.stype != Stype.ID for column in self.columns):
99
+ if self.name == column.name:
100
+ return True
101
+ if (self.name.endswith('s')
102
+ and self.name[:-1] == column.name):
103
+ return True
104
+ return False
105
+
106
+ candidates = [
107
+ column.name for column in self.columns if is_candidate(column)
108
+ ]
109
+
110
+ if primary_key := utils.detect_primary_key(
111
+ table_name=self.name,
112
+ df=self._data,
113
+ candidates=candidates,
114
+ ):
115
+ self.primary_key = primary_key
116
+ logs.append(f"primary key '{primary_key}'")
117
+
118
+ # Try to detect time column if not set:
119
+ if not self.has_time_column():
120
+ candidates = [
121
+ column.name for column in self.columns
122
+ if column.stype == Stype.timestamp
123
+ and column.name != self._end_time_column
124
+ ]
125
+ if time_column := utils.detect_time_column(self._data, candidates):
126
+ self.time_column = time_column
127
+ logs.append(f"time column '{time_column}'")
128
+
129
+ if verbose and len(logs) > 0:
130
+ print(f"Detected {' and '.join(logs)} in table '{self.name}'")
131
+
132
+ return self
133
+
134
+ def _has_source_column(self, name: str) -> bool:
135
+ return name in self._data.columns
136
+
137
+ def _get_source_dtype(self, name: str) -> Dtype:
138
+ return to_dtype(self._data[name])
139
+
140
+ def _get_source_stype(self, name: str, dtype: Dtype) -> Stype:
141
+ return infer_stype(self._data[name], name, dtype)
142
+
143
+ def _get_source_foreign_keys(self) -> List[Tuple[str, str, str]]:
144
+ return []
145
+
146
+ def _infer_primary_key(self, candidates: List[str]) -> Optional[str]:
147
+ return utils.detect_primary_key(
148
+ table_name=self.name,
149
+ df=self._data,
150
+ candidates=candidates,
151
+ )
152
+
153
+ def _infer_time_column(self, candidates: List[str]) -> Optional[str]:
154
+ return utils.detect_time_column(df=self._data, candidates=candidates)
155
+
156
+ def _num_rows(self) -> Optional[int]:
157
+ return len(self._data)
158
+
159
+
160
+ # Data Type ###################################################################
161
+
162
+ PANDAS_TO_DTYPE: Dict[Any, Dtype] = {
163
+ np.dtype('bool'): Dtype.bool,
164
+ pd.BooleanDtype(): Dtype.bool,
165
+ pa.bool_(): Dtype.bool,
166
+ np.dtype('byte'): Dtype.int,
167
+ pd.UInt8Dtype(): Dtype.int,
168
+ np.dtype('int16'): Dtype.int,
169
+ pd.Int16Dtype(): Dtype.int,
170
+ np.dtype('int32'): Dtype.int,
171
+ pd.Int32Dtype(): Dtype.int,
172
+ np.dtype('int64'): Dtype.int,
173
+ pd.Int64Dtype(): Dtype.int,
174
+ np.dtype('float32'): Dtype.float,
175
+ pd.Float32Dtype(): Dtype.float,
176
+ np.dtype('float64'): Dtype.float,
177
+ pd.Float64Dtype(): Dtype.float,
178
+ np.dtype('object'): Dtype.string,
179
+ pd.StringDtype(storage='python'): Dtype.string,
180
+ pd.StringDtype(storage='pyarrow'): Dtype.string,
181
+ pa.string(): Dtype.string,
182
+ pa.binary(): Dtype.binary,
183
+ np.dtype('datetime64[ns]'): Dtype.date,
184
+ np.dtype('timedelta64[ns]'): Dtype.timedelta,
185
+ pa.list_(pa.float32()): Dtype.floatlist,
186
+ pa.list_(pa.int64()): Dtype.intlist,
187
+ pa.list_(pa.string()): Dtype.stringlist,
188
+ }
189
+
190
+
191
+ def to_dtype(ser: pd.Series) -> Dtype:
192
+ """Extracts the :class:`Dtype` from a :class:`pandas.Series`.
193
+
194
+ Args:
195
+ ser: A :class:`pandas.Series` to analyze.
196
+
197
+ Returns:
198
+ The data type.
199
+ """
200
+ if pd.api.types.is_datetime64_any_dtype(ser.dtype):
201
+ return Dtype.date
202
+
203
+ if isinstance(ser.dtype, pd.CategoricalDtype):
204
+ return Dtype.string
205
+
206
+ if pd.api.types.is_object_dtype(ser.dtype):
207
+ index = ser.iloc[:1000].first_valid_index()
208
+ if index is not None and pd.api.types.is_list_like(ser[index]):
209
+ pos = ser.index.get_loc(index)
210
+ assert isinstance(pos, int)
211
+ ser = ser.iloc[pos:pos + 1000].dropna()
212
+
213
+ if not ser.map(pd.api.types.is_list_like).all():
214
+ raise ValueError("Data contains a mix of list-like and "
215
+ "non-list-like values")
216
+
217
+ # Remove all empty Python lists without known data type:
218
+ ser = ser[ser.map(lambda x: not isinstance(x, list) or len(x) > 0)]
219
+
220
+ # Infer unique data types in this series:
221
+ dtypes = ser.apply(lambda x: PANDAS_TO_DTYPE.get(
222
+ np.array(x).dtype, Dtype.string)).unique().tolist()
223
+
224
+ invalid_dtypes = set(dtypes) - {
225
+ Dtype.string,
226
+ Dtype.int,
227
+ Dtype.float,
228
+ }
229
+ if len(invalid_dtypes) > 0:
230
+ raise ValueError(f"Data contains unsupported list data types: "
231
+ f"{list(invalid_dtypes)}")
232
+
233
+ if Dtype.string in dtypes:
234
+ return Dtype.stringlist
235
+
236
+ if dtypes == [Dtype.int]:
237
+ return Dtype.intlist
238
+
239
+ return Dtype.floatlist
240
+
241
+ if ser.dtype not in PANDAS_TO_DTYPE:
242
+ raise ValueError(f"Unsupported data type '{ser.dtype}'")
243
+
244
+ return PANDAS_TO_DTYPE[ser.dtype]