fix: infer CSV/TSV columns during stage registration

KKould · KKould · commit b3c5792e9df7 · 2026-03-18T19:29:29.000+08:00
diff --git a/src/bendpy/src/context.rs b/src/bendpy/src/context.rs
@@ -15,6 +15,8 @@
 use std::sync::Arc;
 
 use databend_common_exception::Result;
+use databend_common_expression::BlockEntry;
+use databend_common_expression::Column;
 use databend_common_meta_app::principal::BUILTIN_ROLE_ACCOUNT_ADMIN;
 use databend_common_version::BUILD_INFO;
 use databend_query::sessions::BuildInfoRef;
@@ -32,33 +34,66 @@ fn resolve_file_path(path: &str) -> String {
     if path.contains("://") {
         return path.to_owned();
     }
+
     if path.starts_with('/') {
         return format!("fs://{}", path);
     }
+
     format!(
         "fs://{}/{}",
         std::env::current_dir().unwrap().to_str().unwrap(),
         path
     )
 }
 
-/// Extract the real filesystem path from a `fs://` URI.
-fn fs_path_from_uri(uri: &str) -> Option<&str> {
-    uri.strip_prefix("fs://")
+fn extract_string_column(
+    entry: &BlockEntry,
+) -> Option<&databend_common_expression::types::StringColumn> {
+    match entry {
+        BlockEntry::Column(Column::String(col)) => Some(col),
+        BlockEntry::Column(Column::Nullable(n)) => match &n.column {
+            Column::String(col) => Some(col),
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+fn build_infer_schema_sql(
+    file_path: &str,
+    file_format: &str,
+    pattern: Option<&str>,
+    connection: Option<&str>,
+) -> String {
+    let connection_clause = connection
+        .map(|c| format!(", connection_name => '{}'", c))
+        .unwrap_or_default();
+    let pattern_clause = pattern
+        .map(|p| format!(", pattern => '{}'", p))
+        .unwrap_or_default();
+
+    format!(
+        "SELECT column_name FROM infer_schema(location => '{}', file_format => '{}'{}{})",
+        file_path,
+        file_format.to_uppercase(),
+        pattern_clause,
+        connection_clause
+    )
 }
 
-/// Read the header line of a CSV file and return column names.
-fn read_csv_column_names(path: &str) -> std::io::Result<Vec<String>> {
-    use std::io::BufRead;
-    let file = std::fs::File::open(path)?;
-    let mut reader = std::io::BufReader::new(file);
-    let mut header = String::new();
-    reader.read_line(&mut header)?;
-    Ok(header
-        .trim()
-        .split(',')
-        .map(|s| s.trim().trim_matches('"').to_string())
-        .collect())
+fn build_position_select(col_names: &[String]) -> PyResult<String> {
+    if col_names.is_empty() {
+        return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
+            "Could not infer schema: no columns found",
+        ));
+    }
+
+    Ok(col_names
+        .iter()
+        .enumerate()
+        .map(|(i, name)| format!("${} AS `{}`", i + 1, name))
+        .collect::<Vec<_>>()
+        .join(", "))
 }
 
 #[pyclass(name = "SessionContext", module = "databend", subclass)]
@@ -214,47 +249,21 @@ impl PySessionContext {
         let pattern_clause = pattern
             .map(|p| format!(", pattern => '{}'", p))
             .unwrap_or_default();
-
         let select_clause = match file_format {
-            "csv" => self.build_column_select(&file_path)?,
+            "csv" | "tsv" => {
+                self.build_column_select(&file_path, file_format, pattern, connection, py)?
+            }
             _ => "*".to_string(),
         };
-
         let sql = format!(
             "create view {} as select {} from '{}' (file_format => '{}'{}{})",
             name, select_clause, file_path, file_format, pattern_clause, connection_clause
         );
+
         let _ = self.sql(&sql, py)?.collect(py)?;
         Ok(())
     }
 
-    /// Read CSV header from local file and build `$1 AS col1, $2 AS col2, ...`.
-    fn build_column_select(&self, file_path: &str) -> PyResult<String> {
-        let fs_path = fs_path_from_uri(file_path).ok_or_else(|| {
-            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
-                "CSV column inference only supports local files (fs://), got: {}",
-                file_path
-            ))
-        })?;
-        let col_names = read_csv_column_names(fs_path).map_err(|e| {
-            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
-                "Failed to read CSV header: {}",
-                e
-            ))
-        })?;
-        if col_names.is_empty() {
-            return Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
-                "Could not infer schema: no columns found",
-            ));
-        }
-        Ok(col_names
-            .iter()
-            .enumerate()
-            .map(|(i, name)| format!("${} AS `{}`", i + 1, name))
-            .collect::<Vec<_>>()
-            .join(", "))
-    }
-
     #[pyo3(signature = (name, access_key_id, secret_access_key, endpoint_url = None, region = None))]
     fn create_s3_connection(
         &mut self,
@@ -398,8 +407,59 @@ impl PySessionContext {
     }
 }
 
