|
| 1 | +// Copyright 2026 Query Farm LLC - https://query.farm |
| 2 | + |
| 3 | +package farm.query.vgi.example.table; |
| 4 | + |
| 5 | +import farm.query.vgi.function.FunctionMetadata; |
| 6 | +import farm.query.vgi.function.FunctionSpec; |
| 7 | +import farm.query.vgi.internal.BatchUtil; |
| 8 | +import farm.query.vgi.protocol.FunctionExample; |
| 9 | +import farm.query.vgi.table.SimpleTableFunction; |
| 10 | +import farm.query.vgi.table.TableBindParams; |
| 11 | +import farm.query.vgi.table.TableInitParams; |
| 12 | +import farm.query.vgi.table.TableProducerState; |
| 13 | +import farm.query.vgi.types.Schemas; |
| 14 | +import farm.query.vgirpc.CallContext; |
| 15 | +import farm.query.vgirpc.OutputCollector; |
| 16 | +import org.apache.arrow.vector.BigIntVector; |
| 17 | +import org.apache.arrow.vector.VarCharVector; |
| 18 | +import org.apache.arrow.vector.types.pojo.Schema; |
| 19 | +import org.apache.arrow.vector.util.Text; |
| 20 | + |
| 21 | +import java.util.List; |
| 22 | + |
| 23 | +/** |
| 24 | + * {@code global_table() -> (n int64, label utf8)} — the table member of the |
| 25 | + * global-registration probe family (see {@code scalar/GlobalScalarFunction} for |
| 26 | + * the family's rationale). |
| 27 | + * |
| 28 | + * <p>A fixed generator taking no arguments: three labelled rows, emitted once. |
| 29 | + * Mirrors vgi-python's {@code GlobalTableFunction}. |
| 30 | + */ |
| 31 | +public final class GlobalTableFunction extends SimpleTableFunction { |
| 32 | + |
| 33 | + private static final Schema OUTPUT = Schemas.of( |
| 34 | + Schemas.nullable("n", Schemas.INT64), |
| 35 | + Schemas.nullable("label", Schemas.UTF8)); |
| 36 | + |
| 37 | + /** Exactly three rows, always. */ |
| 38 | + private static final long ROWS = 3L; |
| 39 | + |
| 40 | + private static final FunctionSpec SPEC = FunctionSpec.builder("global_table") |
| 41 | + .metadata(FunctionMetadata.describe("Global-registration probe (table)") |
| 42 | + .withCategories("test", "global") |
| 43 | + .withExamples(List.of(new FunctionExample( |
| 44 | + "SELECT * FROM vgi_example_global_table()", |
| 45 | + "Table probe published into system.main", null)))) |
| 46 | + .build(); |
| 47 | + |
| 48 | + @Override public FunctionSpec spec() { return SPEC; } |
| 49 | + |
| 50 | + @Override protected Schema outputSchema() { return OUTPUT; } |
| 51 | + |
| 52 | + @Override public long cardinality(TableBindParams params) { return ROWS; } |
| 53 | + |
| 54 | + @Override public TableProducerState createProducer(TableInitParams params) { |
| 55 | + return new State(); |
| 56 | + } |
| 57 | + |
| 58 | + /** One-shot emit latch for the three probe rows. */ |
| 59 | + public static final class State extends TableProducerState { |
| 60 | + |
| 61 | + /** Whether the single batch has been emitted. */ |
| 62 | + public boolean emitted; |
| 63 | + |
| 64 | + /** Required no-arg constructor for state deserialization. */ |
| 65 | + public State() {} |
| 66 | + |
| 67 | + @Override public void produceTick(OutputCollector out, CallContext ctx) { |
| 68 | + if (emitted) { |
| 69 | + out.finish(); |
| 70 | + return; |
| 71 | + } |
| 72 | + emitted = true; |
| 73 | + BatchUtil.emit(OUTPUT, (int) ROWS, out, (root, rows, ignored) -> { |
| 74 | + BigIntVector n = (BigIntVector) root.getVector("n"); |
| 75 | + VarCharVector label = (VarCharVector) root.getVector("label"); |
| 76 | + for (int i = 0; i < rows; i++) { |
| 77 | + n.setSafe(i, i); |
| 78 | + label.setSafe(i, new Text("global_table:" + i)); |
| 79 | + } |
| 80 | + }); |
| 81 | + } |
| 82 | + } |
| 83 | +} |
0 commit comments