feat: support add all null column as metadata-only operation via sql (#3504)

Adds support for adding all-null column via SQL.

If the user passes:
```rs
dataset.add_column(NewColumnTransform::SqlExpressions(vec!["new_col", "CAST(NULL AS int)"]);
```

We'll discover that the intention is to to create an all null column,
and optimize the transform to:
```rs
dataset.add_column(NewColumnTransform::AllNull(Arc::new(
  Schema::new(vec![
    Field::new("new_col", DataType:Int32, true),
  ])
)
```

The motivation here is to be able to expose the capability to add the
all null column as a metadata-only operation through the LanceDB SDKs.
Currently these methods only support passing SQL expressions. A
different option would have been to modify the arguments to the python
table.add_column & typescript table.addColumn, but that seemed like more
work so I wanted to propose this solution first.

Author: albertlockett

python/python/tests/test_schema_evolution.py

@@ -512,3 +512,31 @@ def some_udf(batch):
 
     with pytest.raises(ValueError, match="A checkpoint file cannot be used"):
         frag.merge_columns(some_udf, columns=["a"])
+
+
+def test_add_cols_all_null_with_sql(tmp_path: Path):
+    tab = pa.table(
+        {
+            "a": range(100),
+        }
+    )
+    dataset = lance.write_dataset(
+        tab, tmp_path, max_rows_per_file=50, data_storage_version="stable"
+    )
+    fragments_before = dataset.get_fragments()
+    dataset.add_columns({"b": "CAST(NULL AS INT)"})
+    fragments_after = dataset.get_fragments()
+
+    # assert this was a metadata only operation and no data was written
+    assert len(fragments_before) == len(fragments_after)
+    for frag_before, frag_after in zip(fragments_before, fragments_after):
+        assert frag_before.fragment_id == frag_after.fragment_id
+        assert frag_before.data_files() == frag_after.data_files()
+
+    # assert the schema is as expected
+    assert dataset.schema == pa.schema(
+        {
+            "a": pa.int64(),
+            "b": pa.int32(),
+        }
+    )

rust/lance/src/dataset/schema_evolution.rs

@@ -13,7 +13,6 @@ use futures::stream::{StreamExt, TryStreamExt};
 use lance_arrow::SchemaExt;
 use lance_core::datatypes::{Field, Schema};
 use lance_datafusion::utils::StreamingWriteSource;
-use lance_encoding::version::LanceFileVersion;
 use lance_table::format::Fragment;
 use snafu::location;
 
@@ -23,6 +22,12 @@ use super::{
     Dataset,
 };
 
+mod optimize;
+
+use optimize::{
+    ChainedNewColumnTransformOptimizer, NewColumnTransformOptimizer, SqlToAllNullsOptimizer,
+};
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct BatchInfo {
     pub fragment_id: u32,
@@ -149,6 +154,14 @@ pub(super) async fn add_columns_to_fragments(
         Ok(())
     };
 
+    // Optimize the transforms
+    let mut optimizer = ChainedNewColumnTransformOptimizer::new(vec![]);
+    // ALlNull transform can not performed on legacy files
+    if !dataset.is_legacy_storage() {
+        optimizer.add_optimizer(Box::new(SqlToAllNullsOptimizer::new()));
+    }
+    let transforms = optimizer.optimize(dataset, transforms)?;
+
     let (output_schema, fragments) = match transforms {
         NewColumnTransform::BatchUDF(udf) => {
             check_names(udf.output_schema.as_ref())?;
@@ -262,17 +275,7 @@ pub(super) async fn add_columns_to_fragments(
             // can't add all-null columns as a metadata-only operation. The reason is because we
             // use the NullReader for fragments that have missing columns and we can't mix legacy
             // and non-legacy readers when reading the fragment.
-            if fragments.iter().any(|fragment| {
-                fragment.files.iter().any(|file| {
-                    matches!(
-                        LanceFileVersion::try_from_major_minor(
-                            file.file_major_version,
-                            file.file_minor_version
-                        ),
-                        Ok(LanceFileVersion::Legacy)
-                    )
-                })
-            }) {
+            if dataset.is_legacy_storage() {
                 return Err(Error::NotSupported {
                     source: "Cannot add all-null columns to legacy dataset version.".into(),
                     location: location!(),
@@ -1744,4 +1747,115 @@ mod test {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_new_column_sql_to_all_nulls_transform_optimizer() {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "a",
+            DataType::Int32,
+            false,
+        )]));
+
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from_iter(0..100))],
+        )
+        .unwrap();
+        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
+        let test_dir = tempfile::tempdir().unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+        let mut dataset = Dataset::write(
+            reader,
+            test_uri,
+            Some(WriteParams {
+                max_rows_per_file: 50,
+                max_rows_per_group: 25,
+                data_storage_version: Some(LanceFileVersion::Stable),
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+        dataset.validate().await.unwrap();
+
+        let manifest_before = dataset.manifest.clone();
+
+        // Add all null column
+        dataset
+            .add_columns(
+                NewColumnTransform::SqlExpressions(vec![(
+                    "b".to_string(),
+                    "CAST(NULL AS int)".to_string(),
+                )]),
+                None,
+                None,
+            )
+            .await
+            .unwrap();
+        let manifest_after = dataset.manifest.clone();
+
+        // Check that this is a metadata-only operation (the fragments don't change)
+        assert_eq!(&manifest_before.fragments, &manifest_after.fragments);
+
+        // check that the new field was added to the schema
+        let expected_schema = ArrowSchema::new(vec![
+            ArrowField::new("a", DataType::Int32, false),
+            ArrowField::new("b", DataType::Int32, true),
+        ]);
+        assert_eq!(ArrowSchema::from(dataset.schema()), expected_schema);
+    }
+
+    #[tokio::test]
+    async fn test_new_column_sql_to_all_nulls_transform_optimizer_legacy() {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "a",
+            DataType::Int32,
+            false,
+        )]));
+
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from_iter(0..100))],
+        )
+        .unwrap();
+        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
+        let test_dir = tempfile::tempdir().unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+        let mut dataset = Dataset::write(
+            reader,
+            test_uri,
+            Some(WriteParams {
+                max_rows_per_file: 50,
+                max_rows_per_group: 25,
+                data_storage_version: Some(LanceFileVersion::Legacy),
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+        dataset.validate().await.unwrap();
+
+        // Add all null column ...
+        // This is basically a smoke test to ensure we don't try to use the all-nulls
+        // transform optimizer where it's not supported, and then blow up when we try
+        // to apply the transform
+        dataset
+            .add_columns(
+                NewColumnTransform::SqlExpressions(vec![(
+                    "b".to_string(),
+                    "CAST(NULL AS int)".to_string(),
+                )]),
+                None,
+                None,
+            )
+            .await
+            .unwrap();
+
+        // check that the new field was added to the schema
+        let expected_schema = ArrowSchema::new(vec![
+            ArrowField::new("a", DataType::Int32, false),
+            ArrowField::new("b", DataType::Int32, true),
+        ]);
+        assert_eq!(ArrowSchema::from(dataset.schema()), expected_schema);
+    }
 }

rust/lance/src/dataset/schema_evolution/optimize.rs

@@ -0,0 +1,153 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::sync::Arc;
+
+use arrow_schema::{DataType, Field, Schema};
+use datafusion::prelude::Expr;
+use datafusion::scalar::ScalarValue;
+use lance_datafusion::planner::Planner;
+
+use crate::error::Result;
+use crate::Dataset;
+
+use super::NewColumnTransform;
+
+/// Optimizes a `NewColumnTransform` into
+pub(super) trait NewColumnTransformOptimizer: Send + Sync {
+    /// Optimize the passed `NewColumnTransform` to a more efficient form.
+    fn optimize(
+        &self,
+        dataset: &Dataset,
+        transform: NewColumnTransform,
+    ) -> Result<NewColumnTransform>;
+}
+
+/// A `NewColumnTransformOptimizer` that chains multiple `NewColumnTransformOptimizer`s together.
+pub(super) struct ChainedNewColumnTransformOptimizer {
+    optimizers: Vec<Box<dyn NewColumnTransformOptimizer>>,
+}
+
+impl ChainedNewColumnTransformOptimizer {
+    pub(super) fn new(optimizers: Vec<Box<dyn NewColumnTransformOptimizer>>) -> Self {
+        Self { optimizers }
+    }
+
+    pub(super) fn add_optimizer(&mut self, optimizer: Box<dyn NewColumnTransformOptimizer>) {
+        self.optimizers.push(optimizer);
+    }
+}
+
+/// A `NewColumnTransformOptimizer` that chains multiple `NewColumnTransformOptimizer`s together.
+impl NewColumnTransformOptimizer for ChainedNewColumnTransformOptimizer {
+    fn optimize(
+        &self,
+        dataset: &Dataset,
+        transform: NewColumnTransform,
+    ) -> Result<NewColumnTransform> {
+        let mut transform = transform;
+        for optimizer in &self.optimizers {
+            transform = optimizer.optimize(dataset, transform)?;
+        }
+        Ok(transform)
+    }
+}
+
+/// Optimizes a `NewColumnTransform` that is a SQL expression to a `NewColumnTransform::AllNulls` if
+/// the SQL expression is "NULL". For example
+/// `NewColumnTransform::SqlExpression(vec![("new_col", "CAST(NULL AS int)"])`
+/// would be optimized to
+/// `NewColumnTransform::AllNulls(Schema::new(vec![Field::new("new_col", DataType::Int)]))`.
+///
+pub(super) struct SqlToAllNullsOptimizer;
+
+impl SqlToAllNullsOptimizer {
+    pub(super) fn new() -> Self {
+        Self
+    }
+
+    fn is_all_null(&self, expr: &Expr) -> AllNullsResult {
+        match expr {
+            Expr::Cast(cast) => {
+                if matches!(cast.expr.as_ref(), Expr::Literal(ScalarValue::Null)) {
+                    let data_type = cast.data_type.clone();
+                    AllNullsResult::AllNulls(data_type)
+                } else {
+                    AllNullsResult::NotAllNulls
+                }
+            }
+            _ => AllNullsResult::NotAllNulls,
+        }
+    }
+}
+
+enum AllNullsResult {
+    AllNulls(DataType),
+    NotAllNulls,
+}
+
+impl NewColumnTransformOptimizer for SqlToAllNullsOptimizer {
+    fn optimize(
+        &self,
+        dataset: &Dataset,
+        transform: NewColumnTransform,
+    ) -> Result<NewColumnTransform> {
+        match &transform {
+            NewColumnTransform::SqlExpressions(expressions) => {
+                let arrow_schema = Arc::new(Schema::from(dataset.schema()));
+                let planner = Planner::new(arrow_schema);
+                let mut all_null_schema_fields = vec![];
+                for (name, expr) in expressions {
+                    let expr = planner.parse_expr(expr)?;
+                    if let AllNullsResult::AllNulls(data_type) = self.is_all_null(&expr) {
+                        let field = Field::new(name, data_type, true);
+                        all_null_schema_fields.push(field);
+                    } else {
+                        return Ok(transform);
+                    }
+                }
+
+                let all_null_schema = Schema::new(all_null_schema_fields);
+                Ok(NewColumnTransform::AllNulls(Arc::new(all_null_schema)))
+            }
+            _ => Ok(transform),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use arrow_array::RecordBatchIterator;
+
+    #[tokio::test]
+    async fn test_sql_to_all_null_transform() {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
+        let empty_reader = RecordBatchIterator::new(vec![], schema.clone());
+        let dataset = Arc::new(
+            Dataset::write(empty_reader, "memory://", None)
+                .await
+                .unwrap(),
+        );
+
+        let original = NewColumnTransform::SqlExpressions(vec![
+            ("new_col1".to_string(), "CAST(NULL AS int)".to_string()),
+            ("new_col2".to_string(), "CAST(NULL AS bigint)".to_string()),
+        ]);
+
+        let optimizer = SqlToAllNullsOptimizer::new();
+        let result = optimizer.optimize(&dataset, original).unwrap();
+
+        assert!(matches!(result, NewColumnTransform::AllNulls(_)));
+        if let NewColumnTransform::AllNulls(schema) = result {
+            assert_eq!(schema.fields().len(), 2);
+            assert_eq!(schema.field(0).name(), "new_col1");
+            assert_eq!(schema.field(0).data_type(), &DataType::Int32);
+            assert!(schema.field(0).is_nullable());
+            assert_eq!(schema.field(1).name(), "new_col2");
+            assert_eq!(schema.field(1).data_type(), &DataType::Int64);
+            assert!(schema.field(1).is_nullable());
+        }
+    }
+}