Skip to content

Commit

Permalink
Semi structured flatten - Unit testing support (finos#2475)
Browse files Browse the repository at this point in the history 
* add udf for flattening in H2

* Update h2Extension1_4_200.pure

---------

Co-authored-by: Sai Sriharsha Annepu <[email protected]>
  • Loading branch information
@tanujgirdhar @gs-ssh16 @AFine-gs
2 people authored and AFine-gs committed Jan 9, 2024
1 parent ae7dad6 commit 9feef9f
Show file tree
Hide file tree
Showing 11 changed files with 1,782 additions and 24 deletions.
16 changes: 16 additions & 0 deletions ...xecution/stores/relational/test/semiStructured/TestSnowflakeSemiStructuredFlattening.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,22 @@ public void testSemiStructuredSubAggregationDeep()
Assert.assertEquals(wrapPreAndFinallyExecutionSqlQuery(TDSType, snowflakeExpected), snowflakePlan);
}

@Test
public void testSemiStructuredMultiLevelFlatten()
{
String snowflakePlan = this.buildExecutionPlanString("flatten::semiStructuredMultiLevelFlattening__TabularDataSet_1_", snowflakeMapping, snowflakeRuntime);
String snowflakeExpected =
" Relational\n" +
" (\n" +
" type = TDS[(First Name, String, VARCHAR(100), \"\"), (Firm Name, String, \"\", \"\"), (Firm Address Name Line No, Integer, \"\", \"\")]\n" +
" resultColumns = [(\"First Name\", VARCHAR(100)), (\"Firm Name\", \"\"), (\"Firm Address Name Line No\", \"\")]\n" +
" sql = select \"root\".FIRSTNAME as \"First Name\", \"root\".FIRM_DETAILS['legalName']::varchar as \"Firm Name\", \"ss_flatten_1\".value['lineno'] as \"Firm Address Name Line No\" from PERSON_SCHEMA.PERSON_TABLE as \"root\" inner join lateral flatten(input => \"root\".FIRM_DETAILS['addresses'], outer => true, recursive => false, mode => 'array') as \"ss_flatten_0\" inner join lateral flatten(input => \"ss_flatten_0\".value['lines'], outer => true, recursive => false, mode => 'array') as \"ss_flatten_1\"\n" +
" connection = RelationalDatabaseConnection(type = \"Snowflake\")\n" +
" )\n";
String TDSType = " type = TDS[(First Name, String, VARCHAR(100), \"\"), (Firm Name, String, \"\", \"\"), (Firm Address Name Line No, Integer, \"\", \"\")]\n";
Assert.assertEquals(wrapPreAndFinallyExecutionSqlQuery(TDSType, snowflakeExpected), snowflakePlan);
}

