chalkpy 2.89.22__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/functions/__init__.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import datetime as dt
 import inspect
 from enum import Enum
-from typing import Any, Callable, Literal, Mapping, TypeVar, Union
+from typing import Any, Callable, Literal, Mapping, Optional, TypeVar, Union
 
 import pyarrow as pa
 
@@ -1286,6 +1286,28 @@ def recover(*vals: Any):
     return UnderscoreFunction("recover", *vals)
 
 
+def is_not_null(expr: Any):
+    """
+    Check if a value is not null.
+
+    Parameters
+    ----------
+    expr
+        The value to check for nullity.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class User:
+    ...    id: str
+    ...    nickname: str | None
+    ...    nickname_not_missing: bool = F.is_not_null(_.nickname)
+    """
+    return ~is_null(expr)
+
+
 def is_null(expr: Any):
     """
     Check if a value is null.
@@ -1293,7 +1315,7 @@ def is_null(expr: Any):
     Parameters
     ----------
     expr
-        The value to check if it is null.
+        The value to check for nullity.
 
     Examples
     --------
@@ -1613,6 +1635,69 @@ def sagemaker_predict(
     )
 
 
+def openai_complete(
+    api_key: Underscore | str,
+    prompt: Underscore | str,
+    model: Underscore | str,
+    max_tokens: Underscore | int,
+    temperature: Underscore | float,
+):
+    """
+    Makes a completion request to OpenAI's chat API and returns the response.
+
+    This is a blocking expression that calls OpenAI's API during feature computation.
+    The response includes the completion text along with token usage statistics.
+
+    Parameters
+    ----------
+    api_key
+        The OpenAI API key to use for authentication.
+    prompt
+        The prompt text to send to the model.
+    model
+        The OpenAI model to use (e.g., "gpt-4", "gpt-3.5-turbo").
+    max_tokens
+        The maximum number of tokens to generate in the completion.
+    temperature
+        The sampling temperature to use, between 0 and 2. Higher values make
+        output more random, lower values make it more deterministic.
+
+    Returns
+    -------
+    A struct containing:
+        - completion: The generated text response
+        - prompt_tokens: Number of tokens in the prompt
+        - completion_tokens: Number of tokens in the completion
+        - total_tokens: Total tokens used (prompt + completion)
+        - model: The model used for the completion
+        - finish_reason: Why the completion stopped (e.g., "stop", "length")
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class Document:
+    ...    id: str
+    ...    content: str
+    ...    summary: str = F.openai_complete(
+    ...        api_key="sk-...",
+    ...        prompt=_.content,
+    ...        model="gpt-4",
+    ...        max_tokens=100,
+    ...        temperature=0.7,
+    ...    ).completion
+    """
+    return UnderscoreFunction(
+        "openai_complete",
+        api_key,
+        prompt,
+        model,
+        max_tokens,
+        temperature,
+    )
+
+
 def json_value(expr: Underscore, path: Union[str, Underscore]):
     """
     Extract structured data from a JSON string feature using a JSONPath expression.
@@ -3458,6 +3543,29 @@ def concat(first: Underscore | Any, second: Underscore | Any):
     return UnderscoreFunction("concat", first, second)
 
 
+def array(*args: Underscore | Any):
+    """
+    Creates an array from the given values.
+
+    Parameters
+    ----------
+    args
+        The values to create the array from.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class Name:
+    ...     id: str
+    ...     first_name: str
+    ...     last_name: str
+    ...     name: list[str] = F.array(_.first_name, _.last_name)
+    """
+    return UnderscoreFunction("array_constructor", *args)
+
+
 def array_sort(expr: Underscore | Any, descending: bool = False):
     """
     Returns an array which has the sorted order of the input
@@ -3470,6 +3578,17 @@ def array_sort(expr: Underscore | Any, descending: bool = False):
         The array to sort
     descending
         Whether to sort the array in descending order. Defaults to False.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class LeaderBoard:
+    ...     id: str
+    ...     scores: list[int]
+    ...     sorted_scores_asc: list[int] = F.array_sort(_.scores)
+    ...     sorted_scores_desc: list[int] = F.array_sort(_.scores, descending=True)
     """
     if descending:
         return UnderscoreFunction("array_sort_desc", expr)
