chalkpy 2.96.6__py3-none-any.whl → 2.96.8__py3-none-any.whl

chalk/_version.py

@@ -1 +1 @@
-__version__ = "2.96.6"
+__version__ = "2.96.8"

chalk/client/client.py

@@ -33,6 +33,7 @@ from chalk.client.models import (
     GetRegisteredModelResponse,
     GetRegisteredModelVersionResponse,
     ManualTriggerScheduledQueryResponse,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
@@ -526,6 +527,7 @@ class ChalkClient:
         explain: bool = False,
         num_input_rows: Optional[int] = None,
         headers: Mapping[str, str] | None = None,
+        planner_options: Mapping[str, str | int | bool] | None = None,
     ) -> PlanQueryResponse:
         """Plan a query without executing it.
 
@@ -575,6 +577,10 @@ class ChalkClient:
             The number of input rows that this plan will be run with. If unknown, specify `None`.
         headers
             Additional headers to provide with the request
+        planner_options
+            Dictionary of additional options to pass to the Chalk query engine.
+            Values may be provided as part of conversations with Chalk support
+            to enable or disable specific functionality.
 
         Returns
         -------
@@ -922,7 +928,7 @@ class ChalkClient:
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -1038,8 +1044,8 @@ class ChalkClient:
             If specified, the query will be run asynchronously across a maximum `num_workers` pod workers at any time.
             This parameter is useful if you have a large number of shards and would like to limit the number of pods running at once.
         completion_deadline
-            If specified, shards must complete within 'completion_deadline' duration, or they will be terminated.
-            Terminated shards can be tried.
+            If specified as a timedelta, applies a completion deadline to each shard; each shard's query will fail (allowing retries) if it does not complete within the duration.
+            If specified as an OfflineQueryDeadlineOptions, allows more fine-grained control of shard- or query-level deadlines, with options to retry on failure or not.
         max_retries
             If specified, failed offline query shards will be retried. The retry budget is shared across all shards.
             By default, max_retries=num_shards/

chalk/client/client_async.py

@@ -12,6 +12,7 @@ from chalk.client.models import (
     BulkOnlineQueryResult,
     FeatureReference,
     FeatureStatisticsResponse,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
@@ -486,6 +487,7 @@ class AsyncChalkClient:
         explain: bool = False,
         num_input_rows: Optional[int] = None,
         headers: Mapping[str, str] | None = None,
+        planner_options: Mapping[str, str | int | bool] | None = None,
     ) -> PlanQueryResponse:
         """Plan a query without executing it.
 
@@ -535,6 +537,10 @@ class AsyncChalkClient:
             The number of input rows that this plan will be run with. If unknown, specify `None`.
         headers
             Additional headers to provide with the request
+        planner_options
+            Dictionary of additional options to pass to the Chalk query engine.
+            Values may be provided as part of conversations with Chalk support
+            to enable or disable specific functionality.
 
         Returns
         -------
@@ -642,7 +648,7 @@ class AsyncChalkClient:
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -745,6 +751,9 @@ class AsyncChalkClient:
         num_workers
             If specified, the query will be run asynchronously across a maximum `num_workers` pod workers at any time.
             This parameter is useful if you have a large number of shards and would like to limit the number of pods running at once.
+        completion_deadline
+            If specified as a timedelta, applies a completion deadline to each shard; each shard's query will fail (allowing retries) if it does not complete within the duration.
+            If specified as an OfflineQueryDeadlineOptions, allows more fine-grained control of shard- or query-level deadlines, with options to retry on failure or not.
         query_name
             The name of the query to execute. If provided, will create a new named query or fill in missing parameters from a preexisting execution.
         query_name_version

chalk/client/client_impl.py

