Skip to content

Commit

Permalink
added array data type support (#146)
Browse files Browse the repository at this point in the history
  • Loading branch information
shaharuk-yb authored Aug 28, 2024
1 parent d864c9d commit ec48a36
Show file tree
Hide file tree
Showing 8 changed files with 171 additions and 39 deletions.
28 changes: 28 additions & 0 deletions src/main/java/com/oltpbenchmark/benchmarks/dataloader/Column.java
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@
public class Column {
private String columnName;
private String dataType;
private String baseDataType;
private Integer characterMaximumLength;
private Boolean isIdentity;
private Boolean nullable;

// Getters and Setters
public String getColumnName() {
Expand All @@ -31,6 +33,22 @@ public void setCharacterMaximumLength(Integer characterMaximumLength) {
this.characterMaximumLength = characterMaximumLength;
}

public String getBaseDataType() {
return baseDataType;
}

public void setBaseDataType(String baseDataType) {
this.baseDataType = baseDataType;
}

public Boolean getIdentity() {
return isIdentity;
}

public void setIdentity(Boolean identity) {
isIdentity = identity;
}

public Boolean getIsIdentity() {
return isIdentity;
}
Expand All @@ -39,13 +57,23 @@ public void setIsIdentity(Boolean isIdentity) {
this.isIdentity = isIdentity;
}

public Boolean getNullable() {
return nullable;
}

public void setNullable(Boolean nullable) {
this.nullable = nullable;
}

@Override
public String toString() {
return "Column{" +
"columnName='" + columnName + '\'' +
", dataType='" + dataType + '\'' +
", baseDataType='" + baseDataType + '\'' +
", characterMaximumLength=" + characterMaximumLength +
", isIdentity=" + isIdentity +
", nullable=" + nullable +
'}';
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ protected Loader<DataGenerator> makeLoaderImpl() {
// load properties file

return new DataGeneratorLoader(this, getProperties("datatype-mapping.properties"),
getProperties("pk-mapping.properties"), getFkProperties());
getProperties("pk-mapping.properties"),
getProperties("array-mapping.properties"), getFkProperties());
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,16 +27,18 @@ public class DataGeneratorLoader extends Loader<DataGenerator> {

private final Map<String, PropertyMapping> properties;
private final Map<String, FkPropertyMapping> fkProperties;

private final Map<String, PropertyMapping> arrayProperties;
private final Map<String, PropertyMapping> pkProperties;
private int minLevel = Integer.MAX_VALUE;

public DataGeneratorLoader(DataGenerator benchmark, Map<String, PropertyMapping> properties,
Map<String, PropertyMapping> pkProperties, Map<String, FkPropertyMapping> fkProperties) {
Map<String, PropertyMapping> pkProperties, Map<String, PropertyMapping> arrayProperties,
Map<String, FkPropertyMapping> fkProperties) {
super(benchmark);
this.properties = properties;
this.fkProperties = fkProperties;
this.pkProperties = pkProperties;
this.arrayProperties = arrayProperties;
}

@Override
Expand Down Expand Up @@ -173,9 +175,7 @@ public static String getTableSchemaIfTableExists(String tableName, Connection co

public static List<Column> getTableSchema(String tableName, Connection conn, String tableSchema) {
List<Column> tableSchemaList = new ArrayList<>();
String query = "SELECT column_name, data_type, character_maximum_length, is_identity " +
"FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = ? and table_schema = ?";

String query = "SELECT c.column_name, c.data_type, CASE WHEN c.data_type = 'ARRAY' THEN (SELECT replace(t.typname, '_', '') FROM pg_type t JOIN pg_namespace ns ON t.typnamespace = ns.oid WHERE t.oid = (SELECT typarray FROM pg_type WHERE typname = replace(c.udt_name, '_', ''))) ELSE c.data_type END as base_data_type, c.character_maximum_length, c.is_identity, CASE WHEN c.is_nullable = 'YES' THEN TRUE END as nullable FROM INFORMATION_SCHEMA.COLUMNS c WHERE c.table_name = ? AND c.table_schema = ?";
try (PreparedStatement pstmt = conn.prepareStatement(query)) {
pstmt.setString(1, tableName);
pstmt.setString(2, tableSchema);
Expand All @@ -188,6 +188,8 @@ public static List<Column> getTableSchema(String tableName, Connection conn, Str
column.setDataType(rs.getString("data_type").replaceAll("\\s", ""));
column.setCharacterMaximumLength((Integer) rs.getObject("character_maximum_length"));
column.setIsIdentity("YES".equals(rs.getString("is_identity")));
column.setBaseDataType(rs.getString("base_data_type"));
column.setNullable(rs.getBoolean("nullable"));
tableSchemaList.add(column);
}
} catch (SQLException e) {
Expand Down Expand Up @@ -429,18 +431,25 @@ public Map<String, PropertyMapping> utilsMapping(List<Column> tableSchema, List<
// take care of the rest of the keys
for (Column col : tableSchema) {
if (!columnToUtilMapping.containsKey(col.getColumnName())) {
// skip the columns having array data types
if (col.getDataType().equalsIgnoreCase("array"))
continue;
PropertyMapping pm;
// if column is one of the unique constraint columns, then use primary key util functions for it
if (uniqueConstraintColumns.contains(col.getColumnName()))
pm = pkProperties.get(col.getDataType());
else
pm = properties.get(col.getDataType().toLowerCase());
else {

if (col.getDataType().equalsIgnoreCase("array")) {
// if it is array data type, find utility function with base data type
pm = arrayProperties.get(col.getBaseDataType().toLowerCase());
// if suitable utility function is not present and the column is nullable, it can be skipped
if (pm == null && col.getNullable()) continue;
} else {
// if it is not array type, find utility function is properties.
pm = properties.get(col.getDataType().toLowerCase());
}
}

if (pm == null) {
throw new RuntimeException(String.format("Cannot find suitable utility function for column " +
if (pm == null && !col.getNullable()) {
throw new RuntimeException(String.format("Cannot find suitable utility function for NOT NULL column " +
"`%s` of datatype `%s`. Consider asking #perf team to add a utility function for given " +
"data type", col.getColumnName(), col.getDataType()));
}
Expand Down Expand Up @@ -470,20 +479,20 @@ public Root generateMappingObject(String tableName, int rows, Map<String, Proper

prop.params.forEach(param -> {
if (fkColNames.contains(colName)) {
if (param instanceof Integer)
if (param instanceof Integer || param instanceof Long)
col.params.add(param);
else
col.params.add(param.toString());
} else if (udColumns.containsKey(colName)) {
if (param instanceof Integer)
if (param instanceof Integer || param instanceof Long)
col.params.add(param);
else
col.params.add(param.toString());
} else {
if (param.toString().equalsIgnoreCase("rows")) {
col.params.add(rows);
} else {
col.params.add(Integer.parseInt(param.toString()));
col.params.add(Long.parseLong(param.toString()));
}
}
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -189,22 +189,12 @@ public void load(Connection conn) throws SQLException {
Map<String, Object> columnsDetails = this.columns.get(index);
if (columnsDetails.containsKey("count")) {
for (int i = 0; i < (int) columnsDetails.get("count"); i++) {
columnString.append(columnsDetails.get("name") + String.valueOf(i + 1)).append(",");
if (this.baseutils.get(index).getInstance().getClass().getName().
toLowerCase().indexOf("json") >= 0)
valueString.append("?::JSON,");
else {
valueString.append("?,");
}
columnString.append(columnsDetails.get("name")).append(i + 1).append(",");
typeCastDataTypes(valueString, index);
}
} else {
columnString.append(columnsDetails.get("name")).append(",");
if (this.baseutils.get(index).getInstance().getClass().getName().
toLowerCase().indexOf("json") >= 0)
valueString.append("?::JSON,");
else {
valueString.append("?,");
}
typeCastDataTypes(valueString, index);
}

}
Expand Down Expand Up @@ -240,6 +230,25 @@ public void load(Connection conn) throws SQLException {
numberOfGeneratorFinished += 1;
}

private void typeCastDataTypes(StringBuilder valueString, int index) {
String utilName = this.baseutils.get(index).getInstance().getClass().getName().toLowerCase();
if (utilName.contains("array")) {
if(utilName.contains("integer"))
valueString.append("?::int[],");
else if (utilName.contains("long"))
valueString.append("?::bigint[],");
else if (utilName.contains("double"))
valueString.append("?::double[],");
else if (utilName.contains("text"))
valueString.append("?::text[],");
}
else if (utilName.contains("json"))
valueString.append("?::JSON,");
else {
valueString.append("?,");
}
}

@Override
public void afterLoad() {
if (numberOfGeneratorFinished != sizeOfLoadRule) return;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,8 @@ public class RandomIntegerArrayGen implements BaseUtil {

private List<Integer> integerArray;

private int minValue;
private int maxValue;

private Integer minValue;
private Integer maxValue;
private int arraySize;

public RandomIntegerArrayGen(List<Object> values) {
Expand All @@ -22,8 +21,8 @@ public RandomIntegerArrayGen(List<Object> values) {
this.arraySize = ((Number) values.get(0)).intValue();
if (arraySize <= 0)
throw new RuntimeException("Please enter positive integer array length");
this.minValue = ((Number) values.get(1)).intValue();
this.maxValue = ((Number) values.get(2)).intValue();
this.minValue = (Integer) values.get(1);
this.maxValue = (Integer) values.get(2);
if (minValue > maxValue)
throw new RuntimeException("Please enter correct bounds for max and min value");
}
Expand All @@ -36,8 +35,8 @@ public RandomIntegerArrayGen(List<Object> values, int workerId, int totalWorkers
this.arraySize = ((Number) values.get(0)).intValue();
if (arraySize <= 0)
throw new RuntimeException("Please enter positive integer array length");
this.minValue = ((Number) values.get(1)).intValue();
this.maxValue = ((Number) values.get(2)).intValue();
this.minValue = (Integer) values.get(1);
this.maxValue = (Integer) values.get(2);
if (minValue > maxValue)
throw new RuntimeException("Please enter correct bounds for max and min value");
}
Expand All @@ -47,7 +46,7 @@ public Object run() throws ClassNotFoundException, InvocationTargetException, No
Random random = new Random();
integerArray = new ArrayList<>();
for (int i = 0; i < arraySize; i++) {
int rd = random.nextInt((maxValue - minValue) + 1) + minValue;
Integer rd = random.nextInt((maxValue - minValue) + 1) + minValue;
integerArray.add(rd);
}
return integerArray.toString().replaceFirst("\\[", "{").replace("]", "}");
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
package com.oltpbenchmark.benchmarks.featurebench.utils;

import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class RandomLongArrayGen implements BaseUtil {

private List<Long> longArray;

private Long minValue;
private Long maxValue;
private int arraySize;

public RandomLongArrayGen(List<Object> values) {
if (values.size() != 3) {
throw new RuntimeException("Incorrect number of parameters for util function "
+ this.getClass());
}
this.arraySize = ((Number) values.get(0)).intValue();
if (arraySize <= 0)
throw new RuntimeException("Please enter positive integer array length");
this.minValue = ((Number) values.get(1)).longValue();
this.maxValue = ((Number) values.get(2)).longValue();
if (minValue > maxValue)
throw new RuntimeException("Please enter correct bounds for max and min value");
}

public RandomLongArrayGen(List<Object> values, int workerId, int totalWorkers) {
if (values.size() != 3) {
throw new RuntimeException("Incorrect number of parameters for util function "
+ this.getClass());
}
this.arraySize = ((Number) values.get(0)).intValue();
if (arraySize <= 0)
throw new RuntimeException("Please enter positive integer array length");
this.minValue = ((Number) values.get(1)).longValue();
this.maxValue = ((Number) values.get(2)).longValue();
if (minValue > maxValue)
throw new RuntimeException("Please enter correct bounds for max and min value");
}

@Override
public Object run() throws ClassNotFoundException, InvocationTargetException, NoSuchMethodException, InstantiationException, IllegalAccessException {
Random random = new Random();
longArray = new ArrayList<>();
for (int i = 0; i < arraySize; i++) {
Long rd = random.nextLong((maxValue - minValue) + 1) + minValue;
longArray.add(rd);
}
return longArray.toString().replaceFirst("\\[", "{").replace("]", "}");
}
}

Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
# arraySize, minLength,maxLength
int2=RandomIntegerArrayGen:3:1,1,32767
int4=RandomIntegerArrayGen:3:1,1,2147483647
int8=RandomLongArrayGen:3:1,1,9223372036854775807
smallint=RandomIntegerArrayGen:3:1,1,32767
integer=RandomIntegerArrayGen:3:1,1,2147483647
bigint=RandomLongArrayGen:3:1,1,9223372036854775807

#int2=RandomIntegerArrayGen:3:2,-32768,32767
#int4=RandomIntegerArrayGen:3:2,-2147483648,2147483647
#int8=RandomLongArrayGen:3:2,-9223372036854775808,9223372036854775807
#smallint=RandomIntegerArrayGen:3:2,-32768,32767
#integer=RandomIntegerArrayGen:3:2,-2147483648,2147483647
#bigint=RandomLongArrayGen:3:2,-9223372036854775808,9223372036854775807

#decimal
#numeric
#real
serial=RandomIntegerArrayGen:3:2,1,2147483647
bigserial=RandomLongArrayGen:3:2,1,9223372036854775807
#money
#char
#varchar
#date
#time
#timestamp
#timestamptz
#interval
#boolean
#json
#jsonb
text=RandomTextArrayGen:3:1,20,20
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,7 @@ text=RandomAString:2:1,1000
uuid=RandomUUID:0:


# arraySize, minLength,maxLength
array=RandomTextArrayGen:3:1,20,20


timestampwithtimezone=RandomTimestampWithTimeZone:1:rows
timestampwithouttimezone=RandomTimestampWithoutTimeZone:1:rows
Expand Down

0 comments on commit ec48a36

Please sign in to comment.