@@ -3484,6 +3603,16 @@ def array_stddev(expr: Underscore | Any):
     ----------
     expr
         The array to calculate the standard deviation
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class SensorData:
+    ...     id: str
+    ...     temperature_readings: list[float]
+    ...     temp_stddev: float = F.array_stddev(_.temperature_readings)
     """
     return UnderscoreFunction("array_stddev", expr, False)
 
@@ -3497,6 +3626,16 @@ def array_sample_stddev(expr: Underscore | Any):
     ----------
     expr
         The array to calculate the sample standard deviation
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class ExperimentResults:
+    ...     id: str
+    ...     sample_measurements: list[float]
+    ...     sample_stddev: float = F.array_sample_stddev(_.sample_measurements)
     """
     return UnderscoreFunction("array_stddev", expr, True)
 
@@ -3509,6 +3648,16 @@ def array_sum(expr: Underscore | Any):
     ----------
     expr
         The array to sum
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class Transaction:
+    ...     id: str
+    ...     line_items: list[float]
+    ...     total_amount: float = F.array_sum(_.line_items)
     """
     return UnderscoreFunction("array_sum", expr)
 
@@ -3521,6 +3670,16 @@ def array_average(expr: Underscore | Any):
     ----------
     expr
         The array to average
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class StudentGrades:
+    ...     id: str
+    ...     test_scores: list[float]
+    ...     average_score: float = F.array_average(_.test_scores)
     """
     return UnderscoreFunction("array_average", expr)
 
@@ -3533,6 +3692,16 @@ def array_median(expr: Underscore | Any):
     ----------
     expr
         The array to take the median of
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class HousingMarket:
+    ...     id: str
+    ...     property_prices: list[float]
+    ...     median_price: float = F.array_median(_.property_prices)
     """
     return UnderscoreFunction.with_f_dot_repr("array_median", expr)
 
@@ -3555,6 +3724,17 @@ def array_mode(expr: Underscore | Any, tiebreak: Literal["FIRST", "MAX", "MIN"]
             ``"MIN"`` will return 1, the min of the multimodes;
 
         Defaults to ``"FIRST"`` (the behavior of python's ``statistics.mode()``)
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class SurveyAnalysis:
+    ...     id: str
+    ...     responses: list[int]
+    ...     most_common_response: int = F.array_mode(_.responses)
+    ...     highest_mode: int = F.array_mode(_.responses, tiebreak="MAX")
     """
     int_mode = 0 if tiebreak == "FIRST" else 1 if tiebreak == "MAX" else 2 if tiebreak == "MIN" else -1
     if int_mode == -1:
@@ -3679,33 +3859,64 @@ def array_count_value(expr: Underscore, value: Union[str, Underscore]):
 
 
 def _underscore_lambda(
-    f: Callable[[Underscore], Underscore],
+    f: Callable[..., Underscore],
     *,
-    parameter_type: pa.DataType,
+    parameter_type: Optional[pa.DataType] = None,
+    parameter_types: Optional[list[pa.DataType]] = None,
 ) -> Underscore:
     """
     This is a utility function for constructing lambda expressions in underscore expressions.
-    Note: only accepts functions which take a single positional argument.
+    Accepts functions with any number of positional arguments.
+
+    The caller must specify the parameter type(s) for the callback.
 
-    The caller must specify the parameter type for the callback.
+    Parameters
+    ----------
+    f
+        A callable that takes one or more Underscore arguments and returns an Underscore
+    parameter_type
+        For backward compatibility: the type of the single parameter (if function has one parameter)
+    parameter_types
+        List of parameter types for each argument (if function has multiple parameters)
     """
 