@@ -103,6 +103,7 @@ from chalk.client.models import (
     MultiUploadFeaturesRequest,
     MultiUploadFeaturesResponse,
     OfflineQueryContext,
+    OfflineQueryDeadlineOptions,
     OfflineQueryInput,
     OfflineQueryInputSql,
     OfflineQueryInputUri,
@@ -2229,7 +2230,7 @@ https://docs.chalk.ai/cli/apply
         store_offline: bool = False,
         num_shards: int | None = None,
         num_workers: int | None = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
@@ -3607,7 +3608,7 @@ https://docs.chalk.ai/cli/apply
         num_shards: int | None = None,
         num_workers: int | None = None,
         feature_for_lower_upper_bound: Optional[str] = None,
-        completion_deadline: timedelta | None = None,
+        completion_deadline: Union[timedelta, OfflineQueryDeadlineOptions, None] = None,
         max_retries: int | None = None,
         optional_output_expressions: Optional[List[str]] = None,
         required_output_expressions: Optional[List[str]] = None,
@@ -3651,6 +3652,13 @@ https://docs.chalk.ai/cli/apply
         upper_bound_str = process_bound(upper_bound)
         if branch is ...:
             branch = self._branch
+
+        retyped_completion_deadline: Union[None, str, OfflineQueryDeadlineOptions] = None
+        if isinstance(completion_deadline, OfflineQueryDeadlineOptions):
+            retyped_completion_deadline = completion_deadline.with_chalk_durations()
+        elif isinstance(completion_deadline, timedelta):
+            retyped_completion_deadline = timedelta_to_duration(completion_deadline)
+
         req = CreateOfflineQueryJobRequest(
             output=optional_output,
             output_expressions=optional_output_expressions or [],
@@ -3683,7 +3691,7 @@ https://docs.chalk.ai/cli/apply
             num_shards=num_shards,
             num_workers=num_workers,
             feature_for_lower_upper_bound=feature_for_lower_upper_bound,
-            completion_deadline=timedelta_to_duration(completion_deadline) if completion_deadline is not None else None,
+            completion_deadline=retyped_completion_deadline,
             max_retries=max_retries,
             use_job_queue=use_job_queue,
             overlay_graph=_get_overlay_graph_b64(),
@@ -4438,6 +4446,7 @@ https://docs.chalk.ai/cli/apply
         explain: bool = False,
         num_input_rows: Optional[int] = None,
         headers: Mapping[str, str] | None = None,
+        planner_options: Mapping[str, str | int | bool] | None = None,
     ) -> PlanQueryResponse:
         encoded_inputs = encode_outputs(input).string_outputs
         outputs = encode_outputs(output).string_outputs
@@ -4472,6 +4481,7 @@ https://docs.chalk.ai/cli/apply
             store_plan_stages=store_plan_stages,
             explain=explain,
             num_input_rows=num_input_rows,
+            planner_options=planner_options,
         )
 
         extra_headers: dict[str, str] = {}

chalk/client/models.py

@@ -20,6 +20,7 @@ from chalk.features.tag import EnvironmentId
 from chalk.prompts import Prompt
 from chalk.queries.query_context import ContextJsonDict
 from chalk.utils.df_utils import read_parquet
+from chalk.utils.duration import timedelta_to_duration
 from chalk.utils.missing_dependency import missing_dependency_exception
 
 if TYPE_CHECKING:
@@ -820,6 +821,48 @@ class ResourceRequests(BaseModel):
     """Resource group to use for this job. If not specified, the default resource group will be used."""
 
 
