chalkpy 2.95.8__tar.gz → 2.96.0__tar.gz

{chalkpy-2.95.8 → chalkpy-2.96.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.95.8
+Version: 2.96.0
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai
@@ -279,6 +279,9 @@ Requires-Dist: python-json-logger<4.0.0,>=3.0.0; extra == "dynamodb"
 Requires-Dist: s3fs; extra == "dynamodb"
 Requires-Dist: sqlalchemy[asyncio]<2,>=1.4.26; extra == "dynamodb"
 Requires-Dist: sqlglot<21.2.0,>=19.0.0; extra == "dynamodb"
+Provides-Extra: tracing
+Requires-Dist: opentelemetry-api>=1.29.0; extra == "tracing"
+Requires-Dist: opentelemetry-sdk>=1.29.0; extra == "tracing"
 Provides-Extra: all
 Requires-Dist: PyAthena>=3.0.0; extra == "all"
 Requires-Dist: adlfs; extra == "all"
@@ -303,6 +306,8 @@ Requires-Dist: google-cloud-bigquery<4,>=3.25.0; extra == "all"
 Requires-Dist: google-cloud-storage; extra == "all"
 Requires-Dist: httpx<0.28.0; extra == "all"
 Requires-Dist: openai<1.53,>=1.3.2; extra == "all"
+Requires-Dist: opentelemetry-api>=1.29.0; extra == "all"
+Requires-Dist: opentelemetry-sdk>=1.29.0; extra == "all"
 Requires-Dist: packaging; extra == "all"
 Requires-Dist: polars[timezone]!=1.0,!=1.1,!=1.10,!=1.11,!=1.12,!=1.13,!=1.14,!=1.15,!=1.16,!=1.17,!=1.18,!=1.19,!=1.2,!=1.20,!=1.21,!=1.22,!=1.23,!=1.24,!=1.25,!=1.26,!=1.27,!=1.28,!=1.29,!=1.3,!=1.30,!=1.31,!=1.32,!=1.4,!=1.5,!=1.6,!=1.7,!=1.8,!=1.9,<1.33.1,>=0.17.2; extra == "all"
 Requires-Dist: psycopg2<3,>=2.9.4; extra == "all"

chalkpy-2.96.0/chalk/_version.py

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

{chalkpy-2.95.8 → chalkpy-2.96.0}/chalk/client/client_grpc.py

@@ -943,6 +943,10 @@ class ChalkGRPCClient:
             raise TypeError("One of `input` or `input_sql` is required")
         if input is not None and input_sql is not None:
             raise TypeError("`input` and `input_sql` are mutually exclusive")
+        if input_sql is not None and now is not None:
+            raise TypeError(
+                "When using `input_sql`, `now` is not allowed: instead, to provide a query time, you can have the SQL query output a column named `__ts__`"
+            )
 
         response, call = self._online_query_bulk_grpc_request(
             input=input,

{chalkpy-2.95.8 → chalkpy-2.96.0}/chalk/sql/_internal/integrations/redshift.py

@@ -4,11 +4,26 @@ import contextlib
 import io
 import logging
 import os
+import queue
 import threading
 import typing
 import uuid
-from concurrent.futures import Future, ThreadPoolExecutor
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Type, Union
+from concurrent.futures import ThreadPoolExecutor
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    NewType,
+    Optional,
+    Sequence,
+    Type,
+    Union,
+    cast,
+)
 
 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -27,9 +42,10 @@ from chalk.sql._internal.sql_source import (
 from chalk.sql.finalized_query import FinalizedChalkQuery
 from chalk.utils.df_utils import is_binary_like, read_parquet
 from chalk.utils.environment_parsing import env_var_bool
-from chalk.utils.log_with_context import get_logger
+from chalk.utils.log_with_context import LABELS_KEY, get_logger, get_logging_context
 from chalk.utils.missing_dependency import missing_dependency_exception
-from chalk.utils.threading import DEFAULT_IO_EXECUTOR
+from chalk.utils.threading import DEFAULT_IO_EXECUTOR, MultiSemaphore
+from chalk.utils.tracing import safe_incr, safe_set_gauge
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
@@ -40,6 +56,22 @@ if TYPE_CHECKING:
 _logger = get_logger(__name__)
 _public_logger = chalk.logging.chalk_logger
 
+_WorkerId = NewType("_WorkerId", int)
+
+
+def _get_resolver_tags() -> list[str] | None:
+    """Extract resolver_fqn from log context and return as tags list."""
+    try:
+        log_ctx = get_logging_context()
+        labels = log_ctx.get(LABELS_KEY, {})
+        resolver_fqn = labels.get("resolver_fqn")
+        if resolver_fqn:
+            return [f"resolver_fqn:{resolver_fqn}"]
+    except Exception:
+        # Don't fail if we can't get the resolver_fqn
+        pass
+    return None
+
 
 def get_supported_redshift_unload_types() -> List[Type["TypeEngine"]]:
     """
@@ -287,36 +319,235 @@ class RedshiftSourceImpl(BaseSQLSource):
             unload_destination=unload_destination,
             columns_to_features=columns_to_features,
             yield_empty_batches=query_execution_parameters.yield_empty_batches,
+            max_prefetch_size_bytes=query_execution_parameters.max_prefetch_size_bytes,
+            num_client_prefetch_threads=query_execution_parameters.num_client_prefetch_threads,
         )
 
+    def _download_worker(
+        self,
+        file_handles: queue.Queue[str],
+        sem: MultiSemaphore | None,
+        pa_table_queue: queue.Queue[tuple[pa.Table, int] | _WorkerId],
+        worker_idx: _WorkerId,
+        columns_to_features: Callable[[Sequence[str]], Mapping[str, Feature]],
+    ):
+        """Worker thread that downloads files from S3 with memory control."""
+        assert self._s3_bucket is not None
+        try:
+            while True:
+                try:
+                    filename = file_handles.get_nowait()
+                except queue.Empty:
+                    break
+
+                # Estimate file size from S3 metadata if possible
+                weight: int | None = None
+                try:
+                    with _boto_lock_ctx():
+                        head_response = self._s3_client.head_object(Bucket=self._s3_bucket, Key=filename)
+                        # Boto3 types indicate ContentLength is always int, but cast to Optional for defensive programming
+                        content_length = cast(int | None, head_response.get("ContentLength"))
+                        if content_length is not None:
+                            # Estimate uncompressed size (parquet typically compresses 3-5x)
+                            weight = content_length * 4
+                except Exception as e:
+                    _logger.warning(f"Failed to get file size for {filename}, will estimate after download", exc_info=e)
+
+                # Acquire semaphore before downloading
+                if sem and weight is not None:
+                    if weight > sem.initial_value:
+                        # If the file is larger than the maximum size, truncate to max
+                        weight = sem.initial_value
+                    if weight > 0:
+                        if not sem.acquire(weight):
+                            raise RuntimeError("Failed to acquire semaphore for redshift download")
+                        safe_set_gauge(
+                            "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                        )
+
+                # Download and convert to table
+                tbl = _download_file_to_table(self._s3_client, self._s3_bucket, filename, columns_to_features)
+
+                # If we didn't have a weight estimate, use actual table size
+                if weight is None:
+                    weight = tbl.nbytes
+                    if sem and weight is not None and weight > 0:
+                        if not sem.acquire(weight):
+                            raise RuntimeError("Failed to acquire semaphore for redshift download")
+                        safe_set_gauge(
+                            "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                        )
+
+                # Ensure weight is always an int
+                final_weight: int = weight if weight is not None else 0
+                pa_table_queue.put((tbl, final_weight))
+        finally:
+            # Signal that this worker is done
+            pa_table_queue.put(worker_idx)
+
     def _download_objs_async(
         self,
         unload_destination: str,
         columns_to_features: Callable[[Sequence[str]], Mapping[str, Feature]],
         yield_empty_batches: bool,
+        max_prefetch_size_bytes: int,
+        num_client_prefetch_threads: int,
     ) -> Iterable[pa.RecordBatch]:
+        """Download objects from S3 with byte-bounded memory control."""
         assert self._s3_bucket is not None
         filenames = list(_list_files(self._s3_client, self._s3_bucket, unload_destination))
-        download_futs: list[Future[pa.Table]] = []
+        _public_logger.info(
+            f"Downloading parquet data partitioned into {len(filenames)} files "
+            + f"(max_prefetch_bytes={max_prefetch_size_bytes}, threads={num_client_prefetch_threads})..."
+        )
+
+        if len(filenames) == 0:
+            if yield_empty_batches:
+                # Need to get schema somehow - return empty batch
+                # This matches the original behavior
+                schema: pa.Schema | None = None
+                if schema is not None:
+                    yield pa.RecordBatch.from_pydict({k: [] for k in schema.names}, schema)
+            return
+
+        # Set up queues and semaphore for memory control
+        file_handles: queue.Queue[str] = queue.Queue()
         for filename in filenames:
-            _logger.debug("Scheduling download of file %s", filename)
-            fut = self._executor.submit(
-                _download_file_to_table, self._s3_client, self._s3_bucket, filename, columns_to_features
+            file_handles.put(filename)
+
+        max_weight = max_prefetch_size_bytes if max_prefetch_size_bytes > 0 else None
+        pa_table_queue: queue.Queue[tuple[pa.Table, int] | _WorkerId] = queue.Queue()
+        sem = None if max_weight is None else MultiSemaphore(max_weight)
+
+        # Start download workers
+        futures = {
+            _WorkerId(i): self._executor.submit(
+                self._download_worker,
+                file_handles,
+                sem,
+                pa_table_queue,
+                _WorkerId(i),
+                columns_to_features,
             )
-            download_futs.append(fut)
-        _public_logger.info(f"Downloading parquet data partitioned into {len(download_futs)} files...")
+            for i in range(num_client_prefetch_threads)
+        }
+
         schema: pa.Schema | None = None
         yielded = False
-        for fut in download_futs:
-            tbl = fut.result()
-            if len(tbl) > 0:
-                yield tbl.combine_chunks().to_batches()[0]
-                yielded = True
-            if len(tbl) == 0 and schema is None:
+
+        # Process downloaded tables as they become available
+        while len(futures) > 0:
+            x = pa_table_queue.get()
+            if isinstance(x, int):
+                # Worker finished - remove from futures and check for errors
+                futures.pop(x).result()
+                continue
+
+            tbl, weight = x
+            if schema is None:
                 schema = tbl.schema
-        if not yielded and schema is not None:
+
+            try:
+                if len(tbl) > 0:
+                    yield tbl.combine_chunks().to_batches()[0]
+                    safe_incr("chalk.redshift.downloaded_bytes", tbl.nbytes or 0, tags=_get_resolver_tags())
+                    safe_incr("chalk.redshift.downloaded_rows", tbl.num_rows or 0, tags=_get_resolver_tags())
+                    yielded = True
+            finally:
+                # Release semaphore after yielding
+                if sem is not None and weight > 0:
+                    sem.release(weight)
+                    safe_set_gauge(
+                        "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                    )
+
+        if not yielded and yield_empty_batches and schema is not None:
             yield pa.RecordBatch.from_pydict({k: [] for k in schema.names}, schema)
 
+    def _download_worker_raw(
+        self,
+        file_handles: queue.Queue[str],
+        sem: MultiSemaphore | None,
+        pa_table_queue: queue.Queue[tuple[pa.Table, int] | _WorkerId],
+        worker_idx: _WorkerId,
+        expected_output_schema: pa.Schema,
+    ):
+        """Worker thread that downloads files from S3 with memory control for raw execution."""
+        import pyarrow.compute as pc
+
+        assert self._s3_bucket is not None
+        try:
+            while True:
+                try:
+                    filename = file_handles.get_nowait()
+                except queue.Empty:
+                    break
+
+                # Estimate file size from S3 metadata if possible
+                weight: int | None = None
+                try:
+                    with _boto_lock_ctx():
+                        head_response = self._s3_client.head_object(Bucket=self._s3_bucket, Key=filename)
+                        # Boto3 types indicate ContentLength is always int, but cast to Optional for defensive programming
+                        content_length = cast(int | None, head_response.get("ContentLength"))
+                        if content_length is not None:
+                            # Estimate uncompressed size (parquet typically compresses 3-5x)
+                            weight = content_length * 4
+                except Exception as e:
+                    _logger.warning(f"Failed to get file size for {filename}, will estimate after download", exc_info=e)
+
+                # Acquire semaphore before downloading
+                if sem and weight is not None:
+                    if weight > sem.initial_value:
+                        weight = sem.initial_value
+                    if weight > 0:
+                        if not sem.acquire(weight):
+                            raise RuntimeError("Failed to acquire semaphore for redshift download")
+                        safe_set_gauge(
+                            "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                        )
+
+                # Download parquet file
+                buffer = io.BytesIO()
+                with _boto_lock_ctx():
+                    self._s3_client.download_fileobj(Bucket=self._s3_bucket, Key=filename, Fileobj=buffer)
+                    buffer.seek(0)
+                    if env_var_bool("CHALK_REDSHIFT_POLARS_PARQUET"):
+                        tbl = read_parquet(buffer, use_pyarrow=False).to_arrow()
+                    else:
+                        tbl = pq.read_table(buffer)
+
+                # If we didn't have a weight estimate, use actual table size
+                if weight is None:
+                    weight = tbl.nbytes
+                    if sem and weight is not None and weight > 0:
+                        if not sem.acquire(weight):
+                            raise RuntimeError("Failed to acquire semaphore for redshift download")
+                        safe_set_gauge(
+                            "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                        )
+
+                # Map columns to expected schema
+                arrays: list[pa.Array] = []
+                for field in expected_output_schema:
+                    if field.name in tbl.column_names:
+                        col = tbl.column(field.name)
+                        # Cast to expected type if needed
+                        if col.type != field.type:
+                            col = pc.cast(col, field.type)
+                        arrays.append(col)
+                    else:
+                        # Column not found, create null array
+                        arrays.append(pa.nulls(len(tbl), field.type))
+
+                mapped_tbl = pa.Table.from_arrays(arrays, schema=expected_output_schema)
+                # Ensure weight is always an int
+                final_weight: int = weight if weight is not None else 0
+                pa_table_queue.put((mapped_tbl, final_weight))
+        finally:
+            # Signal that this worker is done
+            pa_table_queue.put(worker_idx)
+
     def execute_query_efficient_raw(
         self,
         finalized_query: FinalizedChalkQuery,
@@ -325,8 +556,6 @@ class RedshiftSourceImpl(BaseSQLSource):
         query_execution_parameters: QueryExecutionParameters,
     ) -> Iterable[pa.RecordBatch]:
         """Execute query efficiently for Redshift and return raw PyArrow RecordBatches."""
-        import pyarrow.compute as pc
-
         temp_query_id = id(finalized_query)
         _public_logger.debug(f"Executing RedShift query [{temp_query_id}]...")
 
@@ -387,47 +616,77 @@ class RedshiftSourceImpl(BaseSQLSource):
                     except Exception:
                         _logger.warning(f"Failed to drop temp table '{temp_table_name}'", exc_info=True)
 
-        # Download files and map to expected schema
+        # Download files with memory control
         assert unload_destination is not None
         assert self._s3_bucket is not None
         filenames = list(_list_files(self._s3_client, self._s3_bucket, unload_destination))
+        _public_logger.info(
+            f"Downloading {len(filenames)} parquet files from Redshift UNLOAD "
+            + f"(max_prefetch_bytes={query_execution_parameters.max_prefetch_size_bytes}, "
+            + f"threads={query_execution_parameters.num_client_prefetch_threads})..."
+        )
 
-        yielded = False
+        if len(filenames) == 0:
+            if query_execution_parameters.yield_empty_batches:
+                arrays = [pa.nulls(0, field.type) for field in expected_output_schema]
+                yield pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
+            return
+
+        # Set up queues and semaphore for memory control
+        file_handles: queue.Queue[str] = queue.Queue()
         for filename in filenames:
-            buffer = io.BytesIO()
-            with _boto_lock_ctx():
-                self._s3_client.download_fileobj(Bucket=self._s3_bucket, Key=filename, Fileobj=buffer)
-                buffer.seek(0)
-                if env_var_bool("CHALK_REDSHIFT_POLARS_PARQUET"):
-                    tbl = read_parquet(buffer, use_pyarrow=False).to_arrow()
-                else:
-                    tbl = pq.read_table(buffer)
+            file_handles.put(filename)
 
-            if len(tbl) == 0:
-                continue
+        max_weight = (
+            query_execution_parameters.max_prefetch_size_bytes
+            if query_execution_parameters.max_prefetch_size_bytes > 0
+            else None
+        )
+        pa_table_queue: queue.Queue[tuple[pa.Table, int] | _WorkerId] = queue.Queue()
+        sem = None if max_weight is None else MultiSemaphore(max_weight)
+
+        # Start download workers
+        futures = {
+            _WorkerId(i): self._executor.submit(
+                self._download_worker_raw,
+                file_handles,
+                sem,
+                pa_table_queue,
+                _WorkerId(i),
+                expected_output_schema,
+            )
+            for i in range(query_execution_parameters.num_client_prefetch_threads)
+        }
 
-            # Map columns to expected schema
-            arrays: list[pa.Array] = []
-            for field in expected_output_schema:
-                if field.name in tbl.column_names:
-                    col = tbl.column(field.name)
-                    # Cast to expected type if needed
-                    if col.type != field.type:
-                        col = pc.cast(col, field.type)
-                    arrays.append(col)
-                else:
-                    # Column not found, create null array
-                    arrays.append(pa.nulls(len(tbl), field.type))
+        yielded = False
 
-            batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
-            yield batch
-            yielded = True
+        # Process downloaded tables as they become available
+        while len(futures) > 0:
+            x = pa_table_queue.get()
+            if isinstance(x, int):
+                # Worker finished - remove from futures and check for errors
+                futures.pop(x).result()
+                continue
+
+            tbl, weight = x
+
+            try:
+                if len(tbl) > 0:
+                    yield tbl.to_batches()[0]
+                    safe_incr("chalk.redshift.downloaded_bytes", tbl.nbytes or 0, tags=_get_resolver_tags())
+                    safe_incr("chalk.redshift.downloaded_rows", tbl.num_rows or 0, tags=_get_resolver_tags())
+                    yielded = True
+            finally:
+                # Release semaphore after yielding
+                if sem is not None and weight > 0:
+                    sem.release(weight)
+                    safe_set_gauge(
+                        "chalk.redshift.remaining_prefetch_bytes", sem.get_value(), tags=_get_resolver_tags()
+                    )
 
         if not yielded and query_execution_parameters.yield_empty_batches:
-            # Create empty batch with expected schema
             arrays = [pa.nulls(0, field.type) for field in expected_output_schema]
-            batch = pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
-            yield batch
+            yield pa.RecordBatch.from_arrays(arrays, schema=expected_output_schema)
 
     @classmethod
     def register_sqlalchemy_compiler_overrides(cls):

{chalkpy-2.95.8 → chalkpy-2.96.0}/chalk/utils/tracing.py

@@ -222,8 +222,8 @@ else:
 if can_use_datadog_statsd:
     from datadog.dogstatsd.base import statsd
 
-    def safe_set_gauge(gauge: str, value: int | float):
-        statsd.gauge(gauge, value)
+    def safe_set_gauge(gauge: str, value: int | float, tags: list[str] | None = None):
+        statsd.gauge(gauge, value, tags=tags)
 
     def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):
         statsd.increment(counter, value, tags)
@@ -233,7 +233,7 @@ if can_use_datadog_statsd:
 
 else:
 
-    def safe_set_gauge(gauge: str, value: int | float):
+    def safe_set_gauge(gauge: str, value: int | float, tags: list[str] | None = None):
         pass
 
     def safe_incr(counter: str, value: int | float, tags: list[str] | None = None):

{chalkpy-2.95.8 → chalkpy-2.96.0}/pyproject.toml

@@ -433,6 +433,11 @@ dynamodb = [
     # [[[end]]]
 ]
 
+tracing = [
+    "opentelemetry-api>=1.29.0",
+    "opentelemetry-sdk>=1.29.0",
+]
+
 all = [
     # "All" includes everything except for dev dependencies
     # [[[cog
@@ -440,7 +445,7 @@ all = [
     # grps = tomli.loads(pathlib.Path("pyproject.toml").read_text())["project"]["optional-dependencies"]
     # final_deps = set()
     # for grp, deps in grps.items():
-    #    if grp in ("runtime", "openai", "cohere", "vertexai", "sql", "athena", "bigquery", "postgresql", "snowflake", "sqlite", "redshift", "mysql", "databricks", "trino", "spanner", "dynamodb", "clickhouse", "mssql"):
+    #    if grp in ("runtime", "openai", "cohere", "vertexai", "sql", "athena", "bigquery", "postgresql", "snowflake", "sqlite", "redshift", "mysql", "databricks", "trino", "spanner", "dynamodb", "clickhouse", "mssql", "tracing"):
     #        final_deps.update(deps)
     # for dep in sorted(final_deps):
     #     print(f'    "{dep}",')
@@ -468,6 +473,8 @@ all = [
     "google-cloud-storage",
     "httpx<0.28.0",
     "openai>=1.3.2,<1.53",
+    "opentelemetry-api>=1.29.0",
+    "opentelemetry-sdk>=1.29.0",
     "packaging",
     "polars[timezone]>=0.17.2,<1.33.1,!=1.0,!=1.1,!=1.2,!=1.3,!=1.4,!=1.5,!=1.6,!=1.7,!=1.8,!=1.9,!=1.10,!=1.11,!=1.12,!=1.13,!=1.14,!=1.15,!=1.16,!=1.17,!=1.18,!=1.19,!=1.20,!=1.21,!=1.22,!=1.23,!=1.24,!=1.25,!=1.26,!=1.27,!=1.28,!=1.29,!=1.30,!=1.31,!=1.32",
     "psycopg2>=2.9.4,<3",

chalkpy-2.95.8/chalk/_version.py

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