+impl PySessionContext {
+    fn build_column_select(
+        &mut self,
+        file_path: &str,
+        file_format: &str,
+        pattern: Option<&str>,
+        connection: Option<&str>,
+        py: Python,
+    ) -> PyResult<String> {
+        let sql = build_infer_schema_sql(file_path, file_format, pattern, connection);
+        let blocks = self.sql(&sql, py)?.collect(py)?;
+
+        let col_names = blocks
+            .blocks
+            .iter()
+            .filter(|b| b.num_rows() > 0)
+            .filter_map(|b| extract_string_column(b.get_by_offset(0)))
+            .flat_map(|col| col.iter().map(|s| s.to_string()))
+            .collect::<Vec<_>>();
+
+        build_position_select(&col_names)
+    }
+}
+
 async fn plan_sql(ctx: &Arc<QueryContext>, sql: &str) -> Result<PyDataFrame> {
     let mut planner = Planner::new(ctx.clone());
     let (plan, _) = planner.plan_sql(sql).await?;
     Ok(PyDataFrame::new(ctx.clone(), plan, default_box_size()))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_resolve_file_path_absolute() {
+        assert_eq!(resolve_file_path("/tmp/data.csv"), "fs:///tmp/data.csv");
+    }
+
+    #[test]
+    fn test_build_position_select() {
+        let col_names = vec!["column_1".to_string(), "column_2".to_string()];
+        let select = build_position_select(&col_names).unwrap();
+        assert_eq!(select, "$1 AS `column_1`, $2 AS `column_2`");
+    }
+
+    #[test]
+    fn test_build_infer_schema_sql_with_pattern_and_connection() {
+        let sql = build_infer_schema_sql("s3://bucket/logs/", "tsv", Some("*.tsv"), Some("my_s3"));
+
+        assert_eq!(
+            sql,
+            "SELECT column_name FROM infer_schema(location => 's3://bucket/logs/', file_format => 'TSV', pattern => '*.tsv', connection_name => 'my_s3')"
+        );
+    }
+}
diff --git a/src/bendpy/tests/test_connections.py b/src/bendpy/tests/test_connections.py
@@ -43,13 +43,6 @@ def register_parquet(self, name, path, pattern=None, connection=None):
             sql = f"create view {name} as select * from '{path}' (file_format => 'parquet'{pattern_clause})"
         self.sql(sql)
 
-    def register_csv(self, name, path, pattern=None, connection=None):
-        file_path = path if connection else (f"fs://{path}" if path.startswith("/") else path)
-        pattern_clause = f", pattern => '{pattern}'" if pattern else ""
-        conn_clause = f", connection => '{connection}'" if connection else ""
-        sql = f"create view {name} as select * from '{file_path}' (file_format => 'csv'{pattern_clause}{conn_clause})"
-        self.sql(sql)
-
     def create_azblob_connection(self, name, endpoint_url, account_name, account_key):
         sql = f"CREATE OR REPLACE CONNECTION {name} STORAGE_TYPE = 'AZBLOB' endpoint_url = '{endpoint_url}' account_name = '{account_name}' account_key = '{account_key}'"
         self.sql(sql)
@@ -245,15 +238,6 @@ def test_register_parquet_with_connection_and_pattern(self):
             expected_sql = "create view sales as select * from 's3://bucket/data/' (file_format => 'parquet', pattern => '*.parquet', connection => 'my_s3')"
             mock_sql.assert_called_once_with(expected_sql)
 
