ASTERIXDB-PR2: Add schema extraction pipeline for NL2SQL++ (SchemaContextBuilder)

asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/ColumnInfo.java

@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.asterix.spidersilk.schema;
+
+/**
+ * Holds metadata for a single field (column) within a Dataset.
+ *
+ * <p>Instances are created by {@link SchemaContextBuilder} while traversing the
+ * AsterixDB ADM type tree, and consumed by:
+ * <ul>
+ *   <li>{@link DatasetSchema#toDescriptionString()} — rendered into a prompt-ready string</li>
+ *   <li>ColumnPruner (PR-4) — scored for relevance and potentially filtered out</li>
+ * </ul>
+ */
+public class ColumnInfo {
+
+    private final String name;
+    private final String type;
+    private final boolean primaryKey;
+
+    /**
+     * @param name       field name as declared in the Dataset's item type
+     * @param type       human-readable type string produced by {@link DatasetSchemaFormatter}
+     * @param primaryKey {@code true} if this field is part of the Dataset's primary / partitioning key
+     */
+    public ColumnInfo(String name, String type, boolean primaryKey) {
+        this.name = name;
+        this.type = type;
+        this.primaryKey = primaryKey;
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public String getType() {
+        return type;
+    }
+
+    public boolean isPrimaryKey() {
+        return primaryKey;
+    }
+
+    /**
+     * Returns a compact description suitable for inclusion in an LLM prompt.
+     * Example: {@code "tweetid: bigint [PK]"} or {@code "message-text: string"}
+     */
+    public String toDescriptionString() {
+        return primaryKey ? name + ": " + type + " [PK]" : name + ": " + type;
+    }
+
+    @Override
+    public String toString() {
+        return "ColumnInfo{name='" + name + "', type='" + type + "', primaryKey=" + primaryKey + '}';
+    }
+}

asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchema.java

@@ -0,0 +1,115 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.asterix.spidersilk.schema;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Represents the schema of a single AsterixDB Dataset, including all field
+ * metadata extracted from its ADM item type.
+ *
+ * <p>This object travels through the three-layer schema injection pipeline:
+ * <ol>
+ *   <li>Created by {@link SchemaContextBuilder} with the full column list</li>
+ *   <li>ColumnPruner (PR-4) calls {@link #setPrunedColumns} to store a filtered subset</li>
+ *   <li>ValueHintsSampler (PR-5) calls {@link #setValueHints} to attach sample values</li>
+ *   <li>{@link #toDescriptionString()} is called by PromptBuilder to produce the final
+ *       prompt fragment for this Dataset</li>
+ * </ol>
+ *
+ * <p>Example output of {@link #toDescriptionString()}:
+ * <pre>
+ * Dataset TweetMessages (tweetid: bigint [PK], sender-location: point,
+ *     send-time: datetime, referred-topics: [string], message-text: string, author-id: bigint)
+ * </pre>
+ */
+public class DatasetSchema {
+
+    private final String datasetName;
+    private final List<ColumnInfo> allColumns;
+
+    /** Set by ColumnPruner; null until pruning is performed. */
+    private List<ColumnInfo> prunedColumns;
+
+    /** Set by ValueHintsSampler; null until sampling is performed. Key = field name. */
+    private Map<String, List<String>> valueHints;
+
+    public DatasetSchema(String datasetName, List<ColumnInfo> allColumns) {
+        this.datasetName = datasetName;
+        this.allColumns = Collections.unmodifiableList(new java.util.ArrayList<>(allColumns));
+    }
+
+    public String getDatasetName() {
+        return datasetName;
+    }
+
+    /** Returns the complete column list as extracted from the ADM item type. */
+    public List<ColumnInfo> getAllColumns() {
+        return allColumns;
+    }
+
+    /**
+     * Returns the pruned column list if {@link #setPrunedColumns} has been called,
+     * otherwise falls back to the full list.
+     */
+    public List<ColumnInfo> getEffectiveColumns() {
+        return prunedColumns != null ? prunedColumns : allColumns;
+    }
+
+    /** Called by ColumnPruner (PR-4) to store the relevance-filtered subset. */
+    public void setPrunedColumns(List<ColumnInfo> prunedColumns) {
+        this.prunedColumns = prunedColumns;
+    }
+
+    /** Called by ValueHintsSampler (PR-5) to attach sample values per field. */
+    public void setValueHints(Map<String, List<String>> valueHints) {
+        this.valueHints = valueHints;
+    }
+
+    public Map<String, List<String>> getValueHints() {
+        return valueHints;
+    }
+
+    /**
+     * Renders this Dataset's schema as a prompt-ready one-liner.
+     * Uses the pruned column list if available, otherwise the full column list.
+     *
+     * <p>Example:
+     * <pre>Dataset TweetMessages (tweetid: bigint [PK], message-text: string, author-id: bigint)</pre>
+     */
+    public String toDescriptionString() {
+        List<ColumnInfo> columns = getEffectiveColumns();
+        StringBuilder sb = new StringBuilder("Dataset ").append(datasetName).append(" (");
+        for (int i = 0; i < columns.size(); i++) {
+            if (i > 0) {
+                sb.append(", ");
+            }
+            sb.append(columns.get(i).toDescriptionString());
+        }
+        sb.append(')');
+        return sb.toString();
+    }
+
+    @Override
+    public String toString() {
+        return "DatasetSchema{name='" + datasetName + "', columns=" + allColumns.size() + '}';
+    }
+}

asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchemaFormatter.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.asterix.spidersilk.schema;
+
+import org.apache.asterix.om.types.AOrderedListType;
+import org.apache.asterix.om.types.ARecordType;
+import org.apache.asterix.om.types.AUnionType;
+import org.apache.asterix.om.types.AUnorderedListType;
+import org.apache.asterix.om.types.IAType;
+
+/**
+ * Converts AsterixDB ADM type objects ({@link IAType}) into human-readable strings
+ * suitable for inclusion in an LLM prompt.
+ *
+ * <p>Type rendering rules:
+ * <ul>
+ *   <li>Primitive types (bigint, string, boolean, …) → lower-cased type name</li>
+ *   <li>{@code ARecordType} (nested object) → {@code {field1: type1, field2: type2}}</li>
+ *   <li>{@code AOrderedListType} (ordered array) → {@code [itemType]}</li>
+ *   <li>{@code AUnorderedListType} (bag/multiset) → {@code {{itemType}}}</li>
+ *   <li>{@code AUnionType} (nullable/missable field) → {@code actualType?}</li>
+ * </ul>
+ *
+ * <p>Recursive formatting is limited to {@value #MAX_DEPTH} levels to prevent
+ * runaway output for deeply nested types.
+ */
+public class DatasetSchemaFormatter {
+
+    private static final int MAX_DEPTH = 4;
+
+    /**
+     * Formats {@code type} as a human-readable string.
+     *
+     * @param type the ADM type to format; {@code null} is rendered as {@code "any"}
+     * @return a compact, prompt-friendly type description
+     */
+    public String formatType(IAType type) {
+        return formatType(type, 0);
+    }
+
+    private String formatType(IAType type, int depth) {
+        if (type == null) {
+            return "any";
+        }
+        if (depth >= MAX_DEPTH) {
+            return "object";
+        }
+        switch (type.getTypeTag()) {
+            case UNION:
+                // Nullable or missable field: unwrap to the actual type and append '?'
+                return formatType(((AUnionType) type).getActualType(), depth) + "?";
+            case OBJECT:
+                return formatRecord((ARecordType) type, depth);
+            case ARRAY:
+                // Ordered list (SQL++ array syntax: [itemType])
+                return "[" + formatType(((AOrderedListType) type).getItemType(), depth + 1) + "]";
+            case MULTISET:
+                // Unordered list / bag (SQL++ multiset syntax: {{itemType}})
+                return "{{" + formatType(((AUnorderedListType) type).getItemType(), depth + 1) + "}}";
+            default:
+                return type.getTypeName().toLowerCase();
+        }
+    }
+
+    /**
+     * Formats a record type as {@code {field1: type1, field2: type2}}.
+     * For top-level fields of a Dataset (depth 0), the outer braces are omitted
+     * because the field list is already wrapped by the Dataset description.
+     */
+    private String formatRecord(ARecordType recordType, int depth) {
+        String[] fieldNames = recordType.getFieldNames();
+        IAType[] fieldTypes = recordType.getFieldTypes();
+        if (fieldNames.length == 0) {
+            return "object";
+        }
+        StringBuilder sb = new StringBuilder();
+        boolean wrapWithBraces = depth > 0;
+        if (wrapWithBraces) {
+            sb.append('{');
+        }
+        for (int i = 0; i < fieldNames.length; i++) {
+            if (i > 0) {
+                sb.append(", ");
+            }
+            sb.append(fieldNames[i]).append(": ").append(formatType(fieldTypes[i], depth + 1));
+        }
+        if (wrapWithBraces) {
+            sb.append('}');
+        }
+        return sb.toString();
+    }
+}