Option to have null defaults for avro sink schema (#21842)

This adds a way to specify an option NULL DEFAULTS in the CREATE SINK sql. If this is set, then the generated avro schema will have nullable fields with default value of null.
MaterializeInc · Sep 26, 2023 · e20ae46 · e20ae46
1 parent becfcb5
commit e20ae46
Show file tree

Hide file tree

Showing 11 changed files with 285 additions and 34 deletions.
diff --git a/src/interchange/src/avro.rs b/src/interchange/src/avro.rs
@@ -17,7 +17,7 @@ mod schema;
 pub use crate::avro::decode::{Decoder, DiffPair};
 pub use crate::avro::encode::{
     encode_datums_as_avro, encode_debezium_transaction_unchecked, get_debezium_transaction_schema,
-    AvroEncoder, AvroSchemaGenerator,
+    AvroEncoder, AvroSchemaGenerator, AvroSchemaOptions,
 };
 pub use crate::avro::envelope_cdc_v2 as cdc_v2;
 pub use crate::avro::schema::{parse_schema, schema_to_relationdesc, ConfluentAvroResolver};
@@ -163,7 +163,8 @@ mod tests {
         ];
         for (typ, datum, expected) in valid_pairings {
             let desc = RelationDesc::empty().with_column("column1", typ.nullable(false));
-            let schema_generator = AvroSchemaGenerator::new(None, None, None, desc, false).unwrap();
+            let schema_generator =
+                AvroSchemaGenerator::new(None, desc, Default::default()).unwrap();
             let avro_value =
                 encode_datums_as_avro(std::iter::once(datum), schema_generator.value_columns());
             assert_eq!(

diff --git a/src/interchange/src/avro/encode.rs b/src/interchange/src/avro/encode.rs
@@ -110,6 +110,18 @@ fn encode_message_unchecked(
     buf
 }
 
+#[derive(Debug, Default)]
+pub struct AvroSchemaOptions {
+    // Optional avro fullname on the generated key schema.
+    pub avro_key_fullname: Option<String>,
+    // Optional avro fullname on the generated value schema.
+    pub avro_value_fullname: Option<String>,
+    // Boolean flag to set null defaults for nullable types
+    pub set_null_defaults: bool,
+    // Boolean flag to indicate debezium envelope
+    pub is_debezium: bool,
+}
+
 /// Generates key and value Avro schemas
 pub struct AvroSchemaGenerator {
     value_columns: Vec<(ColumnName, ColumnType)>,
@@ -119,20 +131,24 @@ pub struct AvroSchemaGenerator {
 
 impl AvroSchemaGenerator {
     pub fn new(
-        key_fullname: Option<&str>,
-        value_fullname: Option<&str>,
         key_desc: Option<RelationDesc>,
         value_desc: RelationDesc,
-        debezium: bool,
+        AvroSchemaOptions {
+            is_debezium,
+            avro_value_fullname,
+            avro_key_fullname,
+            set_null_defaults,
+        }: AvroSchemaOptions,
     ) -> Result<Self, anyhow::Error> {
         let mut value_columns = column_names_and_types(value_desc);
-        if debezium {
+        if is_debezium {
             value_columns = envelopes::dbz_envelope(value_columns);
         }
         let row_schema = build_row_schema_json(
             &value_columns,
-            value_fullname.unwrap_or("envelope"),
+            avro_value_fullname.as_deref().unwrap_or("envelope"),
             &ENVELOPE_CUSTOM_NAMES,
+            set_null_defaults,
         )?;
         let writer_schema = Schema::parse(&row_schema).expect("valid schema constructed");
         let key_info = match key_desc {
@@ -141,8 +157,9 @@ impl AvroSchemaGenerator {
                 let columns = column_names_and_types(key_desc);
                 let row_schema = build_row_schema_json(
                     &columns,
-                    key_fullname.unwrap_or("row"),
+                    avro_key_fullname.as_deref().unwrap_or("row"),
                     &BTreeMap::new(),
+                    set_null_defaults,
                 )?;
                 Some(KeyInfo {
                     schema: Schema::parse(&row_schema).expect("valid schema constructed"),

diff --git a/src/interchange/src/avro/envelope_cdc_v2.rs b/src/interchange/src/avro/envelope_cdc_v2.rs
@@ -320,6 +320,7 @@ mod tests {
             &crate::encode::column_names_and_types(desc),
             "data",
             &BTreeMap::new(),
+            false,
         )
         .unwrap();
         let schema = build_schema(row_schema);

diff --git a/src/interchange/src/json.rs b/src/interchange/src/json.rs
@@ -77,7 +77,7 @@ impl fmt::Debug for JsonEncoder {
                 "schema",
                 &format!(
                     "{:?}",
-                    build_row_schema_json(&self.value_columns, "schema", &BTreeMap::new())
+                    build_row_schema_json(&self.value_columns, "schema", &BTreeMap::new(), false)
                 ),
             )
             .finish()
@@ -247,10 +247,11 @@ impl ToJson for TypedDatum<'_> {
     }
 }
 
-fn build_row_schema_field(
+fn build_row_schema_field_type(
     type_namer: &mut Namer,
     custom_names: &BTreeMap<GlobalId, String>,
     typ: &ColumnType,
+    set_null_defaults: bool,
 ) -> serde_json::Value {
     let mut field_type = match &typ.scalar_type {
         ScalarType::AclItem => json!("string"),
@@ -304,27 +305,29 @@ fn build_row_schema_field(
             "logicalType": "uuid",
         }),
         ty @ (ScalarType::Array(..) | ScalarType::Int2Vector | ScalarType::List { .. }) => {
-            let inner = build_row_schema_field(
+            let inner = build_row_schema_field_type(
                 type_namer,
                 custom_names,
                 &ColumnType {
                     nullable: true,
                     scalar_type: ty.unwrap_collection_element_type().clone(),
                 },
+                set_null_defaults,
             );
             json!({
                 "type": "array",
                 "items": inner
             })
         }
         ScalarType::Map { value_type, .. } => {
-            let inner = build_row_schema_field(
+            let inner = build_row_schema_field_type(
                 type_namer,
                 custom_names,
                 &ColumnType {
                     nullable: true,
                     scalar_type: (**value_type).clone(),
                 },
+                set_null_defaults,
             );
             json!({
                 "type": "map",
@@ -342,7 +345,8 @@ fn build_row_schema_field(
                 json!(name)
             } else {
                 let fields = fields.to_vec();
-                let json_fields = build_row_schema_fields(&fields, type_namer, custom_names);
+                let json_fields =
+                    build_row_schema_fields(&fields, type_namer, custom_names, set_null_defaults);
                 json!({
                     "type": "record",
                     "name": name,
@@ -368,6 +372,9 @@ fn build_row_schema_field(
         ScalarType::MzAclItem => json!("string"),
     };
     if typ.nullable {
+        // Should be revisited if we ever support a different kind of union scheme.
+        // Currently adding the "null" at the beginning means we can set the default
+        // value to "null" if such a preference is set.
         field_type = json!(["null", field_type]);
     }
     field_type
@@ -377,16 +384,32 @@ fn build_row_schema_fields(
     columns: &[(ColumnName, ColumnType)],
     type_namer: &mut Namer,
     custom_names: &BTreeMap<GlobalId, String>,
+    set_null_defaults: bool,
 ) -> Vec<serde_json::Value> {
     let mut fields = Vec::new();
     let mut field_namer = Namer::default();
     for (name, typ) in columns.iter() {
         let (name, _seen) = field_namer.valid_name(name.as_str());
-        let field_type = build_row_schema_field(type_namer, custom_names, typ);
-        fields.push(json!({
-            "name": name,
-            "type": field_type,
-        }));
+        let field_type =
+            build_row_schema_field_type(type_namer, custom_names, typ, set_null_defaults);
+
+        // It's a nullable union if the type is an array and the first option is "null"
+        let is_nullable_union = field_type
+            .as_array()
+            .is_some_and(|array| array.first().is_some_and(|first| first == &json!("null")));
+
+        if set_null_defaults && is_nullable_union {
+            fields.push(json!({
+                "name": name,
+                "type": field_type,
+                "default": null,
+            }));
+        } else {
+            fields.push(json!({
+                "name": name,
+                "type": field_type,
+            }));
+        }
     }
     fields
 }
@@ -396,8 +419,14 @@ pub fn build_row_schema_json(
     columns: &[(ColumnName, ColumnType)],
     name: &str,
     custom_names: &BTreeMap<GlobalId, String>,
+    set_null_defaults: bool,
 ) -> Result<serde_json::Value, anyhow::Error> {
-    let fields = build_row_schema_fields(columns, &mut Namer::default(), custom_names);
+    let fields = build_row_schema_fields(
+        columns,
+        &mut Namer::default(),
+        custom_names,
+        set_null_defaults,
+    );
     let _ = mz_avro::schema::Name::parse_simple(name)?;
     Ok(json!({
         "type": "record",

diff --git a/src/sql-lexer/src/keywords.txt b/src/sql-lexer/src/keywords.txt
@@ -108,6 +108,7 @@ Decimal
 Declare
 Decorrelated
 Default
+Defaults
 Delete
 Delimited
 Delimiter

diff --git a/src/sql-parser/src/ast/defs/ddl.rs b/src/sql-parser/src/ast/defs/ddl.rs
@@ -139,13 +139,15 @@ impl_display_t!(ProtobufSchema);
 pub enum CsrConfigOptionName {
     AvroKeyFullname,
     AvroValueFullname,
+    NullDefaults,
 }
 
 impl AstDisplay for CsrConfigOptionName {
     fn fmt<W: fmt::Write>(&self, f: &mut AstFormatter<W>) {
         f.write_str(match self {
             CsrConfigOptionName::AvroKeyFullname => "AVRO KEY FULLNAME",
             CsrConfigOptionName::AvroValueFullname => "AVRO VALUE FULLNAME",
+            CsrConfigOptionName::NullDefaults => "NULL DEFAULTS",
         })
     }
 }

diff --git a/src/sql-parser/src/parser.rs b/src/sql-parser/src/parser.rs
@@ -2034,7 +2034,7 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_csr_config_option(&mut self) -> Result<CsrConfigOption<Raw>, ParserError> {
-        let name = match self.expect_one_of_keywords(&[AVRO])? {
+        let name = match self.expect_one_of_keywords(&[AVRO, NULL])? {
             AVRO => {
                 let name = match self.expect_one_of_keywords(&[KEY, VALUE])? {
                     KEY => CsrConfigOptionName::AvroKeyFullname,
@@ -2044,6 +2044,10 @@ impl<'a> Parser<'a> {
                 self.expect_keyword(FULLNAME)?;
                 name
             }
+            NULL => {
+                self.expect_keyword(DEFAULTS)?;
+                CsrConfigOptionName::NullDefaults
+            }
             _ => unreachable!(),
         };
         Ok(CsrConfigOption {

diff --git a/src/sql-parser/tests/testdata/ddl b/src/sql-parser/tests/testdata/ddl
@@ -536,6 +536,41 @@ CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC = 'topic') FORMAT BYTE
 =>
 CreateSink(CreateSinkStatement { name: UnresolvedItemName([Ident("foo")]), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: KafkaConnection { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaConfigOption { name: Topic, value: Some(Value(String("topic"))) }] }, key: None }, format: Some(Bytes), envelope: None, with_options: [CreateSinkOption { name: Size, value: Some(Value(String("xlarge"))) }, CreateSinkOption { name: Snapshot, value: Some(Value(Boolean(true))) }] })
 
+parse-statement
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS = TRUE) ENVELOPE UPSERT
+----
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC = 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS = true) ENVELOPE UPSERT
+=>
+CreateSink(CreateSinkStatement { name: UnresolvedItemName([Ident("foo")]), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: KafkaConnection { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaConfigOption { name: Topic, value: Some(Value(String("topic"))) }] }, key: None }, format: Some(Avro(Csr { csr_connection: CsrConnectionAvro { connection: CsrConnection { connection: Name(UnresolvedItemName([Ident("conn2")])), options: [CsrConfigOption { name: NullDefaults, value: Some(Value(Boolean(true))) }] }, key_strategy: None, value_strategy: None, seed: None } })), envelope: Some(Upsert), with_options: [] })
+
+parse-statement
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS TRUE) ENVELOPE UPSERT
+----
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC = 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS = true) ENVELOPE UPSERT
+=>
+CreateSink(CreateSinkStatement { name: UnresolvedItemName([Ident("foo")]), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: KafkaConnection { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaConfigOption { name: Topic, value: Some(Value(String("topic"))) }] }, key: None }, format: Some(Avro(Csr { csr_connection: CsrConnectionAvro { connection: CsrConnection { connection: Name(UnresolvedItemName([Ident("conn2")])), options: [CsrConfigOption { name: NullDefaults, value: Some(Value(Boolean(true))) }] }, key_strategy: None, value_strategy: None, seed: None } })), envelope: Some(Upsert), with_options: [] })
+
+parse-statement
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS) ENVELOPE UPSERT
+----
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC = 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS) ENVELOPE UPSERT
+=>
+CreateSink(CreateSinkStatement { name: UnresolvedItemName([Ident("foo")]), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: KafkaConnection { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaConfigOption { name: Topic, value: Some(Value(String("topic"))) }] }, key: None }, format: Some(Avro(Csr { csr_connection: CsrConnectionAvro { connection: CsrConnection { connection: Name(UnresolvedItemName([Ident("conn2")])), options: [CsrConfigOption { name: NullDefaults, value: None }] }, key_strategy: None, value_strategy: None, seed: None } })), envelope: Some(Upsert), with_options: [] })
+
+parse-statement
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL VALUES) ENVELOPE UPSERT
+----
+error: Expected DEFAULTS, found VALUES
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL VALUES) ENVELOPE UPSERT
+                                                                                                                                      ^
+
+parse-statement
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS = FALSE) ENVELOPE UPSERT
+----
+CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC = 'topic') FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION conn2 (NULL DEFAULTS = false) ENVELOPE UPSERT
+=>
+CreateSink(CreateSinkStatement { name: UnresolvedItemName([Ident("foo")]), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: KafkaConnection { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaConfigOption { name: Topic, value: Some(Value(String("topic"))) }] }, key: None }, format: Some(Avro(Csr { csr_connection: CsrConnectionAvro { connection: CsrConnection { connection: Name(UnresolvedItemName([Ident("conn2")])), options: [CsrConfigOption { name: NullDefaults, value: Some(Value(Boolean(false))) }] }, key_strategy: None, value_strategy: None, seed: None } })), envelope: Some(Upsert), with_options: [] })
+
 parse-statement
 CREATE INDEX foo ON myschema.bar (a, b)
 ----

diff --git a/src/sql/src/plan/statement/ddl.rs b/src/sql/src/plan/statement/ddl.rs
@@ -19,7 +19,7 @@ use std::iter;
 use itertools::Itertools;
 use mz_controller_types::{ClusterId, ReplicaId, DEFAULT_REPLICA_LOGGING_INTERVAL_MICROS};
 use mz_expr::CollectionPlan;
-use mz_interchange::avro::AvroSchemaGenerator;
+use mz_interchange::avro::{AvroSchemaGenerator, AvroSchemaOptions};
 use mz_ore::cast::{self, CastFrom, TryCastFrom};
 use mz_ore::collections::HashSet;
 use mz_ore::str::StrExt;
@@ -2344,7 +2344,8 @@ fn key_constraint_err(desc: &RelationDesc, user_keys: &[ColumnName]) -> PlanErro
 generate_extracted_config!(
     CsrConfigOption,
     (AvroKeyFullname, String),
-    (AvroValueFullname, String)
+    (AvroValueFullname, String),
+    (NullDefaults, bool, Default(false))
 );
 
 fn kafka_sink_builder(
@@ -2444,6 +2445,7 @@ fn kafka_sink_builder(
             let CsrConfigOptionExtracted {
                 avro_key_fullname,
                 avro_value_fullname,
+                null_defaults,
                 ..
             } = options.try_into()?;
 
@@ -2457,14 +2459,19 @@ fn kafka_sink_builder(
                 sql_bail!("Must specify both AVRO KEY FULLNAME and AVRO VALUE FULLNAME when specifying generated schema names");
             }
 
+            let options = AvroSchemaOptions {
+                avro_key_fullname,
+                avro_value_fullname,
+                set_null_defaults: null_defaults,
+                is_debezium: matches!(envelope, SinkEnvelope::Debezium),
+            };
+
             let schema_generator = AvroSchemaGenerator::new(
-                avro_key_fullname.as_deref(),
-                avro_value_fullname.as_deref(),
                 key_desc_and_indices
                     .as_ref()
                     .map(|(desc, _indices)| desc.clone()),
                 value_desc.clone(),
-                matches!(envelope, SinkEnvelope::Debezium),
+                options,
             )?;
             let value_schema = schema_generator.value_writer_schema().to_string();
             let key_schema = schema_generator

diff --git a/src/storage/src/sink/kafka.rs b/src/storage/src/sink/kafka.rs
@@ -21,7 +21,7 @@ use anyhow::{anyhow, bail, Context};
 use differential_dataflow::{Collection, Hashable};
 use futures::{StreamExt, TryFutureExt};
 use maplit::btreemap;
-use mz_interchange::avro::{AvroEncoder, AvroSchemaGenerator};
+use mz_interchange::avro::{AvroEncoder, AvroSchemaGenerator, AvroSchemaOptions};
 use mz_interchange::encode::Encode;
 use mz_interchange::json::JsonEncoder;
 use mz_kafka_util::client::{
@@ -1007,14 +1007,14 @@ where
             key_schema_id,
             value_schema_id,
         }) => {
-            let schema_generator = AvroSchemaGenerator::new(
-                None,
-                None,
-                key_desc,
-                value_desc,
-                matches!(envelope, Some(SinkEnvelope::Debezium)),
-            )
-            .expect("avro schema validated");
+            let options = AvroSchemaOptions {
+                avro_key_fullname: None,
+                avro_value_fullname: None,
+                set_null_defaults: false,
+                is_debezium: matches!(envelope, Some(SinkEnvelope::Debezium)),
+            };
+            let schema_generator = AvroSchemaGenerator::new(key_desc, value_desc, options)
+                .expect("avro schema validated");
             let encoder = AvroEncoder::new(schema_generator, key_schema_id, value_schema_id);
             encode_stream(
                 stream,