-    def test_register_csv_with_connection(self):
-        with unittest.mock.patch.object(self.ctx, "sql") as mock_sql:
-            mock_sql.return_value.collect.return_value = None
-
-            self.ctx.register_csv("users", "s3://bucket/users.csv", connection="my_s3")
-
-            expected_sql = "create view users as select * from 's3://bucket/users.csv' (file_format => 'csv', connection => 'my_s3')"
-            mock_sql.assert_called_once_with(expected_sql)
-
     def test_register_parquet_legacy_mode(self):
         with unittest.mock.patch.object(self.ctx, "sql") as mock_sql:
             mock_sql.return_value.collect.return_value = None
@@ -263,27 +247,6 @@ def test_register_parquet_legacy_mode(self):
             expected_sql = "create view local as select * from '/data/file.parquet' (file_format => 'parquet')"
             mock_sql.assert_called_once_with(expected_sql)
 
-    def test_register_csv_with_pattern_no_connection(self):
-        with unittest.mock.patch.object(self.ctx, "sql") as mock_sql:
-            mock_sql.return_value.collect.return_value = None
-
-            self.ctx.register_csv("logs", "/data/logs/", pattern="*.csv")
-
-            expected_sql = "create view logs as select * from 'fs:///data/logs/' (file_format => 'csv', pattern => '*.csv')"
-            mock_sql.assert_called_once_with(expected_sql)
-
-    def test_register_csv_with_pattern_and_connection(self):
-        with unittest.mock.patch.object(self.ctx, "sql") as mock_sql:
-            mock_sql.return_value.collect.return_value = None
-
-            self.ctx.register_csv(
-                "logs", "s3://bucket/logs/", pattern="*.csv", connection="my_s3"
-            )
-
-            expected_sql = "create view logs as select * from 's3://bucket/logs/' (file_format => 'csv', pattern => '*.csv', connection => 'my_s3')"
-            mock_sql.assert_called_once_with(expected_sql)
-
-
 class TestStages:
     def setup_method(self):
         self.ctx = MockSessionContext()
diff --git a/src/bendpy/tests/test_register_stage_files.py b/src/bendpy/tests/test_register_stage_files.py
@@ -0,0 +1,49 @@
+## Copyright 2021 Datafuse Labs
+##
+## Licensed under the Apache License, Version 2.0 (the "License");
+## you may not use this file except in compliance with the License.
+## You may obtain a copy of the License at
+##
+##     http://www.apache.org/licenses/LICENSE-2.0
+##
+## Unless required by applicable law or agreed to in writing, software
+## distributed under the License is distributed on an "AS IS" BASIS,
+## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+## See the License for the specific language governing permissions and
+## limitations under the License.
+
+from pathlib import Path
+
+from databend import SessionContext
+
+
+ROOT = Path(__file__).resolve().parents[3]
+CSV_DIR = ROOT / "tests" / "data" / "csv"
+TSV_DIR = ROOT / "tests" / "data" / "tsv"
+CSV_PATH = CSV_DIR / "select.csv"
+TSV_PATH = TSV_DIR / "select.tsv"
+
+
+class TestRegisterDelimitedFiles:
+    def setup_method(self):
+        self.ctx = SessionContext()
+
+    def assert_view_rows(self, view_name):
+        df = self.ctx.sql(f"select * from {view_name} order by column_1").to_pandas()
+        assert df.values.tolist() == [[1, None, None], [2, "b", "B"], [3, "c", None]]
+
+    def test_register_csv_select_star(self):
+        self.ctx.register_csv("csv_stage_view", str(CSV_PATH))
+        self.assert_view_rows("csv_stage_view")
+
+    def test_register_csv_select_star_with_pattern(self):
+        self.ctx.register_csv("csv_stage_pattern_view", str(CSV_DIR), pattern="select.csv")
+        self.assert_view_rows("csv_stage_pattern_view")
+
+    def test_register_tsv_select_star(self):
+        self.ctx.register_tsv("tsv_stage_view", str(TSV_PATH))
+        self.assert_view_rows("tsv_stage_view")
+
+    def test_register_tsv_select_star_with_pattern(self):
+        self.ctx.register_tsv("tsv_stage_pattern_view", str(TSV_DIR), pattern="select.tsv")
+        self.assert_view_rows("tsv_stage_pattern_view")
