chalkpy 2.92.2__tar.gz → 2.92.3__tar.gz

{chalkpy-2.92.2 → chalkpy-2.92.3}/PKG-INFO

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

chalkpy-2.92.3/chalk/_version.py

@@ -0,0 +1 @@
+__version__ = "2.92.3"

{chalkpy-2.92.2 → chalkpy-2.92.3}/chalk/client/client_async.py

@@ -9,6 +9,7 @@ from uuid import UUID
 from chalk import DataFrame
 from chalk.client.models import (
     BulkOnlineQueryResponse,
+    BulkOnlineQueryResult,
     FeatureReference,
     FeatureStatisticsResponse,
     OfflineQueryInputUri,
@@ -27,6 +28,7 @@ from chalk.prompts import Prompt
 if TYPE_CHECKING:
     import pandas as pd
     import polars as pl
+    import pyarrow as pa
 
     QueryInput = Mapping[FeatureReference, Any] | pd.DataFrame | pl.DataFrame | DataFrame | str
 
@@ -553,6 +555,62 @@ class AsyncChalkClient:
         """
         ...
 
+    async def _run_serialized_query(
+        self,
+        serialized_plan_bytes: bytes,
+        input: Union[Mapping[FeatureReference, Sequence[Any]], pa.Table],
+        output: Sequence[FeatureReference] = (),
+        staleness: Mapping[FeatureReference, str] | None = None,
+        context: OnlineQueryContext | None = None,
+        query_name: str | None = None,
+        query_name_version: str | None = None,
+        correlation_id: str | None = None,
+        include_meta: bool = False,
+        explain: bool = False,
+        store_plan_stages: bool = False,
+        meta: Mapping[str, str] | None = None,
+        headers: Mapping[str, str] | None = None,
+    ) -> BulkOnlineQueryResult:
+        """Run a query using a pre-serialized plan.
+
+        This is a protected method for internal use and testing.
+
+        Parameters
+        ----------
+        serialized_plan_bytes
+            The serialized BatchPlan protobuf bytes
+        input
+            The input data, either as a mapping of features to values or as a PyArrow table
+        output
+            The output features to compute
+        staleness
+            Maximum staleness overrides for features
+        context
+            Query context including environment and tags
+        query_name
+            The name of the query
+        query_name_version
+            The version of the query
+        correlation_id
+            Correlation ID for logging
+        include_meta
+            Whether to include metadata in the response
+        explain
+            Whether to include explain output
+        store_plan_stages
+            Whether to store plan stages
+        meta
+            Customer metadata tags
+        headers
+            Additional headers to provide with the request
+
+        Returns
+        -------
+        OnlineQueryResult
+            The query result
+        """
+        ...
+
     async def offline_query(
         self,
         input: Union[QueryInput, OfflineQueryInputUri] | None = None,

{chalkpy-2.92.2 → chalkpy-2.92.3}/chalk/client/client_async_impl.py

@@ -18,6 +18,7 @@ from chalk.client.client_impl import ChalkAPIClientImpl, OnlineQueryResponseImpl
 from chalk.client.dataset import Dataset, DatasetImpl
 from chalk.client.models import (
     BulkOnlineQueryResponse,
+    BulkOnlineQueryResult,
     ChalkError,
     FeatureStatisticsResponse,
     PlanQueryResponse,
@@ -153,6 +154,16 @@ class AsyncChalkClientImpl(AsyncChalkClient):
             partial(self._client.plan_query, *args, **kwargs),
         )
 
+    async def _run_serialized_query(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> BulkOnlineQueryResult:
+        return await self._current_loop().run_in_executor(
+            self._executor,
+            partial(self._client._run_serialized_query, *args, **kwargs),
+        )
+
     async def offline_query(
         self,
         *args: Any,

{chalkpy-2.92.2 → chalkpy-2.92.3}/chalk/client/client_impl.py

@@ -743,6 +743,8 @@ class OnlineQueryResponseImpl(OnlineQueryResult):
         self.warnings = warnings
         self.meta = meta
 
+        print(self.data)
+
         for d in self.data:
             if d.value is not None:
                 try:
@@ -4361,6 +4363,186 @@ https://docs.chalk.ai/cli/apply
         )
         return resp
 
+    def _run_serialized_query(
+        self,
+        serialized_plan_bytes: bytes,
+        input: Union[Mapping[FeatureReference, Sequence[Any]], pa.Table],
+        output: Sequence[FeatureReference] = (),
+        staleness: Optional[Mapping[FeatureReference, str]] = None,
+        context: Optional[OnlineQueryContext] = None,
+        query_name: Optional[str] = None,
+        query_name_version: Optional[str] = None,
+        correlation_id: Optional[str] = None,
+        include_meta: bool = False,
+        explain: bool = False,
+        store_plan_stages: bool = False,
+        meta: Optional[Mapping[str, str]] = None,
+        headers: Mapping[str, str] | None = None,
+    ) -> BulkOnlineQueryResult:
+        """Run a query using a pre-serialized plan.
+
+        This is a protected method for internal use and testing.
+
+        Parameters
+        ----------
+        serialized_plan_bytes
+            The serialized BatchPlan protobuf bytes
+        input
+            The input data, either as a mapping of features to values or as a PyArrow table
+        output
+            The output features to compute
+        staleness
+            Maximum staleness overrides for features
+        context
+            Query context including environment and tags
+        query_name
+            The name of the query
+        query_name_version
+            The version of the query
+        correlation_id
+            Correlation ID for logging
+        include_meta
+            Whether to include metadata in the response
+        explain
+            Whether to include explain output
+        store_plan_stages
+            Whether to store plan stages
+        meta
+            Customer metadata tags
+        headers
+            Additional headers to provide with the request
+
+        Returns
+        -------
+        OnlineQueryResult
+            The query result
+        """
+        try:
+            import pyarrow as pa
+            import pyarrow.feather
+        except ImportError:
+            raise missing_dependency_exception("chalkpy[runtime]")
+
+        # Convert input to PyArrow table if needed
+        if isinstance(input, Mapping):
+            # Convert mapping to PyArrow table
+            table_dict = {}
+            for feat_ref, values in input.items():
+                feat_name = str(feat_ref)
+                # Ensure values is a list
+                if not isinstance(values, list):
+                    values = [values]
+                table_dict[feat_name] = values
+            input_table = pa.Table.from_pydict(table_dict)
+        else:
+            input_table = input
+
+        # Encode outputs
+        outputs_encoded = encode_outputs(output).string_outputs if output else []
+
+        # Encode staleness
+        staleness_encoded = {}
+        if staleness is not None:
+            for k, v in staleness.items():
+                if is_feature_set_class(k):
+                    for f in k.features:
+                        staleness_encoded[f.root_fqn] = v
+                else:
+                    staleness_encoded[ensure_feature(k).root_fqn] = v
+
+        # Create FeatherRequestHeader
+        from chalk.client.models import OnlineQueryContext as OQC
+        header_dict = {
+            "outputs": outputs_encoded,
+            "expression_outputs": [],
+            "staleness": staleness_encoded if staleness_encoded else None,
+            "context": (context or OQC()).dict(),
+            "include_meta": include_meta,
+            "explain": explain,
+            "correlation_id": correlation_id,
+            "query_name": query_name,
+            "query_name_version": query_name_version,
+            "meta": meta,
+            "store_plan_stages": store_plan_stages,
+        }
+        header_json = json.dumps(header_dict).encode("utf-8")
+
+        # Serialize the input table to feather format
+        feather_buffer = BytesIO()
+        pa.feather.write_feather(input_table, feather_buffer)
+        feather_bytes = feather_buffer.getvalue()
+
+        # Build the request body:
+        # 1. First 8 bytes: int64 (big-endian) - length of serialized plan
+        # 2. Next N bytes: serialized BatchPlan protobuf
+        # 3. Next 8 bytes: int64 (big-endian) - length of header JSON
+        # 4. Next M bytes: UTF-8 encoded JSON header (FeatherRequestHeader)
+        # 5. Next 8 bytes: int64 (big-endian) - length of feather data
+        # 6. Remaining bytes: feather-encoded input data
+        request_body = BytesIO()
+        request_body.write(len(serialized_plan_bytes).to_bytes(8, byteorder="big"))
+        request_body.write(serialized_plan_bytes)
+        request_body.write(len(header_json).to_bytes(8, byteorder="big"))
+        request_body.write(header_json)
+        request_body.write(len(feather_bytes).to_bytes(8, byteorder="big"))
+        request_body.write(feather_bytes)
+
+        # Make the HTTP request
+        response = self._request(
+            method="POST",
+            uri="/v1/query/run",
+            response=None,  # We'll handle the response manually
+            json=None,
+            data=request_body.getvalue(),
+            environment_override=None,
+            preview_deployment_id=...,
+            branch=...,
+            metadata_request=False,
+            extra_headers=headers,
+        )
+
+        if not isinstance(response, requests.Response):
+            raise TypeError("Expected requests.Response")
+
+        if response.status_code != 200:
+            raise RuntimeError(f"Request failed with status {response.status_code}: {response.text}")
+
+        # Deserialize the response
+        result = OnlineQueryResultFeather.deserialize(response.content)
+
+        # Convert feather bytes back to a dataframe
+        scalars_df = None
+        if result.scalar_data:
+            scalars_table = pa.feather.read_table(BytesIO(result.scalar_data))
+            scalars_df = pa_table_to_pl_df(scalars_table)
+
+        # Parse errors from JSON strings back to ChalkError objects
+        errors = []
+        if result.errors:
+            for error_json in result.errors:
+                try:
+                    error_dict = json.loads(error_json)
+                    errors.append(ChalkError(**error_dict))
+                except Exception:
+                    # If parsing fails, create a generic error
+                    errors.append(ChalkError.create(code=ErrorCode.PARSE_FAILED, message=str(error_json)))
+
+        # Parse meta if present
+        query_meta = None
+        if result.meta:
+            try:
+                query_meta = QueryMeta(**json.loads(result.meta))
+            except Exception:
+                pass
+
+        # Return as BulkOnlineQueryResult
+        return BulkOnlineQueryResult(
+            scalars_df=scalars_df,
+            groups_dfs=None,
+            errors=errors if errors else None,
+            meta=query_meta,
+        )
+
     def _to_value(self, x: FeatureResult):
         f: Feature = Feature.from_root_fqn(x.field)
 

chalkpy-2.92.2/chalk/_version.py

@@ -1 +0,0 @@
-__version__ = "2.92.2"