[improvement][headless]Add partitionTime and primaryKey field into the Schema part of the Text2SQL prompt. #1621

This commit is contained in:
jerryjzhang
2024-09-11 21:29:37 +08:00
parent a82e3c8b1e
commit 82b5fa966a
9 changed files with 83 additions and 74 deletions

View File

@@ -134,13 +134,21 @@ public class DataSetSchema {
}
public boolean containsPartitionDimensions() {
return dimensions.stream().anyMatch(SchemaElement::containsPartitionTime);
return dimensions.stream().anyMatch(SchemaElement::isPartitionTime);
}
public SchemaElement getPartitionDimension() {
for (SchemaElement dimension : dimensions) {
String partitionTimeFormat = dimension.getPartitionTimeFormat();
if (StringUtils.isNotBlank(partitionTimeFormat)) {
if (dimension.isPartitionTime()) {
return dimension;
}
}
return null;
}
public SchemaElement getPrimaryKey() {
for (SchemaElement dimension : dimensions) {
if (dimension.isPrimaryKey()) {
return dimension;
}
}

View File

@@ -63,7 +63,7 @@ public class SchemaElement implements Serializable {
return Objects.hashCode(dataSetId, id, name, bizName, type);
}
public boolean containsPartitionTime() {
public boolean isPartitionTime() {
if (MapUtils.isEmpty(extInfo)) {
return false;
}
@@ -78,6 +78,21 @@ public class SchemaElement implements Serializable {
return DimensionType.isPartitionTime(dimensionTYpe);
}
public boolean isPrimaryKey() {
if (MapUtils.isEmpty(extInfo)) {
return false;
}
Object o = extInfo.get(DimensionConstants.DIMENSION_TYPE);
DimensionType dimensionTYpe = null;
if (o instanceof DimensionType) {
dimensionTYpe = (DimensionType) o;
}
if (o instanceof String) {
dimensionTYpe = DimensionType.valueOf((String) o);
}
return DimensionType.isIdentity(dimensionTYpe);
}
public String getTimeFormat() {
if (MapUtils.isEmpty(extInfo)) {
return null;
@@ -87,7 +102,7 @@ public class SchemaElement implements Serializable {
public String getPartitionTimeFormat() {
String timeFormat = getTimeFormat();
if (StringUtils.isNotBlank(timeFormat) && containsPartitionTime()) {
if (StringUtils.isNotBlank(timeFormat) && isPartitionTime()) {
return timeFormat;
}
return "";

View File

@@ -21,4 +21,8 @@ public enum DimensionType {
public static boolean isPartitionTime(DimensionType type) {
return type == partition_time;
}
public static boolean isIdentity(DimensionType type) {
return type == identify;
}
}