kumoai 2.13.0.dev202511211730__py3-none-any.whl → 2.14.0.dev202512141732__py3-none-any.whl

kumoai/experimental/rfm/sagemaker.py

@@ -2,15 +2,22 @@ import base64
 import json
 from typing import Any, Dict, List, Tuple
 
-import boto3
 import requests
-from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
-from mypy_boto3_sagemaker_runtime.type_defs import InvokeEndpointOutputTypeDef
 
 from kumoai.client import KumoClient
 from kumoai.client.endpoints import Endpoint, HTTPMethod
 from kumoai.exceptions import HTTPException
 
+try:
+    # isort: off
+    from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
+    from mypy_boto3_sagemaker_runtime.type_defs import (
+        InvokeEndpointOutputTypeDef, )
+    # isort: on
+except ImportError:
+    SageMakerRuntimeClient = Any
+    InvokeEndpointOutputTypeDef = Any
+
 
 class SageMakerResponseAdapter(requests.Response):
     def __init__(self, sm_response: InvokeEndpointOutputTypeDef):
@@ -34,6 +41,7 @@ class SageMakerResponseAdapter(requests.Response):
 
 class KumoClient_SageMakerAdapter(KumoClient):
     def __init__(self, region: str, endpoint_name: str):
+        import boto3
         self._client: SageMakerRuntimeClient = boto3.client(
             service_name="sagemaker-runtime", region_name=region)
         self._endpoint_name = endpoint_name

kumoai/pquery/predictive_query.py

@@ -370,9 +370,11 @@ class PredictiveQuery:
         train_table_job_api = global_state.client.generate_train_table_job_api
         job_id: GenerateTrainTableJobID = train_table_job_api.create(
             GenerateTrainTableRequest(
-                dict(custom_tags), pq_id, plan,
-                graph_snapshot_id=self.graph.snapshot(
-                    non_blocking=non_blocking)))
+                dict(custom_tags),
+                pq_id,
+                plan,
+                None,
+            ))
 
         self._train_table = TrainingTableJob(job_id=job_id)
         if non_blocking:
@@ -451,9 +453,11 @@ class PredictiveQuery:
         bp_table_api = global_state.client.generate_prediction_table_job_api
         job_id: GeneratePredictionTableJobID = bp_table_api.create(
             GeneratePredictionTableRequest(
-                dict(custom_tags), pq_id, plan,
-                graph_snapshot_id=self.graph.snapshot(
-                    non_blocking=non_blocking)))
+                dict(custom_tags),
+                pq_id,
+                plan,
+                None,
+            ))
 
         self._prediction_table = PredictionTableJob(job_id=job_id)
         if non_blocking:

kumoai/testing/decorators.py

@@ -25,7 +25,7 @@ def onlyFullTest(func: Callable) -> Callable:
 def has_package(package: str) -> bool:
     r"""Returns ``True`` in case ``package`` is installed."""
     req = Requirement(package)
-    if importlib.util.find_spec(req.name) is None:
+    if importlib.util.find_spec(req.name) is None:  # type: ignore
         return False
 
     try:

kumoai/testing/snow.py

@@ -0,0 +1,50 @@
+import json
+import os
+
+from kumoai.experimental.rfm.backend.snow import Connection
+from kumoai.experimental.rfm.backend.snow import connect as _connect
+
+
+def connect(
+    region: str,
+    id: str,
+    account: str,
+    user: str,
+    warehouse: str,
+    database: str | None = None,
+    schema: str | None = None,
+) -> Connection:
+
+    kwargs = dict(password=os.getenv('SNOWFLAKE_PASSWORD'))
+    if kwargs['password'] is None:
+        import boto3
+        from cryptography.hazmat.primitives import serialization
+
+        client = boto3.client(
+            service_name='secretsmanager',
+            region_name=region,
+        )
+        secret_id = (f'arn:aws:secretsmanager:{region}:{id}:secret:'
+                     f'{account}.snowflakecomputing.com')
+        response = client.get_secret_value(SecretId=secret_id)['SecretString']
+        secret = json.loads(response)
+
+        private_key = serialization.load_pem_private_key(
+            secret['kumo_user_secretkey'].encode(),
+            password=None,
+        )
+        kwargs['private_key'] = private_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+    return _connect(
+        account=account,
+        user=user,
+        warehouse='WH_XS',
+        database='KUMO',
+        schema=schema,
+        session_parameters=dict(CLIENT_TELEMETRY_ENABLED=False),
+        **kwargs,
+    )

