|
51 | 51 | from vgi_rpc.log import Level |
52 | 52 | from vgi_rpc.rpc import OutputCollector |
53 | 53 |
|
54 | | -from vgi.arguments import Arg, Setting, TableInput |
| 54 | +from vgi.arguments import AnyArrow, Arg, Setting, TableInput |
55 | 55 | from vgi.cache_control import CacheControl |
56 | 56 | from vgi.invocation import BindResponse |
57 | 57 | from vgi.metadata import FunctionExample |
@@ -94,6 +94,8 @@ class _LogDrainState(ArrowSerializableDataclass): |
94 | 94 | "GeoEncode3Function", |
95 | 95 | "RowSumFunction", |
96 | 96 | "BlendedDropFunction", |
| 97 | + "BlendedAnyFunction", |
| 98 | + "BlendedAnyVarargsFunction", |
97 | 99 | "SubstreamPartialSumFunction", |
98 | 100 | "MultiBatchFinishFunction", |
99 | 101 | "SumAllColumnsFunction", |
@@ -1880,6 +1882,87 @@ def process( |
1880 | 1882 | cast("VgiOutputCollector", out).emit(out_batch, parent_rows=parent_rows) |
1881 | 1883 |
|
1882 | 1884 |
|
| 1885 | +@dataclass(slots=True, frozen=True, kw_only=True) |
| 1886 | +class _AnyArgs: |
| 1887 | + """One positional ANY-typed input column.""" |
| 1888 | + |
| 1889 | + value: Annotated[AnyArrow, Arg(0, doc="Input column of any type (echoed back unchanged)")] |
| 1890 | + |
| 1891 | + |
| 1892 | +class BlendedAnyFunction(RowTransformFunction[_AnyArgs]): |
| 1893 | + """Blended 1->1 echo of a single ANY-typed input column. |
| 1894 | +
|
| 1895 | + Proves a blended positional arg may be declared ANY: the declared type names |
| 1896 | + no concrete Arrow type, so the client must build the worker's input schema |
| 1897 | + from the type DuckDB actually resolved for the call (a STRUCT built per row, |
| 1898 | + a whole-row struct, a LIST, a plain VARCHAR ...) rather than from the |
| 1899 | + declaration. The output column ``value`` is bound to that resolved input |
| 1900 | + type at bind time and the column is passed through untouched, so a test |
| 1901 | + asserts on DuckDB's own ``typeof(value)`` and on the round-tripped values — |
| 1902 | + nothing here depends on how an SDK happens to render a type name. |
| 1903 | + """ |
| 1904 | + |
| 1905 | + class Meta: |
| 1906 | + name = "blended_any" |
| 1907 | + description = "Blended 1->1 echo of one ANY-typed input column (output typed from the input)" |
| 1908 | + categories = ["blended", "test"] |
| 1909 | + |
| 1910 | + @classmethod |
| 1911 | + def on_bind(cls, params: BindParams[_AnyArgs]) -> BindResponse: |
| 1912 | + input_schema = params.bind_call.input_schema |
| 1913 | + assert input_schema is not None |
| 1914 | + return BindResponse(output_schema=pa.schema([pa.field("value", input_schema.field(0).type)])) |
| 1915 | + |
| 1916 | + @classmethod |
| 1917 | + def process( |
| 1918 | + cls, |
| 1919 | + params: ProcessParams[_AnyArgs], |
| 1920 | + state: None, |
| 1921 | + batch: pa.RecordBatch, |
| 1922 | + out: OutputCollector, |
| 1923 | + ) -> None: |
| 1924 | + out.emit(pa.record_batch([batch.column(0)], schema=params.output_schema)) |
| 1925 | + |
| 1926 | + |
| 1927 | +@dataclass(slots=True, frozen=True, kw_only=True) |
| 1928 | +class _AnyVarargsArgs: |
| 1929 | + """N positional ANY-typed input columns, each free to be a different type.""" |
| 1930 | + |
| 1931 | + values: Annotated[list[AnyArrow], Arg(0, varargs=True, doc="Input columns of any types (echoed back)")] |
| 1932 | + |
| 1933 | + |
| 1934 | +class BlendedAnyVarargsFunction(RowTransformFunction[_AnyVarargsArgs]): |
| 1935 | + """Blended 1->1 echo of N ANY-typed VARARGS input columns. |
| 1936 | +
|
| 1937 | + The varargs counterpart of ``blended_any``: every runtime column may resolve |
| 1938 | + to a different concrete type, so the client must take each column's type |
| 1939 | + from the call rather than from the (ANY) vararg element type. Output columns |
| 1940 | + are ``col0..colN-1`` — the names the client generates for varargs blended |
| 1941 | + input — each bound to its own resolved input type. |
| 1942 | + """ |
| 1943 | + |
| 1944 | + class Meta: |
| 1945 | + name = "blended_any_varargs" |
| 1946 | + description = "Blended 1->1 echo of N ANY-typed varargs input columns (col0..colN-1)" |
| 1947 | + categories = ["blended", "test"] |
| 1948 | + |
| 1949 | + @classmethod |
| 1950 | + def on_bind(cls, params: BindParams[_AnyVarargsArgs]) -> BindResponse: |
| 1951 | + input_schema = params.bind_call.input_schema |
| 1952 | + assert input_schema is not None |
| 1953 | + return BindResponse(output_schema=pa.schema([pa.field(f"col{i}", f.type) for i, f in enumerate(input_schema)])) |
| 1954 | + |
| 1955 | + @classmethod |
| 1956 | + def process( |
| 1957 | + cls, |
| 1958 | + params: ProcessParams[_AnyVarargsArgs], |
| 1959 | + state: None, |
| 1960 | + batch: pa.RecordBatch, |
| 1961 | + out: OutputCollector, |
| 1962 | + ) -> None: |
| 1963 | + out.emit(pa.record_batch(cls.input_columns(batch), schema=params.output_schema)) |
| 1964 | + |
| 1965 | + |
1883 | 1966 | class CachedExplodeFunction(RowTransformFunction[_ExplodeArgs]): |
1884 | 1967 | """Cacheable blended 1->N fan-out advertising vgi.cache.per_value. |
1885 | 1968 |
|
|
0 commit comments