diff --git a/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/ColumnInfo.java b/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/ColumnInfo.java new file mode 100644 index 0000000000..6bd39823bd --- /dev/null +++ b/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. + * + *

Instances are created by {@link SchemaContextBuilder} while traversing the + * AsterixDB ADM type tree, and consumed by: + *

+ */ +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 + '}'; + } +} diff --git a/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchema.java b/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchema.java new file mode 100644 index 0000000000..81eb8e511a --- /dev/null +++ b/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. + * + *

This object travels through the three-layer schema injection pipeline: + *

    + *
  1. Created by {@link SchemaContextBuilder} with the full column list
  2. + *
  3. ColumnPruner (PR-4) calls {@link #setPrunedColumns} to store a filtered subset
  4. + *
  5. ValueHintsSampler (PR-5) calls {@link #setValueHints} to attach sample values
  6. + *
  7. {@link #toDescriptionString()} is called by PromptBuilder to produce the final + * prompt fragment for this Dataset
  8. + *
+ * + *

Example output of {@link #toDescriptionString()}: + *

+ * Dataset TweetMessages (tweetid: bigint [PK], sender-location: point,
+ *     send-time: datetime, referred-topics: [string], message-text: string, author-id: bigint)
+ * 
+ */ +public class DatasetSchema { + + private final String datasetName; + private final List allColumns; + + /** Set by ColumnPruner; null until pruning is performed. */ + private List prunedColumns; + + /** Set by ValueHintsSampler; null until sampling is performed. Key = field name. */ + private Map> valueHints; + + public DatasetSchema(String datasetName, List 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 getAllColumns() { + return allColumns; + } + + /** + * Returns the pruned column list if {@link #setPrunedColumns} has been called, + * otherwise falls back to the full list. + */ + public List getEffectiveColumns() { + return prunedColumns != null ? prunedColumns : allColumns; + } + + /** Called by ColumnPruner (PR-4) to store the relevance-filtered subset. */ + public void setPrunedColumns(List prunedColumns) { + this.prunedColumns = prunedColumns; + } + + /** Called by ValueHintsSampler (PR-5) to attach sample values per field. */ + public void setValueHints(Map> valueHints) { + this.valueHints = valueHints; + } + + public Map> 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. + * + *

Example: + *

Dataset TweetMessages (tweetid: bigint [PK], message-text: string, author-id: bigint)
+ */ + public String toDescriptionString() { + List 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() + '}'; + } +} diff --git a/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchemaFormatter.java b/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/DatasetSchemaFormatter.java new file mode 100644 index 0000000000..bed6e90ac7 --- /dev/null +++ b/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. + * + *

Type rendering rules: + *

    + *
  • Primitive types (bigint, string, boolean, …) → lower-cased type name
  • + *
  • {@code ARecordType} (nested object) → {@code {field1: type1, field2: type2}}
  • + *
  • {@code AOrderedListType} (ordered array) → {@code [itemType]}
  • + *
  • {@code AUnorderedListType} (bag/multiset) → {@code {{itemType}}}
  • + *
  • {@code AUnionType} (nullable/missable field) → {@code actualType?}
  • + *
+ * + *

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(); + } +} diff --git a/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilder.java b/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilder.java new file mode 100644 index 0000000000..0f0ab4dc3f --- /dev/null +++ b/asterixdb/asterix-spidersilk/src/main/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilder.java @@ -0,0 +1,180 @@ +/* + * 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.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import org.apache.asterix.common.metadata.DataverseName; +import org.apache.asterix.common.metadata.MetadataConstants; +import org.apache.asterix.metadata.MetadataManager; +import org.apache.asterix.metadata.MetadataTransactionContext; +import org.apache.asterix.metadata.entities.Dataset; +import org.apache.asterix.om.types.ARecordType; +import org.apache.asterix.om.types.IAType; +import org.apache.asterix.spidersilk.api.Nl2SqlException; +import org.apache.asterix.spidersilk.api.SchemaContext; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +/** + * Builds a {@link SchemaContext} by reading Dataset and type metadata from + * AsterixDB's {@link MetadataManager}. + * + *

For each Dataset in the target Dataverse, this class: + *

    + *
  1. Fetches the ADM item type via {@code MetadataManager.INSTANCE.getDatatype()}
  2. + *
  3. Recursively formats the type tree into a human-readable field list using + * {@link DatasetSchemaFormatter}
  4. + *
  5. Marks primary-key fields from {@code Dataset.getPrimaryKeys()}
  6. + *
  7. Returns a {@link SchemaContext} whose {@code datasetDescriptions} list is + * ready to be consumed by {@link SchemaEmbeddingService} (PR-3)
  8. + *
+ * + *

All metadata reads are wrapped in a single metadata transaction that is + * committed on success and aborted on any failure, following the standard + * AsterixDB metadata access pattern. + * + *

Example output description for one Dataset: + *

+ * Dataset TweetMessages (tweetid: bigint [PK], sender-location: point,
+ *     send-time: datetime, referred-topics: [string], message-text: string, author-id: bigint)
+ * 
+ */ +public class SchemaContextBuilder { + + private static final Logger LOGGER = LogManager.getLogger(); + + private final DatasetSchemaFormatter formatter; + private final String databaseName; + + /** + * Creates a builder that reads from the default AsterixDB database + * ({@code MetadataConstants.DEFAULT_DATABASE}). + */ + public SchemaContextBuilder() { + this(MetadataConstants.DEFAULT_DATABASE); + } + + /** + * Creates a builder that reads from the specified database. + * + * @param databaseName the AsterixDB database name (typically {@code "Default"}) + */ + public SchemaContextBuilder(String databaseName) { + this.databaseName = databaseName; + this.formatter = new DatasetSchemaFormatter(); + } + + /** + * Builds a {@link SchemaContext} containing descriptions for all Datasets + * in the given Dataverse. + * + * @param dataverse the target Dataverse name (e.g. {@code "TinySocial"}) + * @return a populated {@link SchemaContext} ready for embedding and prompt injection + * @throws Nl2SqlException if the Dataverse does not exist or metadata access fails + */ + public SchemaContext build(String dataverse) throws Nl2SqlException { + MetadataTransactionContext txnCtx = null; + try { + DataverseName dataverseName = DataverseName.createSinglePartName(dataverse); + txnCtx = MetadataManager.INSTANCE.beginTransaction(); + + List datasets = MetadataManager.INSTANCE.getDataverseDatasets(txnCtx, databaseName, dataverseName); + + if (datasets.isEmpty()) { + LOGGER.warn("No datasets found in dataverse '{}'", dataverse); + } + + List descriptions = new ArrayList<>(datasets.size()); + for (Dataset dataset : datasets) { + try { + DatasetSchema schema = buildDatasetSchema(txnCtx, dataset); + descriptions.add(schema.toDescriptionString()); + } catch (Exception e) { + // Skip datasets whose type cannot be resolved rather than failing the whole request + LOGGER.warn("Skipping dataset '{}': failed to resolve type — {}", dataset.getDatasetName(), + e.getMessage()); + } + } + + MetadataManager.INSTANCE.commitTransaction(txnCtx); + txnCtx = null; + return new SchemaContext(dataverse, descriptions); + + } catch (Exception e) { + throw new Nl2SqlException("Failed to build schema context for dataverse: " + dataverse, e); + } finally { + if (txnCtx != null) { + try { + MetadataManager.INSTANCE.abortTransaction(txnCtx); + } catch (Exception ignored) { + LOGGER.warn("Failed to abort metadata transaction", ignored); + } + } + } + } + + /** + * Builds a {@link DatasetSchema} for a single Dataset by resolving its item type + * and extracting field metadata. + */ + private DatasetSchema buildDatasetSchema(MetadataTransactionContext txnCtx, Dataset dataset) throws Exception { + // Collect primary-key field names for annotation + Set primaryKeyFields = extractPrimaryKeyFields(dataset); + + // Resolve the ADM item type + IAType itemType = MetadataManager.INSTANCE.getDatatype(txnCtx, dataset.getItemTypeDatabaseName(), + dataset.getItemTypeDataverseName(), dataset.getItemTypeName()).getDatatype(); + + List columns = new ArrayList<>(); + if (itemType instanceof ARecordType) { + ARecordType recordType = (ARecordType) itemType; + String[] fieldNames = recordType.getFieldNames(); + IAType[] fieldTypes = recordType.getFieldTypes(); + for (int i = 0; i < fieldNames.length; i++) { + String typeStr = formatter.formatType(fieldTypes[i]); + boolean isPk = primaryKeyFields.contains(fieldNames[i]); + columns.add(new ColumnInfo(fieldNames[i], typeStr, isPk)); + } + } else { + LOGGER.debug("Dataset '{}' has non-record item type: {}", dataset.getDatasetName(), itemType.getTypeTag()); + } + + return new DatasetSchema(dataset.getDatasetName(), columns); + } + + /** + * Returns the set of top-level field names that form the primary key. + * For composite keys only the first part of each key path is collected, + * which is sufficient for annotation purposes in prompt generation. + */ + private Set extractPrimaryKeyFields(Dataset dataset) { + Set pkFields = new HashSet<>(); + List> primaryKeys = dataset.getPrimaryKeys(); + for (List keyPath : primaryKeys) { + if (!keyPath.isEmpty()) { + pkFields.add(keyPath.get(0)); + } + } + return pkFields; + } +} diff --git a/asterixdb/asterix-spidersilk/src/test/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilderTest.java b/asterixdb/asterix-spidersilk/src/test/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilderTest.java new file mode 100644 index 0000000000..5bbf76743f --- /dev/null +++ b/asterixdb/asterix-spidersilk/src/test/java/org/apache/asterix/spidersilk/schema/SchemaContextBuilderTest.java @@ -0,0 +1,197 @@ +/* + * 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.Arrays; +import java.util.Collections; +import java.util.List; + +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.BuiltinType; +import org.apache.asterix.om.types.IAType; +import org.junit.Assert; +import org.junit.Test; + +/** + * Unit tests for the PR-2 schema extraction components. + * + * These tests exercise {@link DatasetSchemaFormatter} and {@link DatasetSchema} + * using ADM type objects constructed directly in-memory, with no dependency on + * a running AsterixDB instance or MetadataManager. + * + * Integration tests that verify the full {@link SchemaContextBuilder#build(String)} + * path against a live AsterixDB + TinySocial dataset are left for the integration + * test suite (require a running cluster). + */ +public class SchemaContextBuilderTest { + + private final DatasetSchemaFormatter formatter = new DatasetSchemaFormatter(); + + // ------------------------------------------------------------------------- + // DatasetSchemaFormatter tests + // ------------------------------------------------------------------------- + + @Test + public void testFormatPrimitiveTypes() { + Assert.assertEquals("int64", formatter.formatType(BuiltinType.AINT64)); + Assert.assertEquals("string", formatter.formatType(BuiltinType.ASTRING)); + Assert.assertEquals("boolean", formatter.formatType(BuiltinType.ABOOLEAN)); + Assert.assertEquals("double", formatter.formatType(BuiltinType.ADOUBLE)); + } + + @Test + public void testFormatNullType() { + Assert.assertEquals("any", formatter.formatType(null)); + } + + @Test + public void testFormatOrderedList() { + // [string] — ordered list of strings (SQL++ array) + AOrderedListType listType = new AOrderedListType(BuiltinType.ASTRING, "string-list"); + String result = formatter.formatType(listType); + Assert.assertEquals("[string]", result); + } + + @Test + public void testFormatNullableField() { + // string? — union of string + missing (nullable field) + AUnionType unionType = + new AUnionType(Arrays.asList(BuiltinType.ASTRING, BuiltinType.AMISSING), "nullable-string"); + String result = formatter.formatType(unionType); + Assert.assertEquals("string?", result); + } + + @Test + public void testFormatNestedRecord() { + // Nested record: { street: string, city: string } + ARecordType addressType = new ARecordType("AddressType", new String[] { "street", "city" }, + new IAType[] { BuiltinType.ASTRING, BuiltinType.ASTRING }, false); + + // Top-level record with a nested field + ARecordType personType = new ARecordType("PersonType", new String[] { "name", "address" }, + new IAType[] { BuiltinType.ASTRING, addressType }, false); + + // formatType on the top-level record (depth=0) should not wrap in braces + String result = formatter.formatType(personType); + Assert.assertTrue("Should contain nested field 'address'", result.contains("address")); + Assert.assertTrue("Should contain nested field 'street'", result.contains("street")); + Assert.assertTrue("Should contain nested field 'city'", result.contains("city")); + } + + @Test + public void testFormatTweetMessagesSchema() { + // Mimics the TinySocial TweetMessages item type + AOrderedListType topicsType = new AOrderedListType(BuiltinType.ASTRING, "topics-list"); + ARecordType tweetType = new ARecordType("TweetMessageType", + new String[] { "tweetid", "sender-location", "send-time", "referred-topics", "message-text", + "author-id" }, + new IAType[] { BuiltinType.AINT64, BuiltinType.ANY, BuiltinType.ADATETIME, topicsType, + BuiltinType.ASTRING, BuiltinType.AINT64 }, + false); + + String result = formatter.formatType(tweetType); + Assert.assertTrue(result.contains("tweetid")); + Assert.assertTrue(result.contains("int64")); + Assert.assertTrue(result.contains("message-text")); + Assert.assertTrue(result.contains("referred-topics")); + Assert.assertTrue(result.contains("[string]")); + } + + // ------------------------------------------------------------------------- + // ColumnInfo tests + // ------------------------------------------------------------------------- + + @Test + public void testColumnInfoPrimaryKeyDescription() { + ColumnInfo pk = new ColumnInfo("tweetid", "bigint", true); + Assert.assertEquals("tweetid: bigint [PK]", pk.toDescriptionString()); + } + + @Test + public void testColumnInfoNonPrimaryKeyDescription() { + ColumnInfo col = new ColumnInfo("message-text", "string", false); + Assert.assertEquals("message-text: string", col.toDescriptionString()); + } + + // ------------------------------------------------------------------------- + // DatasetSchema tests + // ------------------------------------------------------------------------- + + @Test + public void testDatasetSchemaDescriptionString() { + List columns = Arrays.asList(new ColumnInfo("tweetid", "int64", true), + new ColumnInfo("message-text", "string", false), new ColumnInfo("author-id", "int64", false)); + + DatasetSchema schema = new DatasetSchema("TweetMessages", columns); + String desc = schema.toDescriptionString(); + + Assert.assertTrue(desc.startsWith("Dataset TweetMessages (")); + Assert.assertTrue(desc.contains("tweetid: int64 [PK]")); + Assert.assertTrue(desc.contains("message-text: string")); + Assert.assertTrue(desc.endsWith(")")); + } + + @Test + public void testDatasetSchemaFallsBackToAllColumnsBeforePruning() { + List columns = + Arrays.asList(new ColumnInfo("id", "bigint", true), new ColumnInfo("name", "string", false)); + + DatasetSchema schema = new DatasetSchema("Users", columns); + + // Before pruning, getEffectiveColumns() returns the full list + Assert.assertEquals(2, schema.getEffectiveColumns().size()); + } + + @Test + public void testDatasetSchemaUsesPrunedColumnsAfterPruning() { + List allColumns = Arrays.asList(new ColumnInfo("id", "bigint", true), + new ColumnInfo("name", "string", false), new ColumnInfo("created-at", "datetime", false)); + + DatasetSchema schema = new DatasetSchema("Users", allColumns); + + // Simulate ColumnPruner keeping only id and name + List pruned = + Arrays.asList(new ColumnInfo("id", "bigint", true), new ColumnInfo("name", "string", false)); + schema.setPrunedColumns(pruned); + + Assert.assertEquals(2, schema.getEffectiveColumns().size()); + String desc = schema.toDescriptionString(); + Assert.assertFalse("Pruned field should not appear", desc.contains("created-at")); + } + + @Test + public void testDatasetSchemaImmutableAllColumns() { + List mutable = new java.util.ArrayList<>(); + mutable.add(new ColumnInfo("id", "bigint", true)); + DatasetSchema schema = new DatasetSchema("Foo", mutable); + + // Modifying original list must not affect the schema + mutable.add(new ColumnInfo("extra", "string", false)); + Assert.assertEquals(1, schema.getAllColumns().size()); + } + + @Test + public void testEmptyDatasetDescription() { + DatasetSchema schema = new DatasetSchema("EmptyDataset", Collections.emptyList()); + String desc = schema.toDescriptionString(); + Assert.assertEquals("Dataset EmptyDataset ()", desc); + } +}