[FLINK-37649][datagen] Datagen connector cannot set length for collection type

* [FLINK-37649][datagen] Datagen connector cannot set length for collection type

Author: reswqa

Diff for: flink-table/flink-table-api-java-bridge/src/main/java/org/apache/flink/connector/datagen/table/RandomGeneratorVisitor.java

@@ -74,7 +74,7 @@ public class RandomGeneratorVisitor extends DataGenVisitorBase {
 
     public static final int RANDOM_BYTES_LENGTH_DEFAULT = 100;
 
-    private static final int RANDOM_COLLECTION_LENGTH_DEFAULT = 3;
+    public static final int RANDOM_COLLECTION_LENGTH_DEFAULT = 3;
 
     private static final float NULL_RATE_DEFAULT = 0f;
 
@@ -352,6 +352,7 @@ public DataGeneratorContainer visit(ArrayType arrayType) {
                 RandomGenerator.arrayGenerator(container.getGenerator(), config.get(lenOption));
         Set<ConfigOption<?>> options = container.getOptions();
         options.add(nr);
+        options.add(lenOption);
         return DataGeneratorContainer.of(
                 new DataGeneratorMapper<>(generator, (GenericArrayData::new), config.get(nr)),
                 options.toArray(new ConfigOption<?>[0]));
@@ -374,6 +375,7 @@ public DataGeneratorContainer visit(MultisetType multisetType) {
         Set<ConfigOption<?>> options = container.getOptions();
         ConfigOption<Float> nr = nullRate.floatType().defaultValue(NULL_RATE_DEFAULT);
         options.add(nr);
+        options.add(lenOption);
 
         return DataGeneratorContainer.of(
                 new DataGeneratorMapper<>(mapGenerator, GenericMapData::new, config.get(nr)),
@@ -397,6 +399,7 @@ public DataGeneratorContainer visit(MapType mapType) {
         Set<ConfigOption<?>> options = keyContainer.getOptions();
         options.addAll(valContainer.getOptions());
         options.add(nr);
+        options.add(lenOption);
 
         DataGenerator<Map<Object, Object>> mapGenerator =
                 RandomGenerator.mapGenerator(

Diff for: flink-table/flink-table-api-java-bridge/src/test/java/org/apache/flink/table/factories/DataGenTableSourceFactoryTest.java

@@ -56,6 +56,7 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.apache.flink.connector.datagen.table.RandomGeneratorVisitor.RANDOM_COLLECTION_LENGTH_DEFAULT;
 import static org.apache.flink.core.testutils.FlinkAssertions.anyCauseMatches;
 import static org.apache.flink.table.factories.utils.FactoryMocks.createTableSource;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -81,6 +82,11 @@ class DataGenTableSourceFactoryTest {
                     Column.physical("f2", DataTypes.VARCHAR(30)),
                     Column.physical("f3", DataTypes.VARBINARY(20)),
                     Column.physical("f4", DataTypes.STRING()));
+    private static final ResolvedSchema COLLECTION_SCHEMA =
+            ResolvedSchema.of(
+                    Column.physical("f0", DataTypes.ARRAY(DataTypes.STRING())),
+                    Column.physical("f1", DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())),
+                    Column.physical("f2", DataTypes.MULTISET(DataTypes.INT())));
 
     @Test
     void testDataTypeCoverage() throws Exception {
@@ -347,6 +353,43 @@ void testVariableLengthDataType() throws Exception {
                 "Custom length '21' for variable-length type (VARCHAR/STRING/VARBINARY/BYTES) field 'f3' should be shorter than '20' defined in the schema.");
     }
 
+    @Test
+    void testLengthForCollectionType() throws Exception {
+        DescriptorProperties descriptor = new DescriptorProperties();
+        final int rowsNumber = 200;
+        final int collectionSize = 10;
+        descriptor.putString(FactoryUtil.CONNECTOR.key(), "datagen");
+        descriptor.putLong(DataGenConnectorOptions.NUMBER_OF_ROWS.key(), rowsNumber);
+        // test for default length.
+        List<RowData> results = runGenerator(COLLECTION_SCHEMA, descriptor);
+        assertThat(results).hasSize(rowsNumber);
+        for (RowData row : results) {
+            assertThat(row.getArray(0).size()).isEqualTo(RANDOM_COLLECTION_LENGTH_DEFAULT);
+            assertThat(row.getMap(1).size())
+                    .isEqualTo(RandomGeneratorVisitor.RANDOM_COLLECTION_LENGTH_DEFAULT);
+            assertThat(row.getMap(2).size())
+                    .isEqualTo(RandomGeneratorVisitor.RANDOM_COLLECTION_LENGTH_DEFAULT);
+        }
+
+        // test for provided length.
+        descriptor.putLong(
+                DataGenConnectorOptionsUtil.FIELDS + ".f0." + DataGenConnectorOptionsUtil.LENGTH,
+                collectionSize);
+        descriptor.putLong(
+                DataGenConnectorOptionsUtil.FIELDS + ".f1." + DataGenConnectorOptionsUtil.LENGTH,
+                collectionSize);
+        descriptor.putLong(
+                DataGenConnectorOptionsUtil.FIELDS + ".f2." + DataGenConnectorOptionsUtil.LENGTH,
+                collectionSize);
+        results = runGenerator(COLLECTION_SCHEMA, descriptor);
+        assertThat(results).hasSize(rowsNumber);
+        for (RowData row : results) {
+            assertThat(row.getArray(0).size()).isEqualTo(collectionSize);
+            assertThat(row.getMap(1).size()).isEqualTo(collectionSize);
+            assertThat(row.getMap(2).size()).isEqualTo(collectionSize);
+        }
+    }
+
     @Test
     void testFixedLengthDataType() throws Exception {
         DescriptorProperties descriptor = new DescriptorProperties();