Support VGI protocol 1.5 scan function schemas

Author: rustyconover

vgi/src/main/java/farm/query/vgi/Worker.java

@@ -50,7 +50,7 @@ public final class Worker {
      *  and {@link #globalFunctionPrefix(String)}. A worker that opts out still
      *  carries the fields (empty list, empty prefix): the extension matches the
      *  response schema exactly. */
-    public static final String VGI_PROTOCOL_VERSION = "1.4.0";
+    public static final String VGI_PROTOCOL_VERSION = "1.5.0";
 
     private String catalogName = "vgi";
     private String catalogComment = "";
@@ -519,6 +519,30 @@ public String schemaOf(Object fn) {
         return h == null || h.schemaName() == null ? defaultSchema : h.schemaName();
     }
 
+    /**
+     * Resolve the schema containing a named table function. The table's schema
+     * wins when that function is registered there; otherwise a single
+     * unambiguous registration is returned. Native DuckDB functions and
+     * ambiguous names return {@code null}.
+     *
+     * @param functionName the function named by a scan result
+     * @param tableSchema the schema containing the table being resolved
+     * @param catalogName the auxiliary catalog owner, or {@code null} for this worker
+     * @return the authoritative function schema, or {@code null}
+     */
+    public String resolveTableFunctionSchema(
+            String functionName, String tableSchema, String catalogName) {
+        java.util.LinkedHashSet<String> homes = new java.util.LinkedHashSet<>();
+        for (TableFunction fn : tables) {
+            if (!fn.name().equals(functionName)
+                    || !java.util.Objects.equals(catalogOf(fn), catalogName)) continue;
+            String schema = schemaOf(fn);
+            if (schema.equalsIgnoreCase(tableSchema)) return schema;
+            homes.add(schema);
+        }
+        return homes.size() == 1 ? homes.iterator().next() : null;
+    }
+
     /**
      * The auxiliary catalog {@code fn} is declared in, or {@code null} when it
      * belongs to this worker's own catalog. Ownership is always explicit — a

vgi/src/main/java/farm/query/vgi/catalog/ScanBranch.java

@@ -36,6 +36,8 @@
  *                       no locations is rejected
  * @param formatOptions format branch only — reader options, which BECOME the
  *                       reader's named arguments. Empty for the other kinds
+ * @param schemaName    function branch only — schema containing the named VGI
+ *                       function; {@code null} for native/ambiguous functions
  */
 public record ScanBranch(
         String functionName,
@@ -48,7 +50,28 @@ public record ScanBranch(
         String sourceTable,
         String formatName,
         List<String> formatLocations,
-        Map<String, Object> formatOptions) {
+        Map<String, Object> formatOptions,
+        String schemaName) {
+
+    /**
+     * Source-compatible constructor for pre-1.5 callers. The worker resolves
+     * the function schema before serializing the branch.
+     */
+    public ScanBranch(
+            String functionName,
+            List<Object> positional,
+            Map<String, Object> named,
+            String branchFilter,
+            boolean writable,
+            String sourceCatalog,
+            String sourceSchema,
+            String sourceTable,
+            String formatName,
+            List<String> formatLocations,
+            Map<String, Object> formatOptions) {
+        this(functionName, positional, named, branchFilter, writable, sourceCatalog,
+                sourceSchema, sourceTable, formatName, formatLocations, formatOptions, null);
+    }
 
     /**
      * Validates the branch and defensively copies the collections, normalizing
@@ -118,7 +141,7 @@ public record ScanBranch(
     public static ScanBranch format(String formatName, List<String> locations,
             Map<String, Object> options) {
         return new ScanBranch("", List.of(), Map.of(), null, false, null, null, null,
-                formatName, locations, options);
+                formatName, locations, options, null);
     }
 
     /**
@@ -129,7 +152,7 @@ public static ScanBranch format(String formatName, List<String> locations,
      * @return the branch
      */
     public static ScanBranch of(String functionName, Object... positional) {
-        return new ScanBranch(functionName, List.of(positional), Map.of(), null, false, null, null, null, null, null, null);
+        return new ScanBranch(functionName, List.of(positional), Map.of(), null, false, null, null, null, null, null, null, null);
     }
 
     /**
@@ -141,7 +164,7 @@ public static ScanBranch of(String functionName, Object... positional) {
      * @return the branch
      */
     public static ScanBranch filtered(String functionName, String branchFilter, Object... positional) {
-        return new ScanBranch(functionName, List.of(positional), Map.of(), branchFilter, false, null, null, null, null, null, null);
+        return new ScanBranch(functionName, List.of(positional), Map.of(), branchFilter, false, null, null, null, null, null, null, null);
     }
 
     /**
@@ -152,7 +175,7 @@ public static ScanBranch filtered(String functionName, String branchFilter, Obje
      * @return the writable branch
      */
     public static ScanBranch writable(String functionName, Object... positional) {
-        return new ScanBranch(functionName, List.of(positional), Map.of(), null, true, null, null, null, null, null, null);
+        return new ScanBranch(functionName, List.of(positional), Map.of(), null, true, null, null, null, null, null, null, null);
     }
 
     /**
@@ -169,6 +192,13 @@ public static ScanBranch writable(String functionName, Object... positional) {
     public static ScanBranch catalogTable(
             String sourceCatalog, String sourceSchema, String sourceTable, String branchFilter) {
         return new ScanBranch("", List.of(), Map.of(), branchFilter, false, sourceCatalog, sourceSchema,
-                sourceTable, null, null, null);
+                sourceTable, null, null, null, null);
+    }
+
+    /** Return this branch with an authoritative function schema. */
+    public ScanBranch withSchemaName(String schema) {
+        return new ScanBranch(functionName, positional, named, branchFilter, writable,
+                sourceCatalog, sourceSchema, sourceTable, formatName, formatLocations,
+                formatOptions, schema);
     }
 }

vgi/src/main/java/farm/query/vgi/internal/ScanBranchesResultSerializer.java

@@ -54,7 +54,8 @@ private ScanBranchesResultSerializer() {}
             new Field("format_name", new FieldType(true, UTF8, null), null),
             new Field("format_locations", new FieldType(true, new ArrowType.List(), null),
                     List.of(new Field("item", new FieldType(true, UTF8, null), null))),
-            new Field("format_options", new FieldType(true, BINARY, null), null)));
+            new Field("format_options", new FieldType(true, BINARY, null), null),
+            new Field("schema_name", new FieldType(true, UTF8, null), null)));
 
     private static final Schema RESULT_SCHEMA = new Schema(List.of(
             new Field("branches", new FieldType(false, new ArrowType.List(), null),
@@ -126,6 +127,7 @@ private static byte[] encodeBranch(ScanBranch b, BufferAllocator alloc) {
                 // names, because an option value may be any Arrow type.
                 fo.setSafe(0, ScanFunctionResultEncoder.encodeArguments(List.of(), b.formatOptions()));
             }
+            setNullableString(root, "schema_name", b.schemaName());
             root.setRowCount(1);
             return writeStream(root);
         } catch (Exception e) {

vgi/src/main/java/farm/query/vgi/internal/ScanFunctionResultEncoder.java

@@ -43,17 +43,20 @@ private ScanFunctionResultEncoder() {}
      * @param positional positional scan arguments (may be {@code null})
      * @param named named scan arguments (may be {@code null})
      * @param requiredExtensions DuckDB extension names the scan depends on (may be {@code null})
+     * @param schemaName schema containing the function, or {@code null}
      * @return the 1-row IPC stream bytes
      */
     public static byte[] encode(String functionName, List<Object> positional,
-                                  Map<String, Object> named, List<String> requiredExtensions) {
+                                  Map<String, Object> named, List<String> requiredExtensions,
+                                  String schemaName) {
         BufferAllocator alloc = Allocators.root();
         Schema schema = new Schema(List.of(
                 new Field("function_name", new FieldType(false, UTF8, null), null),
                 new Field("arguments", new FieldType(false, BINARY, null), null),
                 new Field("required_extensions",
                         new FieldType(false, new ArrowType.List(), null),
-                        List.of(new Field("item", new FieldType(true, UTF8, null), null)))));
+                        List.of(new Field("item", new FieldType(true, UTF8, null), null))),
+                new Field("schema_name", new FieldType(true, UTF8, null), null)));
 
         byte[] argsBytes = encodeArguments(positional, named, alloc);
 
@@ -69,6 +72,9 @@ public static byte[] encode(String functionName, List<Object> positional,
             }
             w.endList();
             w.setValueCount(1);
+            VarCharVector schemaVector = (VarCharVector) root.getVector("schema_name");
+            if (schemaName == null) schemaVector.setNull(0);
+            else schemaVector.setSafe(0, new Text(schemaName));
             root.setRowCount(1);
             ByteArrayOutputStream baos = new ByteArrayOutputStream();
             try (ArrowStreamWriter sw = new ArrowStreamWriter(root, null, Channels.newChannel(baos))) {