Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix quoted column name for Iceberg primitive type #891

Merged
merged 1 commit into from
Nov 7, 2024
Merged

Fix quoted column name for Iceberg primitive type #891

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions src/main/java/net/snowflake/ingest/streaming/internal/ParquetRowBuffer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,14 +219,24 @@ public void setupSchema(List<ColumnMetadata> columns) {
isInRepeatedGroup = true;
}

boolean isPrimitiveColumn = path.length == 1;

/* set fieldId to 0 for non-structured columns */
int fieldId = path.length == 1 ? 0 : primitiveType.getId().intValue();
int fieldId = isPrimitiveColumn ? 0 : primitiveType.getId().intValue();
int ordinal = schema.getType(columnDescriptor.getPath()[0]).getId().intValue();

/**
* For non-structured columns, server side performs EP metadata check by columnsDisplayName,
* set it to the original column name to avoid EP validation error. Structured datatype is
* checked by fieldId and ordinal where columnDisplayName doesn't matter.
*/
String columnDisplayName =
isPrimitiveColumn ? columns.get(ordinal - 1).getName() : columnDotPath;

this.statsMap.put(
primitiveType.getId().toString(),
new RowBufferStats(
columnDotPath,
columnDisplayName,
null /* collationDefinitionString */,
ordinal,
fieldId,
Expand All @@ -239,7 +249,7 @@ public void setupSchema(List<ColumnMetadata> columns) {
this.tempStatsMap.put(
primitiveType.getId().toString(),
new RowBufferStats(
columnDotPath,
columnDisplayName,
null /* collationDefinitionString */,
ordinal,
fieldId,
Expand Down
20 changes: 12 additions & 8 deletions src/main/java/net/snowflake/ingest/utils/IcebergDataTypeParser.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ public static org.apache.parquet.schema.Type parseIcebergDataTypeStringToParquet
String name) {
Type icebergType = deserializeIcebergType(icebergDataType);
org.apache.parquet.schema.Type parquetType;
name = sanitizeFieldName(name);
if (icebergType.isPrimitiveType()) {
parquetType =
typeToMessageType.primitive(icebergType.asPrimitiveType(), repetition, id, name);
Expand Down Expand Up @@ -163,13 +164,7 @@ public static Type getTypeFromJson(@Nonnull JsonNode jsonNode) {

int id = JsonUtil.getInt(ID, field);

/* TypeToMessageType throws on empty field name, use a backslash to represent it and escape remaining backslash. */
String name =
JsonUtil.getString(NAME, field)
.replace(EMPTY_FIELD_CHAR, EMPTY_FIELD_CHAR + EMPTY_FIELD_CHAR);
if (name.isEmpty()) {
name = EMPTY_FIELD_CHAR;
}
String name = sanitizeFieldName(JsonUtil.getStringOrNull(NAME, field));
Type type = getTypeFromJson(field.get(TYPE));

String doc = JsonUtil.getStringOrNull(DOC, field);
Expand Down Expand Up @@ -280,7 +275,7 @@ private static org.apache.parquet.schema.Type replaceWithOriginalFieldName(
parquetFieldType,
icebergField.type(),
icebergField.name().equals(EMPTY_FIELD_CHAR)
? "" /* Empty string are encoded as single backslash in #structFromJson. Decode them here. */
? "" /* Empty string are encoded as single backslash in #sanitizeFieldName. Decode them here. */
: icebergField
.name()
.replace(EMPTY_FIELD_CHAR + EMPTY_FIELD_CHAR, EMPTY_FIELD_CHAR)));
Expand All @@ -293,4 +288,13 @@ private static org.apache.parquet.schema.Type replaceWithOriginalFieldName(
}
return builder.as(parquetType.getLogicalTypeAnnotation()).named(fieldName);
}

/* TypeToMessageType throws on empty field name, use a backslash to represent it and escape remaining backslash. */
private static String sanitizeFieldName(String fieldName) {
String name = fieldName.replace(EMPTY_FIELD_CHAR, EMPTY_FIELD_CHAR + EMPTY_FIELD_CHAR);
if (name.isEmpty()) {
name = EMPTY_FIELD_CHAR;
}
return name;
}
}
21 changes: 16 additions & 5 deletions src/test/java/net/snowflake/ingest/streaming/internal/datatypes/AbstractDataTypeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,15 +179,26 @@ protected String createIcebergTable(String dataType) throws SQLException {
conn.createStatement()
.execute(
String.format(
"create or replace iceberg table %s (%s string, %s %s) "
+ "catalog = 'SNOWFLAKE' "
+ "external_volume = 'streaming_ingest' "
+ "base_location = '%s';",
tableName, SOURCE_COLUMN_NAME, VALUE_COLUMN_NAME, dataType, baseLocation));
"create or replace iceberg table %s (%s string, %s %s) %s",
tableName,
SOURCE_COLUMN_NAME,
VALUE_COLUMN_NAME,
dataType,
baseLocation,
getIcebergTableConfig(tableName)));

return tableName;
}

protected String getIcebergTableConfig(String tableName) {
String baseLocation = String.format("SDK_IT/%s/%s", databaseName, tableName);
return String.format(
"catalog = 'SNOWFLAKE' "
+ "external_volume = 'streaming_ingest' "
+ "base_location = '%s';",
baseLocation);
}

protected String getRandomIdentifier() {
return String.format("test_%s", UUID.randomUUID()).replace('-', '_');
}
Expand Down
76 changes: 63 additions & 13 deletions src/test/java/net/snowflake/ingest/streaming/internal/it/ColumnNamesIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
/*
* Copyright (c) 2024 Snowflake Computing Inc. All rights reserved.
*/

package net.snowflake.ingest.streaming.internal.it;

import java.sql.ResultSet;
Expand All @@ -13,17 +17,45 @@
import net.snowflake.ingest.streaming.OpenChannelRequest;
import net.snowflake.ingest.streaming.SnowflakeStreamingIngestChannel;
import net.snowflake.ingest.streaming.internal.datatypes.AbstractDataTypeTest;
import net.snowflake.ingest.utils.Constants;
import net.snowflake.ingest.utils.SFException;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runners.Parameterized;

public class ColumnNamesIT extends AbstractDataTypeTest {
private static final int INGEST_VALUE = 1;

@Parameterized.Parameters(
name = "enableIcebergStreaming={0}, compressionAlgorithm={1}, icebergSerializationPolicy={2}")
public static Object[][] parameters() {
return new Object[][] {
// TODO: Enable this after Iceberg testing is set up on sfctest0 GCP / Azure
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With our plan to run two IT suites in parallel for iceberg/non-iceberg, how will this class be run in one mode in suite 1 and another mode in suite 2?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will split this into two separate tests in later PR.

// {true, "ZSTD", Constants.IcebergSerializationPolicy.COMPATIBLE},
// {true, "ZSTD", Constants.IcebergSerializationPolicy.OPTIMIZED},
{false, "GZIP", null},
{false, "ZSTD", null},
sfc-gh-alhuang marked this conversation as resolved.
Show resolved Hide resolved
};
}

@Parameterized.Parameter(0)
public static boolean enableIcebergStreaming;

@Parameterized.Parameter(1)
public static String compressionAlgorithm;

@Parameterized.Parameter(2)
public static Constants.IcebergSerializationPolicy icebergSerializationPolicy;

@Before
public void before() throws Exception {
super.before();
if (enableIcebergStreaming) {
super.beforeIceberg(compressionAlgorithm, icebergSerializationPolicy);
} else {
super.compressionAlgorithm = compressionAlgorithm;
super.before();
}
}

@Test
Expand Down Expand Up @@ -93,9 +125,11 @@ public void testNullableResolution() throws Exception {
conn.createStatement()
.execute(
String.format(
"create or replace table %s (AbC int, \"AbC\" int, \"abC\" int, ab\\ c int, \"Ab"
+ " c\" int);",
tableName));
"create or replace %s table %s (AbC int, \"AbC\" int, \"abC\" int, ab\\ c int, \"Ab"
+ " c\" int) %s",
enableIcebergStreaming ? "iceberg" : "",
tableName,
enableIcebergStreaming ? getIcebergTableConfig(tableName) : ""));
SnowflakeStreamingIngestChannel channel = openChannel(tableName);
String offsetToken = "token1";
channel.insertRow(new HashMap<>(), offsetToken);
Expand All @@ -119,7 +153,12 @@ public void testNullableResolution() throws Exception {
public void testExtraColNames() throws Exception {
String tableName = "t1";
conn.createStatement()
.execute(String.format("create or replace table %s (\"create\" int);", tableName));
.execute(
String.format(
"create or replace %s table %s (\"create\" int) %s;",
enableIcebergStreaming ? "iceberg" : "",
tableName,
enableIcebergStreaming ? getIcebergTableConfig(tableName) : ""));
SnowflakeStreamingIngestChannel channel =
openChannel(tableName, OpenChannelRequest.OnErrorOption.CONTINUE);

Expand Down Expand Up @@ -152,9 +191,11 @@ public void testMissingNotNullColNames() throws Exception {
conn.createStatement()
.execute(
String.format(
"create or replace table %s (\"CrEaTe\" int not null, a int not null, \"a\" int not"
+ " null, \"create\" int);",
tableName));
"create or replace %s table %s (\"CrEaTe\" int not null, a int not null, \"a\" int"
+ " not null, \"create\" int) %s;",
enableIcebergStreaming ? "iceberg" : "",
tableName,
enableIcebergStreaming ? getIcebergTableConfig(tableName) : ""));
SnowflakeStreamingIngestChannel channel =
openChannel(tableName, OpenChannelRequest.OnErrorOption.CONTINUE);

Expand All @@ -174,9 +215,11 @@ public void testNullValuesForNotNullColNames() throws Exception {
conn.createStatement()
.execute(
String.format(
"create or replace table %s (col1 int not null, a int not null, \"a\" int not"
+ " null, col2 int not null, \"col3\" int);",
tableName));
"create or replace %s table %s (col1 int not null, a int not null, \"a\" int not"
+ " null, col2 int not null, \"col3\" int) %s;",
enableIcebergStreaming ? "iceberg" : "",
tableName,
enableIcebergStreaming ? getIcebergTableConfig(tableName) : ""));
SnowflakeStreamingIngestChannel channel =
openChannel(tableName, OpenChannelRequest.OnErrorOption.CONTINUE);

Expand Down Expand Up @@ -222,7 +265,9 @@ private void testColumnNameSupported(String createTableColumnName, String ingest
}
Assert.assertEquals(1, count);

conn.createStatement().execute(String.format("alter table %s migrate;", tableName));
if (!enableIcebergStreaming) {
conn.createStatement().execute(String.format("alter table %s migrate;", tableName));
}
}

private void testColumnNameSupported(String column) throws SQLException, InterruptedException {
Expand Down Expand Up @@ -252,7 +297,12 @@ private void testInsertRowFails(
private String createSimpleTable(String createTableColumnName) throws SQLException {
String tableName = "a" + UUID.randomUUID().toString().replace("-", "_");
String createTableSql =
String.format("create table %s (%s int);", tableName, createTableColumnName);
String.format(
"create %s table %s (%s int) %s",
enableIcebergStreaming ? "iceberg" : "",
tableName,
createTableColumnName,
enableIcebergStreaming ? getIcebergTableConfig(tableName) : "");
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO we should refactor this to have a method that takes in the table name and column schema list only, and adds the proper prefix (iceberg table vs. table) and proper suffix (catalog and external_volume and base_location).

Instead of every create table call having to worry about both these things.

If this PR is good to merge am ok with this test refactoring in the next PR. Please also pick up any deferred comments from the last PR if any!

conn.createStatement().execute(createTableSql);
return tableName;
}
Expand Down
Loading