-    # The function `f` will be called immediately, to produce an `UnderscoreLambda` object.
-    # It is passed a `lambda_parameter` value with a name derived from the provided function.
-    lambda_param_name = "param1"
+    if parameter_type is not None and parameter_types is not None:
+        raise ValueError("Cannot specify both parameter_type and parameter_types")
+
+    if parameter_type is not None:
+        param_types_list = [parameter_type]
+    elif parameter_types is not None:
+        param_types_list = parameter_types
+    else:
+        raise ValueError("Must specify either parameter_type or parameter_types")
+
     f_sig = inspect.signature(f)
     f_parameters = list(f_sig.parameters.keys())
-    if len(f_parameters) >= 1:
-        lambda_param_name = f_parameters[0]
 
-    # Construct a lambda parameter object, which will be passed into the provided callback.
-    lambda_param_underscore = UnderscoreFunction("lambda_parameter", lambda_param_name, parameter_type)
-    return UnderscoreFunction(
-        "lambda",
-        lambda_param_name,
-        parameter_type,
-        f(lambda_param_underscore),
-    )
+    if len(f_parameters) != len(param_types_list):
+        raise ValueError(f"Function has {len(f_parameters)} parameter(s) but {len(param_types_list)} type(s) provided")
+
+    if len(f_parameters) == 0:
+        raise ValueError("Function must have at least one parameter")
+
+    lambda_param_underscores = []
+
+    for i, (param_name, param_type) in enumerate(zip(f_parameters, param_types_list)):
+        if not param_name:
+            param_name = f"param{i + 1}"
+        lambda_param_underscore = UnderscoreFunction("lambda_parameter", param_name, param_type)
+        lambda_param_underscores.append(lambda_param_underscore)
+
+    result_expr = f(*lambda_param_underscores)
+
+    lambda_args = []
+    for param_name, param_type in zip(f_parameters, param_types_list):
+        lambda_args.append(param_name)
+        lambda_args.append(param_type)
+
+    lambda_args.append(result_expr)
+
+    return UnderscoreFunction("lambda", *lambda_args)
 
 
 def array_filter(
@@ -3736,7 +3947,6 @@ def array_filter(
     ...    id: str
     ...    recent_activities: list[float]
     ...    average_activity: float
-    ...
     ...    recent_high_value_activities: list[float] = F.array_filter(
     ...        _.recent_activities,
     ...        lambda amount: amount > _.average_activity,
@@ -3758,6 +3968,159 @@ def array_filter(
     )
 
 
+def array_transform(
+    arr: Underscore,
+    transform: Callable[[Underscore], Underscore],
+    item_type: Union[pa.DataType, type],
+) -> Underscore:
+    """
+    Applies a custom transform function to each element in an array, returning a new
+    array containing transformed items.
+
+    Parameters
+    ----------
+    arr
+        An array of values
+    transform
+        A Python function producing an underscore expression to be applied to each item
+        in the array.
+    item_type
+        The type of each item in the array. This must be set explicitly.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class Bookstore:
+    ...    id: str
+    ...    prices: list[float]
+    ...    store_discount: float
+    ...    final_price: list[float] = F.array_transform(
+    ...        _.prices,
+    ...        lambda amount: amount * _.store_discount,
+    ...        item_type=float,
+    ...    )
+    """
+
+    if not isinstance(item_type, pa.DataType):
+        item_type = rich_to_pyarrow(
+            item_type,
+            name="array_transform.item_type",
+            respect_nullability=False,
+        )
+
+    return UnderscoreFunction(
+        "array_transform",
+        arr,
+        _underscore_lambda(transform, parameter_type=item_type),
+    )
+
+
+def array_reduce(
+    arr: Underscore,
+    initial_value: Underscore | Any,
+    arr_item_type: Union[pa.DataType, type],
+    reduce: Callable[[Underscore, Underscore], Underscore],
+    accumulator_type: Optional[Union[pa.DataType, type]] = None,
+    output_func: Callable[[Underscore], Underscore] = lambda x: x,
+) -> Underscore:
+    """
+    Reduces an array to a single value by applying a function to each element
+    along with an accumulator.
+
+    Parameters
+    ----------
+    arr
+        An array of values
+    initial_value
+        The initial value for the accumulator
+    reduce
+        A function that takes (accumulator, item) and returns the new accumulator value
+    arr_item_type
+        Type of each item in the array
+    accumulator_type
+        The Optional type of the accumulator result. Typically inferred from initial_value.
+    output_func
+        Optional function to transform the final accumulator value
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class User:
+    ...    id: str
+    ...    scores: list[int]
+    ...    total_score: int = F.array_reduce(
+    ...        arr=_.scores,
+    ...        initial_value=0,
+    ...        arr_item_type=int,
+    ...        reduce=lambda acc, score: acc + score,
+    ...    )
+    """
+    accumulator_type_arrow: Optional[pa.DataType] = None
+    initial_value_type_arrow: Optional[pa.DataType] = None
+    arr_item_type_arrow: Optional[pa.DataType] = None
+
+    if accumulator_type is not None:
+        if not isinstance(accumulator_type, pa.DataType):
+            accumulator_type_arrow = rich_to_pyarrow(
+                accumulator_type,
+                name="array_reduce.accumulator_type",
+                respect_nullability=False,
+            )
+        else:
+            accumulator_type_arrow = accumulator_type
+
+    if initial_value is not None:
+        # Try to infer accumulator type from initial_value if not provided
+        if not isinstance(initial_value, Underscore):
+            # If initial_value is a pyarrow scalar, extract its type
+            try:
+                if isinstance(initial_value, pa.DataType):
+                    initial_value_type_arrow = initial_value
+                elif isinstance(initial_value, pa.Scalar):
+                    accumulator_type = initial_value.type
+                else:
+                    # Try to infer type from Python literal value
+                    inferred_scalar = pa.scalar(initial_value)
+                    accumulator_type = inferred_scalar.type
+            except (TypeError, pa.ArrowInvalid):
+                raise ValueError("Could not infer type of initial_value; please provide accumulator_type explicitly.")
+
+    if accumulator_type_arrow is None and initial_value_type_arrow is None:
+        raise ValueError("initial_value type could not be determined; please provide it explicitly.")
+
+    if initial_value_type_arrow is not None:
+        accumulator_type_arrow = initial_value_type_arrow
+
+    if arr_item_type is None:
+        raise ValueError("arr_item_type must be provided to array_reduce")
+
+    if not isinstance(arr_item_type, pa.DataType):
+        arr_item_type_arrow = rich_to_pyarrow(
+            arr_item_type,
+            name="array_reduce.arr_item_type",
+            respect_nullability=False,
+        )
+    else:
+        arr_item_type_arrow = arr_item_type
+
+    reduce_lambda_param_types = [accumulator_type_arrow, arr_item_type_arrow]
+
+    reduce_lambda = _underscore_lambda(reduce, parameter_types=reduce_lambda_param_types)
+    output_lambda = _underscore_lambda(output_func, parameter_type=accumulator_type_arrow)
+
+    return UnderscoreFunction(
+        "array_reduce",
+        arr,
+        initial_value,
+        reduce_lambda,
+        output_lambda,
+    )
+
+
 def array_max(arr: Underscore):
     """
     Returns the maximum value in an array.
@@ -4139,6 +4502,11 @@ def max(*values: Any):
     )
 
 
+max._chalk__method_chaining_predicate = (  # pyright: ignore[reportFunctionMemberAccess]
+    lambda underscore_call: len(underscore_call._chalk__args) > 0
+)
+
+
 def min(*values: Any):
     """
     Returns the minimum value in a list of values.
@@ -4170,6 +4538,11 @@ def min(*values: Any):
     )
 
 
+min._chalk__method_chaining_predicate = (  # pyright: ignore[reportFunctionMemberAccess]
+    lambda underscore_call: len(underscore_call._chalk__args) > 0
+)
+
+
 def jinja(template: str):
     """
     Runs a Jinja template on the input columns.
@@ -4359,7 +4732,7 @@ def nth_bucket_end(value: Underscore, bucket_duration: str, n: int, initial_buck
 
 
 def inference(
-    model: ModelVersion, inputs: list[Underscore | Any], resource_hint: ResourceHint | None = None
+    model: ModelVersion, inputs: list[Underscore | Any] | Underscore, resource_hint: ResourceHint | None = None
 ) -> Underscore | Feature:
     """
     Run inference on a deployed ML model.
@@ -5228,6 +5601,64 @@ def array_normalize(array: Underscore, p: Underscore | float | None = None):
     return UnderscoreFunction("array_normalize", array, 2.0 if p is None else p)
 
 
+def scale_vector(array: Underscore, p: Underscore | float):
+    """
+    Scales the input vector by the amount p.
+
+    Parameters
+    ----------
+    array
+        The input vector
+    p
+        The factor by which to scale
+
+    Returns
+    -------
+    Array where each element is multiplied by p.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class ProbabilityModel:
+    ...     id: str
+    ...     embedding: Vector[3]
+    ...     scale: float
+    ...     probabilities: Vector[3] = F.scale_vector(_.embedding, _.scale)
+    """
+    return UnderscoreFunction("scale_vector", array, p)
+
+
+def array_add(array1: Underscore, array2: Underscore):
+    """
+    Element-wise addition of two vectors.
+
+    Parameters
+    ----------
+    array1
+        The first input vector
+    array2
+        The second input vector
+
+    Returns
+    -------
+    Array where each element is the sum of corresponding elements from array1 and array2.
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class VectorModel:
+    ...     id: str
+    ...     vec1: Vector[3]
+    ...     vec2: Vector[3]
+    ...     sum_vec: Vector[3] = F.array_add(_.vec1, _.vec2)
+    """
+    return UnderscoreFunction("array_add", array1, array2)
+
+
 def array_position(array: Underscore, element: Underscore):
     """
     Find the position of an element in the array (1-based indexing).
@@ -5431,6 +5862,7 @@ __all__ = (
     "abs",
     "acos",
     "array_agg",
+    "array_add",
     "array_average",
     "array_count_value",
     "array_cum_sum",
@@ -5511,6 +5943,7 @@ __all__ = (
     "inference",
     "is_leap_year",
     "is_month_end",
+    "is_not_null",
     "is_null",
     "is_us_federal_holiday",
     "jaccard_similarity",
@@ -5545,6 +5978,7 @@ __all__ = (
     "normal_cdf",
     "nth_bucket_end",
     "nth_bucket_start",
+    "openai_complete",
     "parse_datetime",
     "partial_ratio",
     "pi",
@@ -5618,4 +6052,5 @@ __all__ = (
     "word_stem",
     "xgboost_regressor",
     "year",
+    "scale_vector",
 )

chalk/functions/holidays.py

@@ -3,9 +3,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from datetime import date
 from enum import Enum
-from typing import Literal, Optional, Union
-
-from typing_extensions import TypeAlias
+from typing import Literal, Optional, TypeAlias, Union
 
 
 class DayOfWeek(int, Enum):

chalk/gitignore/gitignore_parser.py

@@ -114,6 +114,10 @@ def _rule_from_pattern(pattern: str, base_path: Optional[Path] = None, source: O
     regex = _fnmatch_pathname_to_regex(pattern, directory_only)
     if anchored:
         regex = f"^{regex}"
+    else:
+        # For non-anchored patterns, match at path component boundaries
+        # (start of string or after a path separator)
+        regex = f"(^|/){regex}"
     regex = f"(?ms){regex}"
     return IgnoreRule(
         pattern=orig_pattern,
@@ -215,6 +219,6 @@ def _fnmatch_pathname_to_regex(pattern: str, directory_only: bool):
     if directory_only:
         res.append(r"/.*$")
     else:
-        res.append("(/.*|[^/]*)$")
+        res.append("(/.*)?$")
 
     return "".join(res)