partitioning improvement

if fetchNextRealKey does not find a valid key; try fetchPrevRealKey before deeming the partition as unsplittable. This is useful for narrow ranges

Author: jgao54

flow/connectors/mysql/qrep_partition.go

@@ -184,6 +184,9 @@ type rangeProber interface {
 	fetchNextRealKey(
 		ctx context.Context, tableName string, quotedCol string, midpoint string, start string, end string,
 	) (string, bool, error)
+	fetchPrevRealKey(
+		ctx context.Context, tableName string, quotedCol string, midpoint string, start string, end string,
+	) (string, bool, error)
 }
 
 // buildAdaptiveStringPartitions splits an arbitrary string watermark column
@@ -225,6 +228,15 @@ func buildAdaptiveStringPartitions(
 		if err != nil {
 			return nil, fmt.Errorf("failed to fetch next real key for [%s, %s]: %w", p.start, p.end, err)
 		}
+		if !found {
+			// The interpolated midpoint can overshoot every key in the range when
+			// keys occupy a narrow slice of the character space. So also probe
+			// backwards before declaring the partition unsplittable.
+			k, found, err = prober.fetchPrevRealKey(ctx, tableName, quotedCol, mid, p.start, p.end)
+			if err != nil {
+				return nil, fmt.Errorf("failed to fetch prev real key for [%s, %s]: %w", p.start, p.end, err)
+			}
+		}
 		if !found {
 			outputs = append(outputs, p)
 			continue
@@ -267,6 +279,23 @@ func (c *MySqlConnector) fetchNextRealKey(
 	query := fmt.Sprintf(
 		"SELECT %[1]s FROM %[2]s WHERE %[1]s >= '%[3]s' AND %[1]s > '%[4]s' AND %[1]s < '%[5]s' ORDER BY %[1]s LIMIT 1",
 		quotedCol, tableName, mysql.Escape(midpoint), mysql.Escape(start), mysql.Escape(end))
+	return c.fetchRealKey(ctx, query)
+}
+
+// fetchPrevRealKey returns the largest real value of the watermark column that
+// is below the interpolated midpoint and strictly inside (start, end). The
+// bounds are enforced by MySQL using its column's collation. Return false when
+// no such key exists.
+func (c *MySqlConnector) fetchPrevRealKey(
+	ctx context.Context, tableName string, quotedCol string, midpoint string, start string, end string,
+) (string, bool, error) {
+	query := fmt.Sprintf(
+		"SELECT %[1]s FROM %[2]s WHERE %[1]s < '%[3]s' AND %[1]s > '%[4]s' AND %[1]s < '%[5]s' ORDER BY %[1]s DESC LIMIT 1",
+		quotedCol, tableName, mysql.Escape(midpoint), mysql.Escape(start), mysql.Escape(end))
+	return c.fetchRealKey(ctx, query)
+}
+
+func (c *MySqlConnector) fetchRealKey(ctx context.Context, query string) (string, bool, error) {
 	rs, err := c.Execute(ctx, query)
 	if err != nil {
 		return "", false, err
@@ -277,7 +306,7 @@ func (c *MySqlConnector) fetchNextRealKey(
 	}
 	key, err := rs.GetString(0, 0)
 	if err != nil {
-		return "", false, fmt.Errorf("failed to read next real key: %w", err)
+		return "", false, fmt.Errorf("failed to read real key: %w", err)
 	}
 	// GetString is zero-copy over the resultset's row buffer, which rs.Close()
 	// recycles into a pool where the next query's response overwrites it; the

flow/connectors/mysql/qrep_partition_test.go

@@ -263,7 +263,26 @@ func (f *fakeRangeProber) fetchNextRealKey(
 	_ context.Context, _ string, _ string, midpoint string, start string, end string,
 ) (string, bool, error) {
 	for _, k := range f.keys {
-		if f.less(k, midpoint) {
+		if f.less(k, midpoint) { // k < midpoint
+			continue
+		}
+		if !f.less(start, k) { // k <= start
+			continue
+		}
+		if !f.less(k, end) { // k >= end
+			continue
+		}
+		return k, true, nil
+	}
+	return "", false, nil
+}
+
+func (f *fakeRangeProber) fetchPrevRealKey(
+	_ context.Context, _ string, _ string, midpoint string, start string, end string,
+) (string, bool, error) {
+	for i := len(f.keys) - 1; i >= 0; i-- {
+		k := f.keys[i]
+		if !f.less(k, midpoint) { // k >= midpoint
 			continue
 		}
 		if !f.less(start, k) { // k <= start
@@ -386,17 +405,19 @@ func TestBuildAdaptiveStringPartitions_CaseInsensitiveCollation(t *testing.T) {
 	verifyFullCoverage(t, partitions, keys, caseInsensitiveLess)
 }
 
-func TestBuildAdaptiveStringPartitions_NarrowAlphabetLimitation(t *testing.T) {
+func TestBuildAdaptiveStringPartitions_NarrowRange(t *testing.T) {
 	keys := make([]string, 0, 100)
 	for i := 1; i <= 100; i++ {
 		keys = append(keys, fmt.Sprintf("key_%04d", i))
 	}
 
-	// demonstrate midpoint is outside min/max range so fetchNextRealKey always returns false
+	// demonstrate fetchNextRealKey would not return a valid key
+	// (e.g. "key_00_" sorts after "key_0099") so splitting these
+	// keys relies on the fetchPrevRealKey fallback
 	require.Equal(t, "key_00_`", stringMidpoint("key_0001", "key_0100"))
 
 	partitions := runAdaptiveStringPartitions(t, keys, binaryLess, 8)
-	require.Len(t, partitions, 1)
+	require.Len(t, partitions, 8)
 	verifyFullCoverage(t, partitions, keys, binaryLess)
 }