+class OfflineQueryDeadlineOptions(BaseModel):
+    """
+    Specification for setting deadlines for shards of the query or the entire query itself.
+    """
+
+    shard_deadline: Union[timedelta, str, None] = None
+    """
+    Maximum amount of time a query shard can work before being failed.
+    """
+
+    retry_on_shard_deadline: Optional[bool] = None
+    """
+    Whether to retry when the per-shard deadline is triggered. Will default to true.
+    """
+
+    query_deadline: Union[timedelta, str, None] = None
+    """
+    Maximum amount of time that the entire query can work before being failed.
+    """
+
+    retry_on_query_deadline: Optional[bool] = None
+    """
+    Whether to retry when the entire query's deadline is triggered. Will default to false.
+    """
+
+    def with_chalk_durations(self) -> OfflineQueryDeadlineOptions:
+        return OfflineQueryDeadlineOptions(
+            shard_deadline=(
+                timedelta_to_duration(self.shard_deadline)
+                if isinstance(self.shard_deadline, timedelta)
+                else self.shard_deadline
+            ),
+            retry_on_shard_deadline=self.retry_on_shard_deadline,
+            query_deadline=(
+                timedelta_to_duration(self.query_deadline)
+                if isinstance(self.query_deadline, timedelta)
+                else self.query_deadline
+            ),
+            retry_on_query_deadline=self.retry_on_query_deadline,
+        )
+
+
 class CreateOfflineQueryJobRequest(BaseModel):
     output: List[str]
     """A list of output feature root fqns to query"""
@@ -902,7 +945,7 @@ class CreateOfflineQueryJobRequest(BaseModel):
     num_workers: Optional[int] = None
     feature_for_lower_upper_bound: Optional[str] = None
 
-    completion_deadline: Optional[str] = None
+    completion_deadline: Union[None, str, OfflineQueryDeadlineOptions] = None
     max_retries: Optional[int] = None
 
     use_job_queue: bool = False
@@ -915,7 +958,7 @@ class CreateOfflineQueryJobRequest(BaseModel):
     @root_validator
     def _validate_multiple_computers(cls, values: Dict[str, Any]):
         if values["input"] is None or isinstance(
-            values["input"], (UploadedParquetShardedOfflineQueryInput, OfflineQueryInputUri)
+            values["input"], (UploadedParquetShardedOfflineQueryInput, OfflineQueryInputUri, OfflineQueryInputSql)
         ):
             return values
         expected_use_multiple_computers = isinstance(values["input"], tuple)
@@ -1653,6 +1696,7 @@ class PlanQueryRequest(BaseModel):
     explain: bool = False
     store_plan_stages: bool = False
     encoding_options: FeatureEncodingOptions = FeatureEncodingOptions()
+    planner_options: Mapping[str, str | int | bool | float] | None = None
 
 
 class FeatureSchema(BaseModel):

chalk/df/LazyFramePlaceholder.py

@@ -101,6 +101,11 @@ class LazyFramePlaceholder:
 
     __str__ = __repr__
 