@Test
public void testSemiStructuredMultiFlatten()
{
Expand Down
166 changes: 164 additions & 2 deletions ...ain/java/org/finos/legend/engine/plan/execution/stores/relational/LegendH2Extensions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.finos.legend.engine.shared.core.ObjectMapperFactory;
import org.h2.tools.SimpleResultSet;
import org.h2.value.Value;
import org.h2.value.ValueArray;
import org.h2.value.ValueBigint;
import org.h2.value.ValueBoolean;
import org.h2.value.ValueDouble;
Expand All @@ -28,9 +30,15 @@
import org.h2.value.ValueReal;
import org.h2.value.ValueVarchar;

import java.math.BigDecimal;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Types;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;

Expand Down Expand Up @@ -132,12 +140,12 @@ public static String legend_h2_extension_reverse_string(String string)

public static String legend_h2_extension_split_part(String string, String token, Integer part)
{
if (part < 1)
if (part < 1)
{
throw new IllegalArgumentException("Split part must be greater than zero");
}

if (string == null)
if (string == null)
{
return null;
}
Expand All @@ -147,4 +155,158 @@ public static String legend_h2_extension_split_part(String string, String token,

return parts.length > readjustedPart ? parts[readjustedPart] : null;
}

private static HashSet<Object> extractProperty(HashSet<Object> resultSet, Object pathToExtract)
{
HashSet<Object> res = new HashSet<>();
if (pathToExtract instanceof String)
{
String property = (String) pathToExtract;
if (property.equals("*"))
{
for (Object r: resultSet)
{
if (!(r instanceof Iterable))
{
continue;
}
for (Object o : (Iterable<?>) r)
{
res.add(o);
}
}
}
else
{
for (Object r: resultSet)
{
try
{
Object o = ((HashMap)(r)).get(property);
if (o != null)
{
res.add(o);
}
}
catch (Exception e)
{
e.printStackTrace(); // don't stop execution
}
}
}
}
else
{
int index = (int) pathToExtract;
for (Object r: resultSet)
{
try
{
res.add(((ArrayList)(r)).get(index));
}
catch (Exception e)
{
e.printStackTrace();
}

}
}
return res;
}

public static ResultSet legend_h2_extension_flatten_array(Connection conn, String tableName, String toFlattenColumnName, ValueArray jsonPaths)
{
try
{
String sql = String.format("select distinct %s as flat from %s", toFlattenColumnName, tableName);
ResultSet resultSet = conn.createStatement().executeQuery(sql);
HashSet<Object> toFlatten = new HashSet<>();
while (resultSet.next())
{
String json = resultSet.getString("flat");
toFlatten.add(OBJECT_MAPPER.readValue(json, HashMap.class));
}

ArrayList<Object> pathsToExtract = OBJECT_MAPPER.readValue(jsonPaths.getString(), ArrayList.class);

for (Object path: pathsToExtract)
{
toFlatten = extractProperty(toFlatten, path);
}

SimpleResultSet flattenedResultSet = new SimpleResultSet();

flattenedResultSet.addColumn("__INPUT__", Types.VARCHAR, 1000, 0); // using the original array as joinKey

// use first non-null object to infer the type of value
boolean resolvedFlattenedType = false;

for (Object o: toFlatten)
{
if (!(o instanceof Iterable))
{
continue;
}
for (Object value : (Iterable<?>) o)
{
if (value instanceof Map || value instanceof List || value instanceof String)
{
flattenedResultSet.addColumn("VALUE", Types.VARCHAR, 1000, 0);
}
else if (value instanceof Boolean)
{
flattenedResultSet.addColumn("VALUE", Types.BOOLEAN, 0, 0);
}
else if (value instanceof Double || value instanceof Float || value instanceof BigDecimal)
{
flattenedResultSet.addColumn("VALUE", Types.DOUBLE, 20, 20);
}
else if (value instanceof Integer || value instanceof Long)
{
flattenedResultSet.addColumn("VALUE", Types.BIGINT, 0, 0);
}
else
{
throw new RuntimeException("unsupported data type in h2 extension");
}
resolvedFlattenedType = true;
break;
}
if (resolvedFlattenedType)
{
break;
}
}

if (!resolvedFlattenedType)
{
flattenedResultSet.addColumn("VALUE", Types.VARCHAR, 1000, 0);
return flattenedResultSet;
}

for (Object o: toFlatten)
{
if (!(o instanceof Iterable))
{
continue;
}
for (Object value : (Iterable<?>) o)
{
if (value instanceof Map || value instanceof List)
{
flattenedResultSet.addRow(Arrays.asList(OBJECT_MAPPER.writeValueAsString(o), OBJECT_MAPPER.writeValueAsString(value)).toArray());
}
else
{
flattenedResultSet.addRow(Arrays.asList(OBJECT_MAPPER.writeValueAsString(o), value).toArray());
}
}
}
return flattenedResultSet;
}
catch (Exception e)
{
throw new RuntimeException(e);
}
}
}
2 changes: 2 additions & 0 deletions ...stores/relational/connection/authentication/strategy/DefaultH2AuthenticationStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ private static List<String> getLegendH2ExtensionSQLs()
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_base64_decode FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions.legend_h2_extension_base64_decode\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_base64_encode FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions.legend_h2_extension_base64_encode\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_reverse_string FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions.legend_h2_extension_reverse_string\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_flatten_array FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions.legend_h2_extension_flatten_array\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_split_part FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions.legend_h2_extension_split_part\";"
);
}
Expand All @@ -121,6 +122,7 @@ private static List<String> getLegendH2_1_4_200_ExtensionSQLs()
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_reverse_string FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions_1_4_200.legend_h2_extension_reverse_string\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_hash_md5 FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions_1_4_200.legend_h2_extension_hash_md5\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_hash_sha1 FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions_1_4_200.legend_h2_extension_hash_sha1\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_flatten_array FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions_1_4_200.legend_h2_extension_flatten_array\";",
"CREATE ALIAS IF NOT EXISTS legend_h2_extension_split_part FOR \"org.finos.legend.engine.plan.execution.stores.relational.LegendH2Extensions_1_4_200.legend_h2_extension_split_part\";"
);
}
Expand Down
Loading

0 comments on commit 9feef9f

Please sign in to comment.