fix: spaced columns parsing

Signed-off-by: Ion Koutsouris <[email protected]>

Author: ion-elgreco

crates/core/src/delta_datafusion/expr.rs

@@ -333,7 +333,10 @@ impl Display for BinaryExprFormat<'_> {
 impl Display for SqlFormat<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self.expr {
-            Expr::Column(c) => write!(f, "{c}"),
+            Expr::Column(c) => {
+                let c = c.quoted_flat_name();
+                write!(f, "{c}")
+            }
             Expr::Literal(v) => write!(f, "{}", ScalarValueFormat { scalar: v }),
             Expr::Case(case) => {
                 write!(f, "CASE ")?;

crates/core/src/delta_datafusion/mod.rs

@@ -1591,7 +1591,7 @@ impl DeltaDataChecker {
             return Ok(());
         }
         let table = MemTable::try_new(record_batch.schema(), vec![vec![record_batch.clone()]])?;
-
+        table.schema();
         // Use a random table name to avoid clashes when running multiple parallel tasks, e.g. when using a partitioned table
         let table_name: String = uuid::Uuid::new_v4().to_string();
         self.ctx.register_table(&table_name, Arc::new(table))?;
@@ -2408,6 +2408,58 @@ mod tests {
         assert!(result.is_err());
     }
 
+    /// Ensure that constraints when there are spaces in the field name still work
+    ///
+    /// See <https://github.com/delta-io/delta-rs/pull/3374>
+    #[tokio::test]
+    async fn test_constraints_with_spacey_fields() -> DeltaResult<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", ArrowDataType::Utf8, false),
+            Field::new("b bop", ArrowDataType::Int32, false),
+        ]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(arrow::array::StringArray::from(vec![
+                    "a", "b bop", "c", "d",
+                ])),
+                Arc::new(arrow::array::Int32Array::from(vec![1, 10, 10, 100])),
+            ],
+        )?;
+
+        // Valid invariants return Ok(())
+        let constraints = vec![
+            Constraint::new("custom a", "a is not null"),
+            Constraint::new("custom_b", "`b bop` < 1000"),
+        ];
+        assert!(DeltaDataChecker::new_with_constraints(constraints)
+            .check_batch(&batch)
+            .await
+            .is_ok());
+
+        // Violated invariants returns an error with list of violations
+        let constraints = vec![
+            Constraint::new("custom_a", "a is null"),
+            Constraint::new("custom_B", "\"b bop\" < 100"),
+        ];
+        let result = DeltaDataChecker::new_with_constraints(constraints)
+            .check_batch(&batch)
+            .await;
+        assert!(result.is_err());
+        assert!(matches!(result, Err(DeltaTableError::InvalidData { .. })));
+        if let Err(DeltaTableError::InvalidData { violations }) = result {
+            assert_eq!(violations.len(), 2);
+        }
+
+        // Irrelevant constraints return a different error
+        let constraints = vec![Constraint::new("custom_c", "c > 2000")];
+        let result = DeltaDataChecker::new_with_constraints(constraints)
+            .check_batch(&batch)
+            .await;
+        assert!(result.is_err());
+        Ok(())
+    }
+
     #[test]
     fn roundtrip_test_delta_exec_plan() {
         let ctx = SessionContext::new();

crates/core/src/operations/write/mod.rs

@@ -593,7 +593,6 @@ impl std::future::IntoFuture for WriteBuilder {
                         Expression::String(s) => {
                             let df_schema = DFSchema::try_from(schema.as_ref().to_owned())?;
                             parse_predicate_expression(&df_schema, s, &state)?
-                            // this.snapshot.unwrap().parse_predicate_expression(s, &state)?
                         }
                     };
                     (Some(fmt_expr_to_sql(&pred)?), Some(pred))

python/tests/test_constraint.py

@@ -0,0 +1,127 @@
+import pytest
+from arro3.core import Array, DataType, Field, Schema, Table
+
+from deltalake import DeltaTable, write_deltalake
+from deltalake.exceptions import DeltaError, DeltaProtocolError
+
+
+@pytest.fixture()
+def sample_table() -> Table:
+    nrows = 5
+    return Table(
+        {
+            "id": Array(
+                ["1", "2", "3", "4", "5"],
+                Field("id", type=DataType.string(), nullable=True),
+            ),
+            "high price": Array(
+                list(range(nrows)),
+                Field("high price", type=DataType.int64(), nullable=True),
+            ),
+        },
+    )
+
+
+def test_not_corrupting_expression(tmp_path):
+    data = Table.from_pydict(
+        {
+            "b": Array([1], DataType.int64()),
+            "color_column": Array(["red"], DataType.string()),
+        },
+    )
+
+    data2 = Table.from_pydict(
+        {
+            "b": Array([1], DataType.int64()),
+            "color_column": Array(["blue"], DataType.string()),
+        },
+    )
+
+    write_deltalake(
+        tmp_path,
+        data,
+        mode="overwrite",
+        partition_by=["color_column"],
+        predicate="color_column = 'red'",
+    )
+    write_deltalake(
+        tmp_path,
+        data2,
+        mode="overwrite",
+        partition_by=["color_column"],
+        predicate="color_column = 'blue'",
+    )
+
+
+def test_not_corrupting_expression_columns_spaced(tmp_path):
+    data = Table.from_pydict(
+        {
+            "b": Array([1], DataType.int64()),
+            "color column": Array(["red"], DataType.string()),
+        },
+    )
+
+    data2 = Table.from_pydict(
+        {
+            "b": Array([1], DataType.int64()),
+            "color column": Array(["blue"], DataType.string()),
+        },
+    )
+
+    write_deltalake(
+        tmp_path,
+        data,
+        mode="overwrite",
+        # partition_by=["color column"],
+        predicate="`color column` = 'red'",
+    )
+    write_deltalake(
+        tmp_path,
+        data2,
+        mode="overwrite",
+        # partition_by=["color column"],
+        predicate="`color column` = 'blue'",
+    )
+
+
+# fmt: off
+
+@pytest.mark.parametrize("sql_string", [
+    "`high price` >= 0",
+    '"high price" >= 0',
+    "\"high price\" >= 0"
+])
+def test_add_constraint(tmp_path, sample_table: Table, sql_string: str):
+    write_deltalake(tmp_path, sample_table)
+
+    dt = DeltaTable(tmp_path)
+
+    dt.alter.add_constraint({"check_price": sql_string})
+
+    last_action = dt.history(1)[0]
+    assert last_action["operation"] == "ADD CONSTRAINT"
+    assert dt.version() == 1
+    assert dt.metadata().configuration == {
+        "delta.constraints.check_price": '"high price" >= 0'
+    }
+    assert dt.protocol().min_writer_version == 3
+
+    with pytest.raises(DeltaError):
+        # Invalid constraint
+        dt.alter.add_constraint({"check_price": '"high price" < 0'})
+
+    with pytest.raises(DeltaProtocolError):
+        data = Table(
+            {
+                "id": Array(["1"], DataType.string()),
+                "high price": Array([-1], DataType.int64()),
+            },
+            schema=Schema(
+                fields=[
+                    Field("id", type=DataType.string(), nullable=True),
+                    Field("high price", type=DataType.int64(), nullable=True),
+                ]
+            ),
+        )
+
+        write_deltalake(tmp_path, data, mode="append")