+    def _is_equal(self, other: LazyFramePlaceholder) -> bool:
+        # proto equality is janky but it's hard to write a good eq method here given
+        # we have dicts and the proto round trip is slightly lossy on tuples vs lists
+        return self._to_proto() == other._to_proto()
+
     def _to_proto(self) -> dataframe_pb2.DataFramePlan:
         """
         Convert this proto plan to a dataframe.

chalk/features/_encoding/converter.py

@@ -940,10 +940,9 @@ class PrimitiveFeatureConverter(Generic[_TPrim]):
 
     @staticmethod
     def convert_arrow_table_from_proto(proto: pb.TableParquetBytes) -> pa.Table:
-        import pyarrow.parquet
+        import pyarrow.parquet as pq
 
-        pf = pyarrow.parquet.ParquetFile(io.BytesIO(proto.encoded_parquet_bytes))
-        return pyarrow.parquet.read_table(pf)
+        return pq.read_table(io.BytesIO(proto.encoded_parquet_bytes))
 
     @staticmethod
     def _serialize_pa_decimal_to_pb(value: Union[pa.Decimal128Scalar, pa.Decimal256Scalar]) -> pb.ScalarValue:

{chalkpy-2.96.6.dist-info → chalkpy-2.96.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.96.6
+Version: 2.96.8
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalkpy-2.96.6.dist-info → chalkpy-2.96.8.dist-info}/RECORD

@@ -1,5 +1,5 @@
 chalk/__init__.py,sha256=vKsx9-cl5kImlVWGHVRYO6bweBm79NAzGs3l36u71wM,2657
-chalk/_version.py,sha256=Y-Kl04ucuWNOt_WOW6dvhcamZanMwDrSw7NrZ8zKeSs,23
+chalk/_version.py,sha256=XQBn2m_TF6iSSjpgoPtUzJ93WpjTyMNEkfPlAWxDliQ,23
 chalk/cli.py,sha256=ckqqfOI-A2mT23-rnZzDMmblYj-2x1VBX8ebHlIEn9A,5873
 chalk/importer.py,sha256=m4lMn1lSYj_euDq8CS7LYTBnek9JOcjGJf9-82dJHbA,64441
 chalk/prompts.py,sha256=2H9UomLAamdfRTNUdKs9i3VTpiossuyRhntqsAXUhhg,16117
@@ -612,14 +612,14 @@ chalk/_validation/validation.py,sha256=9cCMfZa9-1wxkXLme_ylmD5vIA1qExJD6aqbYvbmK
 chalk/byte_transmit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/byte_transmit/model.py,sha256=LFX8pj9X_CWXeap7fDnMl9YmXsYTgq7jBAbEWkxoYoE,13048
 chalk/client/__init__.py,sha256=wu3WQVzE5gRj6noQQDOdYJUgDaz_9QtbjXH4KuoIlXQ,1782
-chalk/client/client.py,sha256=fqw75x4yiAk3CXryGbuF_qg09WDIlluhYeFdQ-BfTG8,103337
-chalk/client/client_async.py,sha256=nFFTWJbdBlb7zksyjOMBY566tZTAyNXQhCnq06LHWl0,50803
+chalk/client/client.py,sha256=59UYbIq7KHwNm1ZFw5TmCV2JvbJxn45Ar_AjfSkuzQ8,103907
+chalk/client/client_async.py,sha256=wD38RIkwVLaKYKj7K1bWDXnVZkLFYHsbY_5VhtflMvo,51559
 chalk/client/client_async_impl.py,sha256=ZphhgTB49JBWHCGXe-dI0wWWKc9zPcOczy02q_gFy50,6925
 chalk/client/client_grpc.py,sha256=skpSHYCyE8hg0T3Bcl5R8MMi92QjrFtxQRu_NxCSVlw,106952
-chalk/client/client_impl.py,sha256=rvT02yriRo7_WFqbUsHXRjm6DbRyXzw2QjcLwKaW10E,211844
+chalk/client/client_impl.py,sha256=3NDl1Vo7YJ0grraxJeN6vjFj1hBpU7YT_ikolK31D6k,212395
 chalk/client/dataset.py,sha256=LneWwaAOHCjtj7gaJjsSeVNruj-QJ51hjRi62zrFNVE,77561
 chalk/client/exc.py,sha256=kZJ80YbSeSRDmTLTh240j_eRdJFZBa7IaDsNSRoDroU,4145
-chalk/client/models.py,sha256=FDgEqIeBJq25T9AH38jdWkeyRLiAXfLXfXHpomFJM08,63923
+chalk/client/models.py,sha256=YwXywN6OSPOda1DKJuPVe6hgsHzf8ewvX2dYOBOZk7I,65534
 chalk/client/response.py,sha256=m8sQCOj7YVv3mZSZMIC1rIMzFMQ9rfMdBRLg5NRmOOE,53257
 chalk/client/_internal_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/client/_internal_models/check.py,sha256=3Xfo4Ws4rvwjeVg0-5-kejfRfRBJeqHmnRhW-WEz784,917
@@ -636,7 +636,7 @@ chalk/config/_validator.py,sha256=QC1y52m704_bV_TYjq0sdZJ-km8iSkDX1V4sHgw4RJk,13
 chalk/config/auth_config.py,sha256=HAALkQrvDD0i7gaZK5iufS8vDDVbzLIpHLOpcJO1kmw,4498
 chalk/config/project_config.py,sha256=YHB3upvtBJu-tWfNOchBRSc9xGebDbrIpCVmKbBzJ8Q,7217
 chalk/df/ChalkDataFrameImpl.py,sha256=BRwnjQcie3gxaKhKau8HG17NWzS1zdr8SnNVurxF8QY,133
-chalk/df/LazyFramePlaceholder.py,sha256=uTinB2buEFBT8djBCqiPN2hKmr_nU7n4i_DcCLNLuWs,39111
+chalk/df/LazyFramePlaceholder.py,sha256=D4oYEGPfoxpJMmyiDJq3G9Wfl1jF3XSXD71q-GxOrmQ,39398
 chalk/df/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/df/ast_parser.py,sha256=t-DwUxd2hN8LECRSBx85DIv9FtILgMiHyGyCTINfgQw,11199
 chalk/features/__init__.py,sha256=5doD7bFwRthzwdmizbRaPVHiCABiNpiOiAJVFlwqNnA,6943
@@ -673,7 +673,7 @@ chalk/features/_embedding/sentence_transformer.py,sha256=hNYuT9D-16C49lkhIVK_KXZ
 chalk/features/_embedding/utils.py,sha256=gUYUJCt-9XLeioC3QYg42yc5f69ywAr_w_Q-xoF3o54,877
 chalk/features/_embedding/vertexai.py,sha256=PUaMmd7qbK9vXrlce_8BQHg_OM9L982k6LuyM_PU-5w,5362
 chalk/features/_encoding/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-chalk/features/_encoding/converter.py,sha256=YMIFZLAiN7z8JhDj4hsrH4qQTZ0pFBIwaJN6ULVsjd8,66221
+chalk/features/_encoding/converter.py,sha256=YLNX4hxoT5GlP_n_cL5anQlIzfMW0TfBuuK1TBccaiI,66169
 chalk/features/_encoding/http.py,sha256=MVmkgv90O4sOnImc-FrRPVwYR2U7zOGduNkHO4umoe4,2227
 chalk/features/_encoding/inputs.py,sha256=4K4kOCAl0MuobAC_lVJE70toQVrPt2ZlzrlwYALuLHE,10151
 chalk/features/_encoding/json.py,sha256=lHSBWIjNKMDmt_AmiFtf4ZsW5QPIEC6r7NvDW0IdqPI,13804
@@ -827,8 +827,8 @@ chalk/utils/tracing.py,sha256=NiiM-9dbuJhSCv6R1npR1uYNKWlkqTR6Ygm0Voi2NrY,13078
 chalk/utils/weak_set_by_identity.py,sha256=VmikA_laYwFeOphCwXJIuyOIkrdlQe0bSzaXq7onoQw,953
 chalk/utils/pydanticutil/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 chalk/utils/pydanticutil/pydantic_compat.py,sha256=O575lLYJ5GvZC4HMzR9yATxf9XwjC6NrDUXbNwZidlE,3031
-chalkpy-2.96.6.dist-info/METADATA,sha256=5dqjMynKBZ6H6Y6tbYNazvRRwyCLZblyJFsi6E_PA1E,27754
-chalkpy-2.96.6.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-chalkpy-2.96.6.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
-chalkpy-2.96.6.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
-chalkpy-2.96.6.dist-info/RECORD,,
+chalkpy-2.96.8.dist-info/METADATA,sha256=wLMi2G7i8tczLoHoAAYiMWG9g2oujf4DoUmq3hqkP98,27754
+chalkpy-2.96.8.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+chalkpy-2.96.8.dist-info/entry_points.txt,sha256=Vg23sd8icwq-morJrljVFr-kQnMbm95rZfZj5wsZGis,42
+chalkpy-2.96.8.dist-info/top_level.txt,sha256=1Q6_19IGYfNxSw50W8tYKEJ2t5HKQ3W9Wiw4ia5yg2c,6
+chalkpy-2.96.8.dist-info/RECORD,,