kumoai/utils/__init__.py

@@ -1,8 +1,10 @@
+from .sql import quote_ident
 from .progress_logger import ProgressLogger, InteractiveProgressLogger
 from .forecasting import ForecastVisualizer
 from .datasets import from_relbench
 
 __all__ = [
+    'quote_ident',
     'ProgressLogger',
     'InteractiveProgressLogger',
     'ForecastVisualizer',

kumoai/utils/sql.py

@@ -0,0 +1,3 @@
+def quote_ident(name: str) -> str:
+    r"""Quotes a SQL identifier."""
+    return '"' + name.replace('"', '""') + '"'

{kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.14.0.dev202512141732.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.13.0.dev202511211730
+Version: 2.14.0.dev202512141732
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT
@@ -23,13 +23,11 @@ Requires-Dist: requests>=2.28.2
 Requires-Dist: urllib3
 Requires-Dist: plotly
 Requires-Dist: typing_extensions>=4.5.0
-Requires-Dist: kumo-api==0.46.0
+Requires-Dist: kumo-api==0.49.0
 Requires-Dist: tqdm>=4.66.0
 Requires-Dist: aiohttp>=3.10.0
 Requires-Dist: pydantic>=1.10.21
 Requires-Dist: rich>=9.0.0
-Requires-Dist: mypy-boto3-sagemaker-runtime
-Requires-Dist: boto3
 Provides-Extra: doc
 Requires-Dist: sphinx; extra == "doc"
 Requires-Dist: sphinx-book-theme; extra == "doc"
@@ -40,13 +38,16 @@ Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
 Requires-Dist: requests-mock; extra == "test"
-Provides-Extra: test-sagemaker
-Requires-Dist: sagemaker; extra == "test-sagemaker"
-Requires-Dist: pandas==2.1.4; extra == "test-sagemaker"
-Requires-Dist: pyarrow==12.0.1; extra == "test-sagemaker"
+Provides-Extra: sqlite
+Requires-Dist: adbc_driver_sqlite; extra == "sqlite"
+Provides-Extra: snowflake
+Requires-Dist: snowflake-connector-python; extra == "snowflake"
+Requires-Dist: pyyaml; extra == "snowflake"
 Provides-Extra: sagemaker
 Requires-Dist: boto3<2.0,>=1.30.0; extra == "sagemaker"
 Requires-Dist: mypy-boto3-sagemaker-runtime<2.0,>=1.34.0; extra == "sagemaker"
+Provides-Extra: test-sagemaker
+Requires-Dist: sagemaker<3.0; extra == "test-sagemaker"
 Dynamic: license-file
 Dynamic: requires-dist
 

{kumoai-2.13.0.dev202511211730.dist-info → kumoai-2.14.0.dev202512141732.dist-info}/RECORD

@@ -1,7 +1,7 @@
-kumoai/__init__.py,sha256=L3yOOtpSdwe3PYQlJBLkiQd3Ypp8iB5ChXkzprk3Si4,10546
+kumoai/__init__.py,sha256=Nn9YH_x9kAeEFn8RWbP95slZow0qFnakPZZ1WADe1hY,10843
 kumoai/_logging.py,sha256=U2_5ROdyk92P4xO4H2WJV8EC7dr6YxmmnM-b7QX9M7I,886
 kumoai/_singleton.py,sha256=UTwrbDkoZSGB8ZelorvprPDDv9uZkUi1q_SrmsyngpQ,836
-kumoai/_version.py,sha256=SG9nFVn5zQGjYC2bJHf14ClhCtvj0GzhN9puMBHh-sE,39
+kumoai/_version.py,sha256=jDuUtpQbekehTWtqQXeU4rYieU0OA-7URATbFO1iyvo,39
 kumoai/databricks.py,sha256=e6E4lOFvZHXFwh4CO1kXU1zzDU3AapLQYMxjiHPC-HQ,476
 kumoai/exceptions.py,sha256=b-_sdbAKOg50uaJZ65GmBLdTo4HANdjl8_R0sJpwaN0,833
 kumoai/formatting.py,sha256=jA_rLDCGKZI8WWCha-vtuLenVKTZvli99Tqpurz1H84,953
@@ -19,7 +19,7 @@ kumoai/client/endpoints.py,sha256=iF2ZD25AJCIVbmBJ8tTZ8y1Ch0m6nTp18ydN7h4WiTk,53
 kumoai/client/graph.py,sha256=zvLEDExLT_RVbUMHqVl0m6tO6s2gXmYSoWmPF6YMlnA,3831
 kumoai/client/jobs.py,sha256=iu_Wrta6BQMlV6ZtzSnmhjwNPKDMQDXOsqVVIyWodqw,17074
 kumoai/client/online.py,sha256=pkBBh_DEC3GAnPcNw6bopNRlGe7EUbIFe7_seQqZRaw,2720
-kumoai/client/pquery.py,sha256=R2hc-M8vPoyIDH0ywLwFVxCznVAqpZz3w2HszjdNW-o,6891
+kumoai/client/pquery.py,sha256=IQ8As-OOJOkuMoMosphOsA5hxQYLCbzOQJO7RezK8uY,7091
 kumoai/client/rfm.py,sha256=NxKk8mH2A-B58rSXhDWaph4KeiSyJYDq-RO-vAHh7es,3726
 kumoai/client/source_table.py,sha256=VCsCcM7KYcnjGP7HLTb-AOSEGEVsJTWjk8bMg1JdgPU,2101
 kumoai/client/table.py,sha256=cQG-RPm-e91idEgse1IPJDvBmzddIDGDkuyrR1rq4wU,3235
@@ -49,37 +49,52 @@ kumoai/connector/glue_connector.py,sha256=HivT0QYQ8-XeB4QLgWvghiqXuq7jyBK9G2R1py
 kumoai/connector/s3_connector.py,sha256=3kbv-h7DwD8O260Q0h1GPm5wwQpLt-Tb3d_CBSaie44,10155
 kumoai/connector/snowflake_connector.py,sha256=K0s-H9tW3rve8g2x1PbyxvzSpkROfGQZz-Qa4PoT4UE,9022
 kumoai/connector/source_table.py,sha256=QLT8bEYaxeMwy-b168url0VfnkTrs5K6VKLbxTI4hEY,17539
-kumoai/connector/utils.py,sha256=PUjunLpfqMZsrPDo2EmnyJRBl_mt-E6ugv2kNkf5Rn8,64011
+kumoai/connector/utils.py,sha256=wlqQxMmPvnFNoCcczGkKYjSu05h8OhWh4fhTzQm_2bQ,64694
 kumoai/encoder/__init__.py,sha256=VPGs4miBC_WfwWeOXeHhFomOUocERFavhKf5fqITcds,182
 kumoai/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kumoai/experimental/rfm/__init__.py,sha256=wKfMKTxfuJNH1GCWGZ7-288HXil0tsCuXqg-BAFctZE,6812
+kumoai/experimental/rfm/__init__.py,sha256=slliYcrh80xPtQQ_nnsp3ny9IbmHCyirmdZUfKTdME4,6064
 kumoai/experimental/rfm/authenticate.py,sha256=FiuHMvP7V3zBZUlHMDMbNLhc-UgDZgz4hjVSTuQ7DRw,18888
-kumoai/experimental/rfm/local_graph.py,sha256=2iJDlsGVzqCe1bD_puXWlhwGkn7YnQyJ4p4C-fwCZNE,30076
-kumoai/experimental/rfm/local_graph_sampler.py,sha256=5DbhL9h0usFKSJfnx7HjLMPcG54qwJ48M2tmONqxXyY,6672
-kumoai/experimental/rfm/local_graph_store.py,sha256=8BqonuaMftAAsjgZpB369i5AeNd1PkisMbbEqc0cKBo,13847
-kumoai/experimental/rfm/local_pquery_driver.py,sha256=aO7Jfwx9gxGKYvpqxZx1LLWdI1MhuZQOPtAITxoOQO0,26162
-kumoai/experimental/rfm/local_table.py,sha256=r8xZ33Mjs6JD8ud6h23tZ99Dag2DvZ4h6tWjmGrKQg4,19605
-kumoai/experimental/rfm/rfm.py,sha256=8SvGWfMuRYJgiz5OTplu7m47mDrHAjQ2mRZtRASnSCk,48136
-kumoai/experimental/rfm/sagemaker.py,sha256=e0rRQ28WcgAk_ALqUyU20d193c8_68rCkSengZIHu3Y,4823
-kumoai/experimental/rfm/utils.py,sha256=3IiBvT_aLBkkcJh3H11_50yt_XlEzHR0cm9Kprrtl8k,11123
-kumoai/experimental/rfm/infer/__init__.py,sha256=xQ8_SuejIzXyn2J7bIKX3pXumFtRuEfBtE5oEDUDJjI,293
+kumoai/experimental/rfm/graph.py,sha256=awVJSk4cWRMacS5CJvJtR8TR56FEbrJPcQCukNydQOc,40392
+kumoai/experimental/rfm/rfm.py,sha256=j4LfCbmHPs9RrwRNjTEE-sPUC6RcHNcymKBtG_Rt-M4,49670
+kumoai/experimental/rfm/sagemaker.py,sha256=_hTrFg4qfXe7uzwqSEG_wze-IFkwn7qde9OpUodCpbc,4982
+kumoai/experimental/rfm/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kumoai/experimental/rfm/backend/local/__init__.py,sha256=2s9sSA-E-8pfkkzCH4XPuaSxSznEURMfMgwEIfYYPsg,1014
+kumoai/experimental/rfm/backend/local/graph_store.py,sha256=5cHuExHljU_Z56KV3s-PwzeiLuPKgh2mCxcjTMmPZ8E,11928
+kumoai/experimental/rfm/backend/local/sampler.py,sha256=wJXQdiOOCIjqpNkDoXgPV0cnsKjeJ2Zi77-pJ8zAEM4,10738
+kumoai/experimental/rfm/backend/local/table.py,sha256=-R_9nncosByAfSMfUt6HgCUNoW_MLGJW3F5SnAd4Ru0,3744
+kumoai/experimental/rfm/backend/snow/__init__.py,sha256=BYfsiuJ4Ee30GjG9EuUtitMHXnRfvVKi85zNlIwldV4,993
+kumoai/experimental/rfm/backend/snow/sampler.py,sha256=d5hyKyiNxnjmINGxM7K5nCVnN7ZaknM1QdHLt8EBt5c,4287
+kumoai/experimental/rfm/backend/snow/table.py,sha256=bnBGgkXmgLcSZlr0t4pYo9aSk7DO5PuSDM06HBgA5BE,4841
+kumoai/experimental/rfm/backend/sqlite/__init__.py,sha256=cA-PZL1oTaLxthZbfLSudexImtF6jRsGkdjSp-66dCM,914
+kumoai/experimental/rfm/backend/sqlite/sampler.py,sha256=n3etgKH7viiIgMXHFLM3KQvKvCpylXRTw0WBNoSKd7s,4057
+kumoai/experimental/rfm/backend/sqlite/table.py,sha256=9F2_6iqgM3h9jkIcecHnEwQq-UG3RBmR9Y9XPcFLR04,4161
+kumoai/experimental/rfm/base/__init__.py,sha256=IdQwFFEeXZnTBkj1I65qQDL6QlcbmjVoTP6-quCtmFc,481
+kumoai/experimental/rfm/base/column.py,sha256=izCJmufJcd1RSi-ptFMfrue-JYag38MJxizka7ya0-A,2319
+kumoai/experimental/rfm/base/sampler.py,sha256=aCD98t0CUhAvGXEFv24Vq2g4otuclpKkkyL1rMR_mFg,31449
+kumoai/experimental/rfm/base/source.py,sha256=RqlI_kBoRV0ADb8KdEKn15RNHMdFUzEVzb57lIoyBM4,294
+kumoai/experimental/rfm/base/table.py,sha256=neGldEZaweoJ8VRgnEnaSpAISSkSTkgXxItuuywBM4E,20010
+kumoai/experimental/rfm/infer/__init__.py,sha256=krdMFN8iKZlSFOl-M5MW1KuSviQV3H1E18jj2uB8g6Q,469
 kumoai/experimental/rfm/infer/categorical.py,sha256=VwNaKwKbRYkTxEJ1R6gziffC8dGsEThcDEfbi-KqW5c,853
+kumoai/experimental/rfm/infer/dtype.py,sha256=ZZ6ztqJnTR1CaC2z5Uhf0o0rSdNThnss5tem5JNQkck,2607
 kumoai/experimental/rfm/infer/id.py,sha256=ZIO0DWIoiEoS_8MVc5lkqBfkTWWQ0yGCgjkwLdaYa_Q,908
 kumoai/experimental/rfm/infer/multicategorical.py,sha256=0-cLpDnGryhr76QhZNO-klKokJ6MUSfxXcGdQ61oykY,1102
+kumoai/experimental/rfm/infer/pkey.py,sha256=ubNqW1LIjLKiXbjXELAY3g6n2f3u2Eis_uC2DEiXFiU,4393
+kumoai/experimental/rfm/infer/time_col.py,sha256=7R5Itl8RRBOr61qLpRTanIqrUVZFZcAXzDA9lCw4nx4,1820
 kumoai/experimental/rfm/infer/timestamp.py,sha256=vM9--7eStzaGG13Y-oLYlpNJyhL6f9dp17HDXwtl_DM,1094
 kumoai/experimental/rfm/pquery/__init__.py,sha256=X0O3EIq5SMfBEE-ii5Cq6iDhR3s3XMXB52Cx5htoePw,152
 kumoai/experimental/rfm/pquery/executor.py,sha256=f7-pJhL0BgFU9E4o4gQpQyArOvyrZtwxFmks34-QOAE,2741
-kumoai/experimental/rfm/pquery/pandas_executor.py,sha256=kiBJq7uVGbasG7TiqsubEl6ey3UYzZiM4bwxILqp_54,18487
+kumoai/experimental/rfm/pquery/pandas_executor.py,sha256=wYI9a3smClR2pQGwsYRdmpOm0PlUsbtyW9wpAVpCEe4,18492
 kumoai/graph/__init__.py,sha256=n8X4X8luox4hPBHTRC9R-3JzvYYMoR8n7lF1H4w4Hzc,228
 kumoai/graph/column.py,sha256=t7wBmcx0VYKXjIoESU9Nq-AisiJOdlqd80t8zby1R8Y,4189
 kumoai/graph/graph.py,sha256=iyp4klPIMn2ttuEqMJvsrxKb_tmz_DTnvziIhCegduM,38291
 kumoai/graph/table.py,sha256=nZqYX8xlyAz6kVtlE2vf9BAIOCoWeFNIfbGbReDCb7k,33888
 kumoai/pquery/__init__.py,sha256=uTXr7t1eXcVfM-ETaM_1ImfEqhrmaj8BjiIvy1YZTL8,533
 kumoai/pquery/prediction_table.py,sha256=QPDH22X1UB0NIufY7qGuV2XW7brG3Pv--FbjNezzM2g,10776
-kumoai/pquery/predictive_query.py,sha256=oUqwdOWLLkPM-G4PhpUk_6mwSJGBtaD3t37Wp5Oow8M,24971
+kumoai/pquery/predictive_query.py,sha256=UXn1s8ztubYZMNGl4ijaeidMiGlFveb1TGw9qI5-TAo,24901
 kumoai/pquery/training_table.py,sha256=elmPDZx11kPiC_dkOhJcBUGtHKgL32GCBvZ9k6U0pMg,15809
 kumoai/testing/__init__.py,sha256=goHIIo3JE7uHV7njo4_aTd89mVVR74BEAZ2uyBaOR0w,170
-kumoai/testing/decorators.py,sha256=RiFrJcP-ym-mB1BYSGC26bBiryxoR9-GwL1G4EHc2sc,1591
+kumoai/testing/decorators.py,sha256=83tMifuPTpUqX7zHxMttkj1TDdB62EBtAP-Fjj72Zdo,1607
+kumoai/testing/snow.py,sha256=ubx3yJP0UHxsNiar1-jNdv8ZfszKc8Js3_Gg70uf008,1487
 kumoai/trainer/__init__.py,sha256=zUdFl-f-sBWmm2x8R-rdVzPBeU2FaMzUY5mkcgoTa1k,939
 kumoai/trainer/baseline_trainer.py,sha256=LlfViNOmswNv4c6zJJLsyv0pC2mM2WKMGYx06ogtEVc,4024
 kumoai/trainer/config.py,sha256=-2RfK10AsVVThSyfWtlyfH4Fc4EwTdu0V3yrDRtIOjk,98
@@ -87,12 +102,13 @@ kumoai/trainer/job.py,sha256=Wk69nzFhbvuA3nEvtCstI04z5CxkgvQ6tHnGchE0Lkg,44938
 kumoai/trainer/online_serving.py,sha256=9cddb5paeZaCgbUeceQdAOxysCtV5XP-KcsgFz_XR5w,9566
 kumoai/trainer/trainer.py,sha256=hBXO7gwpo3t59zKFTeIkK65B8QRmWCwO33sbDuEAPlY,20133
 kumoai/trainer/util.py,sha256=bDPGkMF9KOy4HgtA-OwhXP17z9cbrfMnZGtyGuUq_Eo,4062
-kumoai/utils/__init__.py,sha256=wGDC_31XJ-7ipm6eawjLAJaP4EfmtNOH8BHzaetQ9Ko,268
+kumoai/utils/__init__.py,sha256=cF5ACzp1X61sqhlCHc6biQk6fc4gW_oyhGsBrjx-SoM,316
 kumoai/utils/datasets.py,sha256=ptKIUoBONVD55pTVNdRCkQT3NWdN_r9UAUu4xewPa3U,2928
 kumoai/utils/forecasting.py,sha256=-nDS6ucKNfQhTQOfebjefj0wwWH3-KYNslIomxwwMBM,7415
 kumoai/utils/progress_logger.py,sha256=pngEGzMHkiOUKOa6fbzxCEc2xlA4SJKV4TDTVVoqObM,5062
-kumoai-2.13.0.dev202511211730.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
-kumoai-2.13.0.dev202511211730.dist-info/METADATA,sha256=S4KjjwNYqtLxfkX6vqdiOoo-iUPcfoXnKGJif23K9jU,2475
-kumoai-2.13.0.dev202511211730.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-kumoai-2.13.0.dev202511211730.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
-kumoai-2.13.0.dev202511211730.dist-info/RECORD,,
+kumoai/utils/sql.py,sha256=f6lR6rBEW7Dtk0NdM26dOZXUHDizEHb1WPlBCJrwoq0,118
+kumoai-2.14.0.dev202512141732.dist-info/licenses/LICENSE,sha256=TbWlyqRmhq9PEzCaTI0H0nWLQCCOywQM8wYH8MbjfLo,1102
+kumoai-2.14.0.dev202512141732.dist-info/METADATA,sha256=afHVHJGv5rxq96TO9d3VqyuN_Nhvf-dFGx3nDJo41lk,2510
+kumoai-2.14.0.dev202512141732.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+kumoai-2.14.0.dev202512141732.dist-info/top_level.txt,sha256=YjU6UcmomoDx30vEXLsOU784ED7VztQOsFApk1SFwvs,7
+kumoai-2.14.0.dev202512141732.dist-info/RECORD,,

kumoai/experimental/rfm/local_graph_sampler.py

@@ -1,182 +0,0 @@
-from typing import Dict, List, Optional, Tuple
-
-import numpy as np
-import pandas as pd
-from kumoapi.rfm.context import EdgeLayout, Link, Subgraph, Table
-from kumoapi.typing import Stype
-
-import kumoai.kumolib as kumolib
-from kumoai.experimental.rfm.local_graph_store import LocalGraphStore
-from kumoai.experimental.rfm.utils import normalize_text
-
-
-class LocalGraphSampler:
-    def __init__(self, graph_store: LocalGraphStore) -> None:
-        self._graph_store = graph_store
-        self._sampler = kumolib.NeighborSampler(
-            self._graph_store.node_types,
-            self._graph_store.edge_types,
-            {
-                '__'.join(edge_type): colptr
-                for edge_type, colptr in self._graph_store.colptr_dict.items()
-            },
-            {
-                '__'.join(edge_type): row
-                for edge_type, row in self._graph_store.row_dict.items()
-            },
-            self._graph_store.time_dict,
-        )
-
-    def __call__(
-        self,
-        entity_table_names: Tuple[str, ...],
-        node: np.ndarray,
-        time: np.ndarray,
-        num_neighbors: List[int],
-        exclude_cols_dict: Dict[str, List[str]],
-    ) -> Subgraph:
-
-        (
-            row_dict,
-            col_dict,
-            node_dict,
-            batch_dict,
-            num_sampled_nodes_dict,
-            num_sampled_edges_dict,
-        ) = self._sampler.sample(
-            {
-                '__'.join(edge_type): num_neighbors
-                for edge_type in self._graph_store.edge_types
-            },
-            {},  # time interval based sampling
-            entity_table_names[0],
-            node,
-            time // 1000**3,  # nanoseconds to seconds
-        )
-
-        table_dict: Dict[str, Table] = {}
-        for table_name, node in node_dict.items():
-            batch = batch_dict[table_name]
-
-            if len(node) == 0:
-                continue
-
-            df = self._graph_store.df_dict[table_name]
-
-            num_sampled_nodes = num_sampled_nodes_dict[table_name].tolist()
-            stype_dict = {  # Exclude target columns:
-                column_name: stype
-                for column_name, stype in
-                self._graph_store.stype_dict[table_name].items()
-                if column_name not in exclude_cols_dict.get(table_name, [])
-            }
-            primary_key: Optional[str] = None
-            if table_name in entity_table_names:
-                primary_key = self._graph_store.pkey_name_dict.get(table_name)
-
-            columns: List[str] = []
-            if table_name in entity_table_names:
-                columns += [self._graph_store.pkey_name_dict[table_name]]
-            columns += list(stype_dict.keys())
-
-            if len(columns) == 0:
-                table_dict[table_name] = Table(
-                    df=pd.DataFrame(index=range(len(node))),
-                    row=None,
-                    batch=batch,
-                    num_sampled_nodes=num_sampled_nodes,
-                    stype_dict=stype_dict,
-                    primary_key=primary_key,
-                )
-                continue
-
-            row: Optional[np.ndarray] = None
-            if table_name in self._graph_store.end_time_column_dict:
-                # Set end time to NaT for all values greater than anchor time:
-                df = df.iloc[node].reset_index(drop=True)
-                col_name = self._graph_store.end_time_column_dict[table_name]
-                ser = df[col_name]
-                value = ser.astype('datetime64[ns]').astype(int).to_numpy()
-                mask = value > time[batch]
-                df.loc[mask, col_name] = pd.NaT
-            else:
-                # Only store unique rows in `df` above a certain threshold:
-                unique_node, inverse = np.unique(node, return_inverse=True)
-                if len(node) > 1.05 * len(unique_node):
-                    df = df.iloc[unique_node].reset_index(drop=True)
-                    row = inverse
-                else:
-                    df = df.iloc[node].reset_index(drop=True)
-
-            # Filter data frame to minimal set of columns:
-            df = df[columns]
-
-            # Normalize text (if not already pre-processed):
-            for column_name, stype in stype_dict.items():
-                if stype == Stype.text:
-                    df[column_name] = normalize_text(df[column_name])
-
-            table_dict[table_name] = Table(
-                df=df,
-                row=row,
-                batch=batch,
-                num_sampled_nodes=num_sampled_nodes,
-                stype_dict=stype_dict,
-                primary_key=primary_key,
-            )
-
-        link_dict: Dict[Tuple[str, str, str], Link] = {}
-        for edge_type in self._graph_store.edge_types:
-            edge_type_str = '__'.join(edge_type)
-
-            row = row_dict[edge_type_str]
-            col = col_dict[edge_type_str]
-
-            if len(row) == 0:
-                continue
-
-            # Do not store reverse edge type if it is a replica:
-            rev_edge_type = Subgraph.rev_edge_type(edge_type)
-            rev_edge_type_str = '__'.join(rev_edge_type)
-            if (rev_edge_type in link_dict
-                    and np.array_equal(row, col_dict[rev_edge_type_str])
-                    and np.array_equal(col, row_dict[rev_edge_type_str])):
-                link = Link(
-                    layout=EdgeLayout.REV,
-                    row=None,
-                    col=None,
-                    num_sampled_edges=(
-                        num_sampled_edges_dict[edge_type_str].tolist()),
-                )
-                link_dict[edge_type] = link
-                continue
-
-            layout = EdgeLayout.COO
-            if np.array_equal(row, np.arange(len(row))):
-                row = None
-            if np.array_equal(col, np.arange(len(col))):
-                col = None
-
-            # Store in compressed representation if more efficient:
-            num_cols = table_dict[edge_type[2]].num_rows
-            if col is not None and len(col) > num_cols + 1:
-                layout = EdgeLayout.CSC
-                colcount = np.bincount(col, minlength=num_cols)
-                col = np.empty(num_cols + 1, dtype=col.dtype)
-                col[0] = 0
-                np.cumsum(colcount, out=col[1:])
-
-            link = Link(
-                layout=layout,
-                row=row,
-                col=col,
-                num_sampled_edges=(
-                    num_sampled_edges_dict[edge_type_str].tolist()),
-            )
-            link_dict[edge_type] = link
-
-        return Subgraph(
-            anchor_time=time,
-            table_dict=table_dict,
-            link_dict=link_dict,
-        )