provideKeyGenParityArgs() {
+ // Covers every keygen + partition-formatter combo the dispatcher in
+ // HoodieDatasetBulkInsertHelper#buildKeygenColumns can land on. SimpleKeyGenerator cases
+ // exercise hive-style / url-encode / slash-separated date flags. Tier 2 covers only the
+ // {default, hive-style} subset for SimpleKeyGen; url-encode and slash-separated push to Tier 3.
+ return Stream.of(
+ // Tier 1
+ Arguments.of("nonpartitioned-tier1", NonpartitionedKeyGenerator.class.getName(), "_row_key", "", false, false, false),
+ // Tier 2
+ Arguments.of("simple-default-tier2", SimpleKeyGenerator.class.getName(), "_row_key", "partition", false, false, false),
+ Arguments.of("simple-hive-tier2", SimpleKeyGenerator.class.getName(), "_row_key", "partition", true, false, false),
+ // Tier 3 fallbacks: Simple under url-encode / slash-sep, plus every other keygen class.
+ Arguments.of("simple-slash-tier3", SimpleKeyGenerator.class.getName(), "_row_key", "partition", false, true, false),
+ Arguments.of("simple-hive-slash-tier3", SimpleKeyGenerator.class.getName(), "_row_key", "partition", true, true, false),
+ Arguments.of("simple-url-tier3", SimpleKeyGenerator.class.getName(), "_row_key", "partition", false, false, true),
+ Arguments.of("simple-hive-url-tier3", SimpleKeyGenerator.class.getName(), "_row_key", "partition", true, false, true),
+ Arguments.of("complex-single-tier3", ComplexKeyGenerator.class.getName(), "_row_key", "partition", false, false, false),
+ Arguments.of("complex-multi-tier3", ComplexKeyGenerator.class.getName(), "_row_key,ts", "partition,_hoodie_is_deleted", false, false, false),
+ Arguments.of("timestamp-based-tier3", TimestampBasedKeyGenerator.class.getName(), "_row_key", "ts", false, false, false),
+ Arguments.of("custom-tier3", CustomKeyGenerator.class.getName(), "_row_key", "partition:SIMPLE", false, false, false));
+ }
+
+ /**
+ * Asserts that record-key / partition-path values produced by {@link HoodieDatasetBulkInsertHelper}
+ * match those produced by the canonical Avro path ({@link BuiltinKeyGenerator#getRecordKey(GenericRecord)}).
+ *
+ * The Avro path is the ground truth shared by read- and write-side keygen invocations, so parity
+ * against it (rather than RDD-vs-UDF parity) is what actually matters for correctness.
+ */
+ @ParameterizedTest(name = "{0}")
+ @MethodSource("provideKeyGenParityArgs")
+ public void testKeyGenParityAgainstAvroGroundTruth(String label,
+ String keyGenClass,
+ String recordKeyFields,
+ String partitionPathFields,
+ boolean hiveStylePartitioning,
+ boolean slashSepPartitioning,
+ boolean urlEncodePartitioning) {
+ Map props = new HashMap<>();
+ props.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), keyGenClass);
+ props.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), recordKeyFields);
+ props.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), partitionPathFields);
+ props.put(HoodieWriteConfig.TBL_NAME.key(), label + "_parity_table");
+ props.put(KeyGeneratorOptions.HIVE_STYLE_PARTITIONING_ENABLE.key(), String.valueOf(hiveStylePartitioning));
+ props.put(KeyGeneratorOptions.SLASH_SEPARATED_DATE_PARTITIONING.key(), String.valueOf(slashSepPartitioning));
+ props.put(KeyGeneratorOptions.URL_ENCODE_PARTITIONING.key(), String.valueOf(urlEncodePartitioning));
+ if (keyGenClass.equals(TimestampBasedKeyGenerator.class.getName())) {
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_TYPE_FIELD.key(), "EPOCHMILLISECONDS");
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_OUTPUT_DATE_FORMAT.key(), "yyyy-MM-dd");
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_TIMEZONE_FORMAT.key(), "UTC");
+ }
+ HoodieWriteConfig config = getConfigBuilder(schemaStr).withProps(props).combineInput(false, false).build();
+
+ List rows = DataSourceTestUtils.generateRandomRows(20);
+ Dataset dataset = sqlContext.createDataFrame(rows, structType);
+ Dataset result = HoodieDatasetBulkInsertHelper.prepareForBulkInsert(dataset, config,
+ new HoodieTableConfig(), new NonSortPartitionerWithRows(), "000000001");
+
+ int recordKeyIdx = result.schema().fieldIndex(HoodieRecord.RECORD_KEY_METADATA_FIELD);
+ int partitionPathIdx = result.schema().fieldIndex(HoodieRecord.PARTITION_PATH_METADATA_FIELD);
+ int rowKeyIdx = result.schema().fieldIndex("_row_key");
+
+ Map actual = new HashMap<>();
+ for (Row r : result.collectAsList()) {
+ actual.put(r.getString(rowKeyIdx),
+ new String[] {r.getString(recordKeyIdx), r.getString(partitionPathIdx)});
+ }
+
+ TypedProperties keyGenProps = new TypedProperties();
+ keyGenProps.putAll(props);
+ BuiltinKeyGenerator groundTruthKeyGen =
+ (BuiltinKeyGenerator) org.apache.hudi.common.util.ReflectionUtils.loadClass(keyGenClass, keyGenProps);
+ scala.Function1 toAvro = AvroConversionUtils.createConverterToAvro(
+ structType, "trip", "example.schema");
+
+ assertEquals(rows.size(), actual.size(), "Row count mismatch — possible duplicate record keys");
+ for (Row inputRow : rows) {
+ GenericRecord avro = toAvro.apply(inputRow);
+ String expectedRecordKey = groundTruthKeyGen.getRecordKey(avro);
+ String expectedPartitionPath = groundTruthKeyGen.getPartitionPath(avro);
+ String[] observed = actual.get(inputRow.getString(0));
+ assertEquals(expectedRecordKey, observed[0],
+ "record key mismatch for keygen=" + label + " row=" + inputRow.getString(0));
+ assertEquals(expectedPartitionPath, observed[1],
+ "partition path mismatch for keygen=" + label + " row=" + inputRow.getString(0));
+ }
+ }
+
+ /**
+ * Tier 1 / Tier 2 fast paths use {@code col(field).cast(String)} for the record key. Confirms that
+ * a non-string record-key column (e.g. {@code ts: long}) is materialised as the string form of the
+ * underlying value, matching the canonical keygen output.
+ */
+ @Test
+ public void testFastPathCastsNonStringRecordKey() {
+ // Tier 1: Nonpartitioned, single record-key field. Using `ts` (long) as the record key.
+ Map nonpartitionedProps = new HashMap<>();
+ nonpartitionedProps.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), NonpartitionedKeyGenerator.class.getName());
+ nonpartitionedProps.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "ts");
+ nonpartitionedProps.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "");
+ nonpartitionedProps.put(HoodieWriteConfig.TBL_NAME.key(), "nonpartitioned_cast_tbl");
+ assertFastPathRecordKeyCast(nonpartitionedProps, /* expectedPartitionPath */ "");
+
+ // Tier 2: Simple, single record-key + partition-path. Using `ts` (long) as the record key.
+ Map simpleProps = new HashMap<>();
+ simpleProps.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), SimpleKeyGenerator.class.getName());
+ simpleProps.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "ts");
+ simpleProps.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "partition");
+ simpleProps.put(HoodieWriteConfig.TBL_NAME.key(), "simple_cast_tbl");
+ assertFastPathRecordKeyCast(simpleProps, /* expectedPartitionPath */ null);
+ }
+
+ private void assertFastPathRecordKeyCast(Map props, String expectedPartitionPath) {
+ HoodieWriteConfig config = getConfigBuilder(schemaStr).withProps(props).combineInput(false, false).build();
+ List rows = DataSourceTestUtils.generateRandomRows(10);
+ Dataset dataset = sqlContext.createDataFrame(rows, structType);
+ Dataset result = HoodieDatasetBulkInsertHelper.prepareForBulkInsert(dataset, config,
+ new HoodieTableConfig(), new NonSortPartitionerWithRows(), "000000001");
+
+ int recordKeyIdx = result.schema().fieldIndex(HoodieRecord.RECORD_KEY_METADATA_FIELD);
+ int partitionPathIdx = result.schema().fieldIndex(HoodieRecord.PARTITION_PATH_METADATA_FIELD);
+ int tsIdx = result.schema().fieldIndex("ts");
+ int partitionIdx = result.schema().fieldIndex("partition");
+
+ for (Row r : result.collectAsList()) {
+ Object tsVal = r.get(tsIdx);
+ assertEquals(String.valueOf(tsVal), r.getString(recordKeyIdx),
+ "record key should be string form of ts column");
+ String expected = expectedPartitionPath != null ? expectedPartitionPath : String.valueOf(r.get(partitionIdx));
+ assertEquals(expected, r.getString(partitionPathIdx),
+ "partition path mismatch");
+ }
+ }
+
+ /**
+ * Sanity check that {@link HoodieDatasetBulkInsertHelper} composes the Catalyst plan from
+ * pure column projections for Tier 1 and Tier 2 paths — no UDFs, no toRdd round-trip — so the
+ * fast paths actually benefit from Catalyst codegen. We verify this by inspecting the resulting
+ * Dataset's logical plan for absence of {@code ScalaUDF}/{@code UserDefinedFunction} nodes.
+ */
+ @Test
+ public void testFastPathAvoidsUdf() {
+ // Tier 1
+ Map tier1Props = new HashMap<>();
+ tier1Props.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), NonpartitionedKeyGenerator.class.getName());
+ tier1Props.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "_row_key");
+ tier1Props.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "");
+ tier1Props.put(HoodieWriteConfig.TBL_NAME.key(), "tier1_plan_tbl");
+ assertNoScalaUdfInPlan(tier1Props, "tier1");
+
+ // Tier 2
+ Map tier2Props = new HashMap<>();
+ tier2Props.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), SimpleKeyGenerator.class.getName());
+ tier2Props.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "_row_key");
+ tier2Props.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "partition");
+ tier2Props.put(HoodieWriteConfig.TBL_NAME.key(), "tier2_plan_tbl");
+ assertNoScalaUdfInPlan(tier2Props, "tier2");
+ }
+
+ private void assertNoScalaUdfInPlan(Map props, String label) {
+ HoodieWriteConfig config = getConfigBuilder(schemaStr).withProps(props).combineInput(false, false).build();
+ List rows = DataSourceTestUtils.generateRandomRows(5);
+ Dataset dataset = sqlContext.createDataFrame(rows, structType);
+ Dataset result = HoodieDatasetBulkInsertHelper.prepareForBulkInsert(dataset, config,
+ new HoodieTableConfig(), new NonSortPartitionerWithRows(), "000000001");
+ String plan = result.queryExecution().analyzed().toString();
+ assertTrue(!plan.toLowerCase().contains("scalaudf"),
+ label + " path should not contain ScalaUDF in its logical plan; plan was:\n" + plan);
+ }
+
+ /**
+ * The Tier 2 partition-path projection must substitute {@code __HIVE_DEFAULT_PARTITION__} for
+ * empty/null partition values, matching {@link org.apache.hudi.keygen.StringPartitionPathFormatter#handleEmpty}.
+ * Runs the case under both default and hive-style formatter flags, and against an injected
+ * row whose {@code partition} value is the empty string.
+ *
+ * The slash-separated formatter branch deliberately skips {@code handleEmpty} (per
+ * {@code PartitionPathFormatterBase#combine}), so we don't exercise that combo here -- the
+ * cross-flag parity test already covers it against the canonical Avro path.
+ */
+ @Test
+ public void testTier2EmptyPartitionValueSubstitutedWithHiveDefault() {
+ Row emptyPartitionRow = RowFactory.create("rk-empty", "", 1700000000000L, false);
+
+ Map defaultProps = new HashMap<>();
+ defaultProps.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), SimpleKeyGenerator.class.getName());
+ defaultProps.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "_row_key");
+ defaultProps.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "partition");
+ defaultProps.put(HoodieWriteConfig.TBL_NAME.key(), "empty_partition_default_tbl");
+ assertEmptyPartitionSubstituted(defaultProps, emptyPartitionRow,
+ PartitionPathEncodeUtils.DEFAULT_PARTITION_PATH);
+
+ Map hiveProps = new HashMap<>(defaultProps);
+ hiveProps.put(HoodieWriteConfig.TBL_NAME.key(), "empty_partition_hive_tbl");
+ hiveProps.put(KeyGeneratorOptions.HIVE_STYLE_PARTITIONING_ENABLE.key(), "true");
+ assertEmptyPartitionSubstituted(hiveProps, emptyPartitionRow,
+ "partition=" + PartitionPathEncodeUtils.DEFAULT_PARTITION_PATH);
+ }
+
+ private void assertEmptyPartitionSubstituted(Map props,
+ Row inputRow,
+ String expectedPartitionPath) {
+ HoodieWriteConfig config = getConfigBuilder(schemaStr).withProps(props).combineInput(false, false).build();
+ Dataset dataset = sqlContext.createDataFrame(java.util.Collections.singletonList(inputRow), structType);
+ Dataset result = HoodieDatasetBulkInsertHelper.prepareForBulkInsert(dataset, config,
+ new HoodieTableConfig(), new NonSortPartitionerWithRows(), "000000001");
+
+ Row out = result.collectAsList().get(0);
+ int partitionPathIdx = result.schema().fieldIndex(HoodieRecord.PARTITION_PATH_METADATA_FIELD);
+ assertEquals(expectedPartitionPath, out.getString(partitionPathIdx),
+ "empty partition value should be substituted with " + PartitionPathEncodeUtils.DEFAULT_PARTITION_PATH);
+ }
+
+ /**
+ * Guards against a UDF-path-only failure mode: the RDD path threaded the driver's {@link SQLConf}
+ * onto executor tasks via {@code injectSQLConf}; the UDF path does not. If any keygen reads
+ * {@code spark.sql.session.timeZone} during execution (e.g. when converting {@link java.sql.Timestamp}
+ * values through Catalyst), executor-local conf would silently override the driver's choice.
+ *
+ * This test forces {@code spark.sql.session.timeZone} to a non-default value, runs the helper, and
+ * asserts the resulting keys match the Avro ground truth computed under the same timezone. A divergence
+ * here would indicate the UDF path is sensitive to executor JVM defaults.
+ */
+ @Test
+ public void testUdfPathRespectsDriverSessionTimezone() {
+ String originalTz = sqlContext.getConf("spark.sql.session.timeZone", "UTC");
+ sqlContext.setConf("spark.sql.session.timeZone", "America/Los_Angeles");
+ try {
+ Map props = new HashMap<>();
+ props.put(DataSourceWriteOptions.KEYGENERATOR_CLASS_NAME().key(), TimestampBasedKeyGenerator.class.getName());
+ props.put(DataSourceWriteOptions.RECORDKEY_FIELD().key(), "_row_key");
+ props.put(DataSourceWriteOptions.PARTITIONPATH_FIELD().key(), "ts");
+ props.put(HoodieWriteConfig.TBL_NAME.key(), "tz_parity_table");
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_TYPE_FIELD.key(), "EPOCHMILLISECONDS");
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_OUTPUT_DATE_FORMAT.key(), "yyyy-MM-dd-HH");
+ props.put(TimestampKeyGeneratorConfig.TIMESTAMP_TIMEZONE_FORMAT.key(), "America/Los_Angeles");
+
+ HoodieWriteConfig config = getConfigBuilder(schemaStr).withProps(props).combineInput(false, false).build();
+ List rows = DataSourceTestUtils.generateRandomRows(10);
+ Dataset dataset = sqlContext.createDataFrame(rows, structType);
+ Dataset result = HoodieDatasetBulkInsertHelper.prepareForBulkInsert(dataset, config,
+ new HoodieTableConfig(), new NonSortPartitionerWithRows(), "000000001");
+
+ int partitionPathIdx = result.schema().fieldIndex(HoodieRecord.PARTITION_PATH_METADATA_FIELD);
+ int rowKeyIdx = result.schema().fieldIndex("_row_key");
+ Map actualPartition = new HashMap<>();
+ for (Row r : result.collectAsList()) {
+ actualPartition.put(r.getString(rowKeyIdx), r.getString(partitionPathIdx));
+ }
+
+ TypedProperties keyGenProps = new TypedProperties();
+ keyGenProps.putAll(props);
+ BuiltinKeyGenerator groundTruth =
+ (BuiltinKeyGenerator) org.apache.hudi.common.util.ReflectionUtils.loadClass(
+ TimestampBasedKeyGenerator.class.getName(), keyGenProps);
+ scala.Function1 toAvro = AvroConversionUtils.createConverterToAvro(
+ structType, "trip", "example.schema");
+
+ for (Row inputRow : rows) {
+ GenericRecord avro = toAvro.apply(inputRow);
+ assertEquals(groundTruth.getPartitionPath(avro), actualPartition.get(inputRow.getString(0)),
+ "partition path diverged for row " + inputRow.getString(0) + " under LA timezone");
+ }
+ } finally {
+ sqlContext.setConf("spark.sql.session.timeZone", originalTz);
+ }
+ }
+
class StageCheckBulkParallelismListener extends SparkListener {
private boolean checkFlag = false;