vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/scalar/geo.py

@@ -0,0 +1,300 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Geospatial scalar fixtures (geo_distance_*, geo_centroid_*)."""
+
+from __future__ import annotations
+
+from typing import Annotated, Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from vgi.arguments import Param, Returns
+from vgi.metadata import FunctionExample
+from vgi.scalar_function import ScalarFunction
+
+_POINT_STRUCT_TYPE = pa.struct([("lat", pa.float64()), ("lon", pa.float64())])
+
+
+def _euclidean_distance(
+    lat1: pa.Array[Any], lon1: pa.Array[Any], lat2: pa.Array[Any], lon2: pa.Array[Any]
+) -> pa.DoubleArray:
+    """Compute Euclidean distance: sqrt((lat2-lat1)^2 + (lon2-lon1)^2)."""
+    dlat = pc.subtract(lat2, lat1)
+    dlon = pc.subtract(lon2, lon1)
+    return pc.sqrt(pc.add(pc.multiply(dlat, dlat), pc.multiply(dlon, dlon)))  # type: ignore[return-value]
+
+
+def _compute_centroid(lat_arrays: list[pa.Array[Any]], lon_arrays: list[pa.Array[Any]]) -> pa.StructArray:
+    """Compute centroid (average lat, average lon) from parallel lat/lon arrays."""
+    n = len(lat_arrays)
+    lat_sum: pa.Array[Any] = lat_arrays[0]
+    lon_sum: pa.Array[Any] = lon_arrays[0]
+    for i in range(1, n):
+        lat_sum = pc.add(lat_sum, lat_arrays[i])
+        lon_sum = pc.add(lon_sum, lon_arrays[i])
+    divisor = pa.scalar(n, type=pa.float64())
+    avg_lat = pc.divide(lat_sum, divisor)
+    avg_lon = pc.divide(lon_sum, divisor)
+    return pa.StructArray.from_arrays([avg_lat, avg_lon], names=["lat", "lon"])
+
+
+class GeoDistanceStructFunction(ScalarFunction):
+    """Euclidean distance between two struct points.
+
+    Each point is a struct with lat and lon fields.
+
+    Example:
+        SQL:    SELECT geo_distance_struct(p1, p2) FROM points
+        Input:  p1={lat: 0.0, lon: 0.0}, p2={lat: 3.0, lon: 4.0}
+        Output: result=5.0
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_distance_struct"
+        description = "Euclidean distance between two struct points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_distance_struct({lat: 0, lon: 0}, {lat: 3, lon: 4})",
+                description="Distance between origin and (3, 4)",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        p1: Annotated[
+            pa.StructArray,
+            Param(doc="First point {lat, lon}", arrow_type=_POINT_STRUCT_TYPE),
+        ],
+        p2: Annotated[
+            pa.StructArray,
+            Param(doc="Second point {lat, lon}", arrow_type=_POINT_STRUCT_TYPE),
+        ],
+    ) -> Annotated[pa.DoubleArray, Returns()]:
+        """Compute Euclidean distance between two points."""
+        return _euclidean_distance(p1.field("lat"), p1.field("lon"), p2.field("lat"), p2.field("lon"))
+
+
+class GeoDistanceListFunction(ScalarFunction):
+    """Euclidean distance between two list points.
+
+    Each point is a list of two float64 values [lat, lon].
+
+    Example:
+        SQL:    SELECT geo_distance_list(p1, p2) FROM points
+        Input:  p1=[0.0, 0.0], p2=[3.0, 4.0]
+        Output: result=5.0
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_distance_list"
+        description = "Euclidean distance between two list points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_distance_list([0, 0], [3, 4])",
+                description="Distance between origin and (3, 4)",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        p1: Annotated[  # type: ignore[type-arg]
+            pa.ListArray,
+            Param(doc="First point [lat, lon]", arrow_type=pa.list_(pa.float64())),
+        ],
+        p2: Annotated[  # type: ignore[type-arg]
+            pa.ListArray,
+            Param(doc="Second point [lat, lon]", arrow_type=pa.list_(pa.float64())),
+        ],
+    ) -> Annotated[pa.DoubleArray, Returns()]:
+        """Compute Euclidean distance between two points."""
+        return _euclidean_distance(
+            pc.list_element(p1, 0),
+            pc.list_element(p1, 1),
+            pc.list_element(p2, 0),
+            pc.list_element(p2, 1),
+        )
+
+
+class GeoDistanceFixedFunction(ScalarFunction):
+    """Euclidean distance between two fixed-size list points.
+
+    Each point is a fixed-size list of two float64 values [lat, lon].
+
+    Example:
+        SQL:    SELECT geo_distance_fixed(p1, p2) FROM points
+        Input:  p1=[0.0, 0.0], p2=[3.0, 4.0]
+        Output: result=5.0
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_distance_fixed"
+        description = "Euclidean distance between two fixed-size list points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_distance_fixed([0, 0], [3, 4])",
+                description="Distance between origin and (3, 4)",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        p1: Annotated[  # type: ignore[type-arg]
+            pa.FixedSizeListArray,
+            Param(doc="First point [lat, lon]", arrow_type=pa.list_(pa.float64(), 2)),
+        ],
+        p2: Annotated[  # type: ignore[type-arg]
+            pa.FixedSizeListArray,
+            Param(doc="Second point [lat, lon]", arrow_type=pa.list_(pa.float64(), 2)),
+        ],
+    ) -> Annotated[pa.DoubleArray, Returns()]:
+        """Compute Euclidean distance between two points."""
+        return _euclidean_distance(
+            pc.list_element(p1, 0),
+            pc.list_element(p1, 1),
+            pc.list_element(p2, 0),
+            pc.list_element(p2, 1),
+        )
+
+
+class GeoCentroidStructFunction(ScalarFunction):
+    """Centroid of N struct points (varargs).
+
+    Computes the average lat and average lon across all input point columns.
+
+    Example:
+        SQL:    SELECT geo_centroid_struct(p1, p2) FROM points
+        Input:  p1={lat: 0.0, lon: 0.0}, p2={lat: 4.0, lon: 6.0}
+        Output: result={lat: 2.0, lon: 3.0}
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_centroid_struct"
+        description = "Centroid of N struct points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_centroid_struct(p1, p2) FROM points",
+                description="Compute centroid of two struct points",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        points: Annotated[
+            list[pa.StructArray],
+            Param(
+                doc="Point columns {lat, lon}",
+                arrow_type=_POINT_STRUCT_TYPE,
+                varargs=True,
+            ),
+        ],
+    ) -> Annotated[pa.StructArray, Returns(arrow_type=_POINT_STRUCT_TYPE)]:
+        """Compute centroid of all points."""
+        return _compute_centroid(
+            [p.field("lat") for p in points],
+            [p.field("lon") for p in points],
+        )
+
+
+class GeoCentroidListFunction(ScalarFunction):
+    """Centroid of N list points (varargs).
+
+    Computes the average lat and average lon across all input point columns,
+    where each point is a list of [lat, lon].
+
+    Example:
+        SQL:    SELECT geo_centroid_list(p1, p2) FROM points
+        Input:  p1=[0.0, 0.0], p2=[4.0, 6.0]
+        Output: result={lat: 2.0, lon: 3.0}
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_centroid_list"
+        description = "Centroid of N list points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_centroid_list(p1, p2) FROM points",
+                description="Compute centroid of two list points",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        points: Annotated[  # type: ignore[type-arg]
+            list[pa.ListArray],
+            Param(
+                doc="Point columns [lat, lon]",
+                arrow_type=pa.list_(pa.float64()),
+                varargs=True,
+            ),
+        ],
+    ) -> Annotated[pa.StructArray, Returns(arrow_type=_POINT_STRUCT_TYPE)]:
+        """Compute centroid of all points."""
+        return _compute_centroid(
+            [pc.list_element(p, 0) for p in points],
+            [pc.list_element(p, 1) for p in points],
+        )
+
+
+class GeoCentroidFixedFunction(ScalarFunction):
+    """Centroid of N fixed-size list points (varargs).
+
+    Computes the average lat and average lon across all input point columns,
+    where each point is a fixed-size list of [lat, lon].
+
+    Example:
+        SQL:    SELECT geo_centroid_fixed(p1, p2) FROM points
+        Input:  p1=[0.0, 0.0], p2=[4.0, 6.0]
+        Output: result={lat: 2.0, lon: 3.0}
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "geo_centroid_fixed"
+        description = "Centroid of N fixed-size list points"
+        examples = [
+            FunctionExample(
+                sql="SELECT geo_centroid_fixed(p1, p2) FROM points",
+                description="Compute centroid of two fixed-size list points",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        points: Annotated[  # type: ignore[type-arg]
+            list[pa.FixedSizeListArray],
+            Param(
+                doc="Point columns [lat, lon]",
+                arrow_type=pa.list_(pa.float64(), 2),
+                varargs=True,
+            ),
+        ],
+    ) -> Annotated[pa.StructArray, Returns(arrow_type=_POINT_STRUCT_TYPE)]:
+        """Compute centroid of all points."""
+        return _compute_centroid(
+            [pc.list_element(p, 0) for p in points],
+            [pc.list_element(p, 1) for p in points],
+        )

vgi/_test_fixtures/scalar/null_handling.py

@@ -0,0 +1,107 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Null-handling and conditional-message scalar fixtures."""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from vgi.arguments import ConstParam, Param, Returns
+from vgi.metadata import FunctionExample, NullHandling
+from vgi.scalar_function import ScalarFunction
+
+
+class ConditionalMessageFunction(ScalarFunction):
+    """Returns a repeated message when condition is true, empty string otherwise.
+
+    This example demonstrates multiple ConstParam parameters:
+    - repeat_count (int): How many times to repeat the message
+    - message (string): The message to repeat
+    - condition (boolean column): Whether to apply the message
+
+    The constant parameters come first, followed by the column parameter.
+
+    Example:
+        SQL:    SELECT conditional_message(3, 'Hi! ', is_active) FROM users
+        Input:  is_active=[true, false, true]
+        Args:   repeat_count=3, message='Hi! '
+        Output: result=['Hi! Hi! Hi! ', '', 'Hi! Hi! Hi! ']
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "conditional_message"
+        description = "Returns repeated message when condition is true"
+        examples = [
+            FunctionExample(
+                sql="SELECT conditional_message(3, 'Alert! ', flag) FROM items",
+                description="Show alert message for flagged items",
+            ),
+            FunctionExample(
+                sql="SELECT conditional_message(2, '⭐', is_featured) FROM products",
+                description="Add stars to featured products",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        repeat_count: Annotated[int, ConstParam("Number of times to repeat")],
+        message: Annotated[str, ConstParam("Message to repeat")],
+        condition: Annotated[pa.BooleanArray, Param(doc="Apply message condition")],
+    ) -> Annotated[pa.StringArray, Returns()]:
+        """Return repeated message when condition is true, empty string otherwise."""
+        repeated_message = message * repeat_count
+        result: pa.StringArray = pc.if_else(condition, repeated_message, "")  # type: ignore[assignment]
+        return result
+
+
+# Type for config struct: {label: string, version: int64}
+_CONFIG_STRUCT_TYPE = pa.struct([("label", pa.string()), ("version", pa.int64())])
+
+
+class NullHandlingFunction(ScalarFunction):
+    """Demonstrates special null handling in a scalar function.
+
+    This function returns the input value if it's not null, or -5000 if it is null.
+    It demonstrates how to use NullHandling.SPECIAL to receive null values
+    instead of having them automatically converted to null output.
+
+    This example uses type inference with pa.Int64Array and Meta.null_handling.
+
+    Example:
+        SQL:    SELECT null_handling(value) FROM data
+        Input:  value=[1, None, 3]
+        Output: result=[1, -5000, 3]
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "null_handling"
+        description = "Returns value or -5000 if null"
+        null_handling = NullHandling.SPECIAL
+        examples = [
+            FunctionExample(
+                sql="SELECT null_handling(value) FROM data",
+                description="Replace null values with -5000",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        value: Annotated[pa.Int64Array, Param(doc="Integer value to process")],
+    ) -> Annotated[pa.Int64Array, Returns()]:
+        """Return value if not null, otherwise -5000."""
+        # Use if_else: if value is null, return -5000, otherwise return the value
+        result: pa.Int64Array = pc.if_else(  # type: ignore[assignment]
+            pc.is_null(value), pa.scalar(-5000, type=pa.int64()), value
+        )
+        return result

vgi/_test_fixtures/scalar/random_demo.py

@@ -0,0 +1,171 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Random/seeded scalar fixtures (random_int, random_bytes, bernoulli, hash_seed)."""
+
+from __future__ import annotations
+
+from typing import Annotated
+
+import pyarrow as pa
+
+from vgi.arguments import ConstParam, OutputLength, Param, Returns
+from vgi.metadata import FunctionExample, FunctionStability
+from vgi.scalar_function import ScalarFunction
+
+
+class RandomIntFunction(ScalarFunction):
+    """Generates random integers for each row (demonstrates VOLATILE stability).
+
+    This function demonstrates FunctionStability.VOLATILE - calling it twice
+    with the same input will produce different results. The database optimizer
+    cannot cache or reuse results from volatile functions.
+
+    This example uses type inference with pa.Int64Array and Meta.stability.
+
+    Other stability options:
+    - CONSISTENT: Same input always produces same output (deterministic)
+    - CONSISTENT_WITHIN_QUERY: Same within a query, may vary across queries
+
+    Example:
+        SQL:    SELECT random_int(min_col, max_col) FROM data
+        Input:  min_col=[1, 10, 100], max_col=[10, 100, 1000]
+        Output: result=[7, 55, 823]  (random values per row, different each time)
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "random_int"
+        description = "Generate random integers (demonstrates VOLATILE stability)"
+        stability = FunctionStability.VOLATILE
+        examples = [
+            FunctionExample(
+                sql="SELECT random_int(min_col, max_col) FROM data",
+                description="Generate random integers between min and max values",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        min_val: Annotated[pa.Int64Array, Param(doc="Minimum value (inclusive)")],
+        max_val: Annotated[pa.Int64Array, Param(doc="Maximum value (inclusive)")],
+    ) -> Annotated[pa.Int64Array, Returns()]:
+        """Generate random integers for each row."""
+        import numpy as np
+
+        # Use np.random.default_rng().integers(..., endpoint=True) so we don't
+        # have to add 1 to max_val (which overflows int64 when max_val is
+        # INT64_MAX, wrapping to a negative value and triggering "high <= 0").
+        rng = np.random.default_rng()
+        result = rng.integers(min_val.to_numpy(), max_val.to_numpy(), endpoint=True)
+        return pa.array(result, type=pa.int64())
+
+
+class BernoulliFunction(ScalarFunction):
+    """Generates random booleans for each row (demonstrates VOLATILE stability).
+
+    This function demonstrates how to generate output without any input parameters.
+    It will produce a random 0 or 1 for each row in the output.
+
+    Example:
+        SQL:    SELECT bernoulli() FROM data
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "bernoulli"
+        description = "Generate random booleans (demonstrates VOLATILE stability)"
+        stability = FunctionStability.VOLATILE
+        examples = [
+            FunctionExample(
+                sql="SELECT bernoulli() FROM data",
+                description="Generate samples from the bernoulli distribution",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        _length: Annotated[int, OutputLength()],
+    ) -> Annotated[pa.BooleanArray, Returns()]:
+        """Generate random booleans for each row."""
+        import random
+
+        values = [bool(random.randint(0, 1)) for _ in range(_length)]
+        return pa.array(values, type=pa.bool_())
+
+
+class HashSeedFunction(ScalarFunction):
+    """Generates deterministic integers from a constant seed.
+
+    Demonstrates the single-ConstParam pattern: one constant argument
+    folded at plan time, no column parameters.
+
+    Example:
+        SQL:    SELECT hash_seed(42) FROM data
+        Input:  (no column input)
+        Args:   seed=42
+        Output: result=[42, 43, 44, ...]  (seed + row_index)
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "hash_seed"
+        description = "Generate deterministic integers from a constant seed"
+        stability = FunctionStability.CONSISTENT
+        examples = [
+            FunctionExample(
+                sql="SELECT hash_seed(42) FROM data",
+                description="Generate deterministic integers seeded at 42",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        seed: Annotated[int, ConstParam("Seed value")],
+        _length: Annotated[int, OutputLength()],
+    ) -> Annotated[pa.Int64Array, Returns()]:
+        """Generate deterministic integers: seed + row_index for each row."""
+        return pa.array([seed + i for i in range(_length)], type=pa.int64())
+
+
+class RandomBytesFunction(ScalarFunction):
+    """Generates deterministic pseudo-random binary blobs from a seed."""
+
+    class Meta:
+        """Function metadata."""
+
+        name = "random_bytes"
+        description = "Generate pseudo-random binary blobs from seed and length"
+        stability = FunctionStability.CONSISTENT
+        examples = [
+            FunctionExample(
+                sql="SELECT random_bytes(42, 16) FROM data",
+                description="Generate a deterministic 16-byte blob per input row",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        seed: Annotated[int, ConstParam("Seed for pseudo-random byte generation")],
+        byte_length: Annotated[int, ConstParam("Output blob length in bytes")],
+        _length: Annotated[int, OutputLength()],
+    ) -> Annotated[pa.BinaryArray, Returns()]:
+        """Generate pseudo-random binary blobs for each row."""
+        import random
+
+        if byte_length < 0:
+            raise ValueError("byte_length must be >= 0")
+        rng = random.Random(seed)
+        return pa.array(
+            [bytes(rng.getrandbits(8) for _ in range(byte_length)) for _ in range(_length)],
+            type=pa.binary(),
+        )

vgi/_test_fixtures/scalar/settings_secrets.py

@@ -0,0 +1,102 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Setting/secret/auth-aware scalar fixtures (multiply_by_setting, return_secret_value, who_am_i)."""
+
+from __future__ import annotations
+
+import json
+from typing import Annotated, Any
+
+import pyarrow as pa
+import pyarrow.compute as pc
+
+from vgi.arguments import Auth, OutputLength, Param, Returns, Secret, Setting
+from vgi.auth import AuthContext
+from vgi.metadata import FunctionExample
+from vgi.scalar_function import ScalarFunction
+
+
+class MultiplyBySettingFunction(ScalarFunction):
+    """Generates the input value multiplied by a setting."""
+
+    class Meta:
+        """Function metadata."""
+
+        name = "multiply_by_setting"
+        description = "Multiply the input value by a setting value"
+        examples = [
+            FunctionExample(
+                sql="SELECT multiply_by_setting(5)",
+                description="Multiply the input value by a setting's value",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        value: Annotated[pa.Int64Array, Param(doc="Integer value to multiply")],
+        multiplier: Annotated[pa.Scalar[Any] | None, Setting()],
+    ) -> Annotated[pa.Int64Array, Returns()]:
+        """Generate the result for each row."""
+        assert multiplier is not None
+        return pc.multiply(multiplier, value)
+
+
+class ReturnSecretValueFunction(ScalarFunction):
+    """Return the value of a secret.
+
+    Example:
+        SQL:    SELECT return_secret_value()
+
+    """
+
+    class Meta:
+        """Function metadata."""
+
+        name = "return_secret_value"
+        description = "Return a secret's value"
+        examples = [
+            FunctionExample(
+                sql="SELECT return_secret_value()",
+                description="Return a secret's value",
+            ),
+        ]
+
+    @classmethod
+    def compute(
+        cls,
+        vgi_example_secret: Annotated[dict[str, pa.Scalar[Any]], Secret("vgi_example")],
+        _length: Annotated[int, OutputLength()],
+    ) -> Annotated[pa.StringArray, Returns()]:
+        """Generate the result for each row."""
+        # Convert pa.Scalar values to Python for JSON serialization
+        secret_dict = {k: v.as_py() for k, v in vgi_example_secret.items()}
+        return pa.array(
+            [json.dumps(secret_dict) for _ in range(_length)],
+            type=pa.string(),
+        )
+
+
+class WhoAmIFunction(ScalarFunction):
+    """Return the authenticated principal name.
+
+    Demonstrates the Auth annotation for accessing auth context in compute().
+    Over stdio transport (or when no auth is configured), returns "anonymous".
+
+    SQL: ``SELECT whoami(1)``
+    """
+
+    class Meta:
+        """Metadata for the whoami function."""
+
+        name = "whoami"
+
+    @classmethod
+    def compute(
+        cls,
+        x: Annotated[pa.Int64Array, Param(doc="dummy input")],
+        auth: Annotated[AuthContext, Auth()],
+    ) -> Annotated[pa.StringArray, Returns()]:
+        """Return the authenticated principal name."""
+        name = auth.principal or "anonymous"
+        return pa.array([name] * len(x))