Skip to content

Commit

Permalink
SNOW-1348487 Fix Tests on Preprod (#106)
Browse files Browse the repository at this point in the history 
* fix tests on preprod

* check preprod
  • Loading branch information
@sfc-gh-bli
sfc-gh-bli authored May 15, 2024
1 parent 1e55932 commit 50d03e3
Show file tree
Hide file tree
Showing 5 changed files with 189 additions and 67 deletions.
201 changes: 135 additions & 66 deletions src/test/java/com/snowflake/snowpark_test/JavaRowSuite.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -339,81 +339,150 @@ public void testSpecialMap() {

@Test
public void testGetList() {
DataFrame df = getSession().sql("select [[1, 2], [3]]::ARRAY(ARRAY(NUMBER)) AS arr1");
StructType schema = df.schema();
assert schema.get(0).dataType() instanceof ArrayType;
assert ((ArrayType) schema.get(0).dataType()).getElementType() instanceof ArrayType;

List<?> list = df.collect()[0].getList(0);
assert list.size() == 2;

List<?> list1 = (List<?>) list.get(0);
List<?> list2 = (List<?>) list.get(1);

assert list1.size() == 2;
assert list2.size() == 1;

assert (Long) list1.get(0) == 1;
assert (Long) list1.get(1) == 2;
assert (Long) list2.get(0) == 3;
try {
if (isPreprodAccount()) {
getSession()
.sql("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true")
.show();
getSession()
.sql("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true")
.show();
getSession()
.sql("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true")
.show();
}
DataFrame df = getSession().sql("select [[1, 2], [3]]::ARRAY(ARRAY(NUMBER)) AS arr1");
StructType schema = df.schema();
assert schema.get(0).dataType() instanceof ArrayType;
assert ((ArrayType) schema.get(0).dataType()).getElementType() instanceof ArrayType;

List<?> list = df.collect()[0].getList(0);
assert list.size() == 2;

List<?> list1 = (List<?>) list.get(0);
List<?> list2 = (List<?>) list.get(1);

assert list1.size() == 2;
assert list2.size() == 1;

assert (Long) list1.get(0) == 1;
assert (Long) list1.get(1) == 2;
assert (Long) list2.get(0) == 3;
} finally {
if (isPreprodAccount()) {
getSession().sql("alter session unset ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE").show();
getSession()
.sql("alter session unset IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE")
.show();
getSession()
.sql("alter session unset FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT")
.show();
}
}
}

@Test
public void testGetMap() {
DataFrame df =
try {
if (isPreprodAccount()) {
getSession()
.sql(
"select {'1':{'a':1,'b':2},'2':{'c':3}} :: MAP(NUMBER, MAP(VARCHAR, NUMBER)) as map");
StructType schema = df.schema();
assert schema.get(0).dataType() instanceof MapType;
assert ((MapType) schema.get(0).dataType()).getKeyType() instanceof LongType;
assert ((MapType) schema.get(0).dataType()).getValueType() instanceof MapType;

Map<?, ?> map = df.collect()[0].getMap(0);
Map<?, ?> map1 = (Map<?, ?>) map.get(1L);
assert map1.size() == 2;
assert (Long) map1.get("a") == 1;
assert (Long) map1.get("b") == 2;

Map<?, ?> map2 = (Map<?, ?>) map.get(2L);
assert map2.size() == 1;
assert (Long) map2.get("c") == 3;
.sql("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true")
.show();
getSession()
.sql("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true")
.show();
getSession()
.sql("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true")
.show();
}
DataFrame df =
getSession()
.sql(
"select {'1':{'a':1,'b':2},'2':{'c':3}} :: MAP(NUMBER, MAP(VARCHAR, NUMBER)) as map");
StructType schema = df.schema();
assert schema.get(0).dataType() instanceof MapType;
assert ((MapType) schema.get(0).dataType()).getKeyType() instanceof LongType;
assert ((MapType) schema.get(0).dataType()).getValueType() instanceof MapType;

Map<?, ?> map = df.collect()[0].getMap(0);
Map<?, ?> map1 = (Map<?, ?>) map.get(1L);
assert map1.size() == 2;
assert (Long) map1.get("a") == 1;
assert (Long) map1.get("b") == 2;

Map<?, ?> map2 = (Map<?, ?>) map.get(2L);
assert map2.size() == 1;
assert (Long) map2.get("c") == 3;
} finally {
if (isPreprodAccount()) {
getSession().sql("alter session unset ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE").show();
getSession()
.sql("alter session unset IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE")
.show();
getSession()
.sql("alter session unset FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT")
.show();
}
}
}

@Test
public void testGetRow() {
DataFrame df =
try {
if (isPreprodAccount()) {
getSession()
.sql(
"select {'a': {'b': {'d':10,'c': 'txt'}}} :: OBJECT(a OBJECT(b OBJECT(c VARCHAR, d NUMBER))) as obj1");
StructType schema = df.schema();
schema.printTreeString();
assert schema.get(0).dataType() instanceof StructType;
assert schema.get(0).name().equals("OBJ1");
StructType sub1 = (StructType) schema.get(0).dataType();
assert sub1.size() == 1;
assert sub1.get(0).dataType() instanceof StructType;
assert sub1.get(0).name().equals("A");
StructType sub2 = (StructType) sub1.get(0).dataType();
assert sub2.size() == 1;
assert sub2.get(0).dataType() instanceof StructType;
assert sub2.get(0).name().equals("B");
StructType sub3 = (StructType) sub2.get(0).dataType();
assert sub3.size() == 2;
assert sub3.get(0).dataType() instanceof StringType;
assert sub3.get(0).name().equals("C");
assert sub3.get(1).dataType() instanceof LongType;
assert sub3.get(1).name().equals("D");

Row[] rows1 = df.collect();
assert rows1.length == 1;
Row row1 = rows1[0].getObject(0);
assert row1.size() == 1;
Row row2 = row1.getObject(0);
assert row2.size() == 1;
Row row3 = row2.getObject(0);
assert row3.size() == 2;
assert row3.getString(0).equals("txt");
assert row3.getLong(1) == 10;
.sql("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true")
.show();
getSession()
.sql("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true")
.show();
getSession()
.sql("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true")
.show();
}
DataFrame df =
getSession()
.sql(
"select {'a': {'b': {'d':10,'c': 'txt'}}} :: OBJECT(a OBJECT(b OBJECT(c VARCHAR, d NUMBER))) as obj1");
StructType schema = df.schema();
schema.printTreeString();
assert schema.get(0).dataType() instanceof StructType;
assert schema.get(0).name().equals("OBJ1");
StructType sub1 = (StructType) schema.get(0).dataType();
assert sub1.size() == 1;
assert sub1.get(0).dataType() instanceof StructType;
assert sub1.get(0).name().equals("A");
StructType sub2 = (StructType) sub1.get(0).dataType();
assert sub2.size() == 1;
assert sub2.get(0).dataType() instanceof StructType;
assert sub2.get(0).name().equals("B");
StructType sub3 = (StructType) sub2.get(0).dataType();
assert sub3.size() == 2;
assert sub3.get(0).dataType() instanceof StringType;
assert sub3.get(0).name().equals("C");
assert sub3.get(1).dataType() instanceof LongType;
assert sub3.get(1).name().equals("D");

Row[] rows1 = df.collect();
assert rows1.length == 1;
Row row1 = rows1[0].getObject(0);
assert row1.size() == 1;
Row row2 = row1.getObject(0);
assert row2.size() == 1;
Row row3 = row2.getObject(0);
assert row3.size() == 2;
assert row3.getString(0).equals("txt");
assert row3.getLong(1) == 10;
} finally {
if (isPreprodAccount()) {
getSession().sql("alter session unset ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE").show();
getSession()
.sql("alter session unset IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE")
.show();
getSession()
.sql("alter session unset FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT")
.show();
}
}
}
}
16 changes: 16 additions & 0 deletions src/test/java/com/snowflake/snowpark_test/TestBase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.snowflake.snowpark.TestUtils;
import com.snowflake.snowpark_java.JavaToScalaConvertor;
import com.snowflake.snowpark_java.Session;
import java.util.Optional;

public abstract class TestBase extends TestFunctions {

Expand All @@ -20,6 +21,21 @@ protected Session getSession() {
return _session;
}

Optional<Boolean> isPreprodAccount = Optional.empty();

protected boolean isPreprodAccount() {
if (!isPreprodAccount.isPresent()) {
isPreprodAccount =
Optional.of(
!getSession()
.sql("select current_account()")
.collect()[0]
.getString(0)
.contains("SFCTEST0"));
}
return isPreprodAccount.get();
}

protected void runQuery(String sql) {
getSession().sql(sql).collect();
}
Expand Down
15 changes: 14 additions & 1 deletion src/test/scala/com/snowflake/snowpark/APIInternalSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,23 @@ class APIInternalSuite extends TestData {
runQuery(s"CREATE TEMPORARY STAGE $tmpStageName", session)
// upload the file to stage
uploadFileToStage(tmpStageName, testFileCsv, compress = false)

if (isPreprodAccount) {
session.sql("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true").show()
session
.sql("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true")
.show()
session
.sql("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true")
.show()
}
}

override def afterAll: Unit = {
if (isPreprodAccount) {
session.sql("alter session unset ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE").show()
session.sql("alter session unset IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE").show()
session.sql("alter session unset FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT").show()
}
// drop the temporary stages
runQuery(s"DROP STAGE IF EXISTS $tmpStageName", session)

Expand Down
3 changes: 3 additions & 0 deletions src/test/scala/com/snowflake/snowpark/SNTestBase.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,9 @@ trait SNTestBase extends FunSuite with BeforeAndAfterAll with SFTestUtils with S
def runQueryReturnStatement(sql: String, sess: Session): Statement =
TestUtils.runQueryReturnStatement(sql, sess)

lazy val isPreprodAccount: Boolean =
!session.sql("select current_account()").collect().head.getString(0).contains("SFCTEST0")

case class TypeMap(colName: String, sfType: String, jdbcType: Int, tsType: DataType)
lazy val typeMap = List(
TypeMap("number", "number(10,2)", Types.DECIMAL, DecimalType(10, 2)),
Expand Down
21 changes: 21 additions & 0 deletions src/test/scala/com/snowflake/snowpark_test/DataTypeSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,27 @@ import java.util.TimeZone

// Test DataTypes out of com.snowflake.snowpark package.
class DataTypeSuite extends SNTestBase {
override def beforeAll: Unit = {
super.beforeAll
if (isPreprodAccount) {
session.sql("alter session set ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE=true").show()
session
.sql("alter session set IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE = true")
.show()
session
.sql("alter session set FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT=true")
.show()
}
}

override def afterAll: Unit = {
if (isPreprodAccount) {
session.sql("alter session unset ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE").show()
session.sql("alter session unset IGNORE_CLIENT_VESRION_IN_STRUCTURED_TYPES_RESPONSE").show()
session.sql("alter session unset FORCE_ENABLE_STRUCTURED_TYPES_NATIVE_ARROW_FORMAT").show()
}
super.afterAll
}
test("IntegralType") {
def verifyIntegralType(tpe: DataType): Unit = {
assert(TestUtils.isIntegralType(tpe))
Expand Down

0 comments on commit 50d03e3

Please sign in to comment.