feat: infer_schema expands csv and ndjson support

KKould · KKould · commit ef8e64f4a848 · 2025-08-18T16:15:09.000+08:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -229,8 +229,10 @@ arrow = { version = "55" }
 arrow-array = { version = "55" }
 arrow-buffer = { version = "55" }
 arrow-cast = { version = "55", features = ["prettyprint"] }
+arrow-csv = { version = "55" }
 arrow-data = { version = "55" }
 arrow-flight = { version = "55", features = ["flight-sql-experimental", "tls"] }
+arrow-json = { version = "55" }
 arrow-ipc = { version = "55" }
 arrow-ord = { version = "55" }
 arrow-schema = { version = "55", features = ["serde"] }
diff --git a/src/query/catalog/src/table_args.rs b/src/query/catalog/src/table_args.rs
@@ -119,6 +119,12 @@ pub fn bool_value(value: &Scalar) -> Result<bool> {
     }
 }
 
+pub fn i64_value(value: &Scalar) -> Result<i64> {
+    value.get_i64().ok_or_else(|| {
+        ErrorCode::BadArguments(format!("invalid value {value} expect to be i64 literal."))
+    })
+}
+
 pub fn string_literal(val: &str) -> Scalar {
     Scalar::String(val.to_string())
 }
diff --git a/src/query/service/Cargo.toml b/src/query/service/Cargo.toml
@@ -23,7 +23,9 @@ io-uring = [
 anyhow = { workspace = true }
 arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
+arrow-csv = { workspace = true }
 arrow-flight = { workspace = true }
+arrow-json = { workspace = true }
 arrow-ipc = { workspace = true, features = ["lz4", "zstd"] }
 arrow-schema = { workspace = true }
 arrow-select = { workspace = true }
diff --git a/src/query/service/src/table_functions/infer_schema/infer_schema_table.rs b/src/query/service/src/table_functions/infer_schema/infer_schema_table.rs
@@ -32,7 +32,7 @@ use databend_common_meta_app::schema::TableInfo;
 use databend_common_meta_app::schema::TableMeta;
 use databend_common_pipeline_core::Pipeline;
 
-use super::parquet::ParquetInferSchemaSource;
+use super::source::InferSchemaSource;
 use crate::sessions::TableContext;
 use crate::table_functions::infer_schema::table_args::InferSchemaArgsParsed;
 use crate::table_functions::TableFunction;
@@ -114,9 +114,7 @@ impl Table for InferSchemaTable {
         _put_cache: bool,
     ) -> Result<()> {
         pipeline.add_source(
-            |output| {
-                ParquetInferSchemaSource::create(ctx.clone(), output, self.args_parsed.clone())
-            },
+            |output| InferSchemaSource::create(ctx.clone(), output, self.args_parsed.clone()),
             1,
         )?;
         Ok(())
diff --git a/src/query/service/src/table_functions/infer_schema/mod.rs b/src/query/service/src/table_functions/infer_schema/mod.rs
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 mod infer_schema_table;
-mod parquet;
+mod source;
 mod table_args;
 
 pub use infer_schema_table::InferSchemaTable;
diff --git a/src/query/service/src/table_functions/infer_schema/source.rs b/src/query/service/src/table_functions/infer_schema/source.rs
@@ -13,8 +13,12 @@
 // limitations under the License.
 
 use std::collections::BTreeMap;
+use std::io::Cursor;
 use std::sync::Arc;
 
+use arrow_csv::reader::Format;
+use arrow_json::reader::infer_json_schema;
+use arrow_schema::Schema as ArrowSchema;
 use databend_common_ast::ast::FileLocation;
 use databend_common_ast::ast::UriLocation;
 use databend_common_catalog::table_context::TableContext;
@@ -26,7 +30,8 @@ use databend_common_expression::types::UInt64Type;
 use databend_common_expression::DataBlock;
 use databend_common_expression::FromData;
 use databend_common_expression::TableSchema;
-use databend_common_meta_app::principal::StageFileFormatType;
+use databend_common_meta_app::principal::CsvFileFormatParams;
+use databend_common_meta_app::principal::FileFormatParams;
 use databend_common_meta_app::principal::StageType;
 use databend_common_pipeline_core::processors::OutputPort;
 use databend_common_pipeline_core::processors::ProcessorPtr;
@@ -37,24 +42,25 @@ use databend_common_storage::init_stage_operator;
 use databend_common_storage::read_parquet_schema_async_rs;
 use databend_common_storage::StageFilesInfo;
 use databend_common_users::Object;
+use opendal::Operator;
 use opendal::Scheme;
 
 use crate::table_functions::infer_schema::infer_schema_table::INFER_SCHEMA;
 use crate::table_functions::infer_schema::table_args::InferSchemaArgsParsed;
 
-pub(crate) struct ParquetInferSchemaSource {
+pub(crate) struct InferSchemaSource {
     is_finished: bool,
     ctx: Arc<dyn TableContext>,
     args_parsed: InferSchemaArgsParsed,
 }
 
-impl ParquetInferSchemaSource {
+impl InferSchemaSource {
     pub fn create(
         ctx: Arc<dyn TableContext>,
         output: Arc<OutputPort>,
         args_parsed: InferSchemaArgsParsed,
     ) -> Result<ProcessorPtr> {
-        AsyncSourcer::create(ctx.clone(), output, ParquetInferSchemaSource {
+        AsyncSourcer::create(ctx.clone(), output, InferSchemaSource {
             is_finished: false,
             ctx,
             args_parsed,
@@ -63,7 +69,7 @@ impl ParquetInferSchemaSource {
 }
 
 #[async_trait::async_trait]
-impl AsyncSource for ParquetInferSchemaSource {
+impl AsyncSource for InferSchemaSource {
     const NAME: &'static str = INFER_SCHEMA;
 
     #[async_backtrace::framed]
@@ -127,9 +133,9 @@ impl AsyncSource for ParquetInferSchemaSource {
             Some(f) => self.ctx.get_file_format(f).await?,
             None => stage_info.file_format_params.clone(),
         };
-        let schema = match (first_file.as_ref(), file_format_params.get_type()) {
+        let schema = match (first_file.as_ref(), file_format_params) {
             (None, _) => return Ok(None),
-            (Some(first_file), StageFileFormatType::Parquet) => {
+            (Some(first_file), FileFormatParams::Parquet(_)) => {
                 let arrow_schema = read_parquet_schema_async_rs(
                     &operator,
                     &first_file.path,
@@ -138,6 +144,27 @@ impl AsyncSource for ParquetInferSchemaSource {
                 .await?;
                 TableSchema::try_from(&arrow_schema)?
             }
+            (Some(first_file), FileFormatParams::Csv(params)) => {
+                let arrow_schema = read_csv_metadata_async(
+                    &first_file.path,
+                    &operator,
+                    Some(first_file.size),
+                    self.args_parsed.max_records,
+                    &params,
+                )
+                .await?;
+                TableSchema::try_from(&arrow_schema)?
+            }
+            (Some(first_file), FileFormatParams::NdJson(_)) => {
+                let arrow_schema = read_json_metadata_async(
+                    &first_file.path,
+                    &operator,
+                    Some(first_file.size),
+                    self.args_parsed.max_records,
+                )
+                .await?;
+                TableSchema::try_from(&arrow_schema)?
+            }
             _ => {
                 return Err(ErrorCode::BadArguments(
                     "infer_schema is currently limited to format Parquet",
@@ -168,3 +195,52 @@ impl AsyncSource for ParquetInferSchemaSource {
         Ok(Some(block))
     }
 }
+
+pub async fn read_csv_metadata_async(
+    path: &str,
+    operator: &Operator,
+    file_size: Option<u64>,
+    max_records: Option<usize>,
+    params: &CsvFileFormatParams,
+) -> Result<ArrowSchema> {
+    let file_size = match file_size {
+        None => operator.stat(path).await?.content_length(),
+        Some(n) => n,
+    };
+    let escape = if params.escape.is_empty() {
+        None
+    } else {
+        Some(params.escape.as_bytes()[0])
+    };
+
+    // TODO: It would be better if it could be read in the form of Read trait
+    let buf = operator.read_with(path).range(..file_size).await?.to_vec();
+    let mut format = Format::default()
+        .with_delimiter(params.field_delimiter.as_bytes()[0])
+        .with_quote(params.quote.as_bytes()[0])
+        .with_header(params.headers != 0);
+
+    if let Some(escape) = escape {
+        format = format.with_escape(escape);
+    }
+    let (schema, _) = format.infer_schema(Cursor::new(&buf), max_records)?;
+
+    Ok(schema)
+}
+
+pub async fn read_json_metadata_async(
+    path: &str,
+    operator: &Operator,
+    file_size: Option<u64>,
+    max_records: Option<usize>,
+) -> Result<ArrowSchema> {
+    let file_size = match file_size {
+        None => operator.stat(path).await?.content_length(),
+        Some(n) => n,
+    };
+    // TODO: It would be better if it could be read in the form of Read trait
+    let buf = operator.read_with(path).range(..file_size).await?.to_vec();
+    let (schema, _) = infer_json_schema(Cursor::new(&buf), max_records)?;
+
+    Ok(schema)
+}
diff --git a/src/query/service/src/table_functions/infer_schema/table_args.rs b/src/query/service/src/table_functions/infer_schema/table_args.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use databend_common_catalog::table_args::i64_value;
 use databend_common_catalog::table_args::TableArgs;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -24,6 +25,7 @@ pub(crate) struct InferSchemaArgsParsed {
     pub(crate) connection_name: Option<String>,
     pub(crate) file_format: Option<String>,
     pub(crate) files_info: StageFilesInfo,
+    pub(crate) max_records: Option<usize>,
 }
 
 impl InferSchemaArgsParsed {
@@ -38,6 +40,7 @@ impl InferSchemaArgsParsed {
             files: None,
             pattern: None,
         };
+        let mut max_records = None;
 
         for (k, v) in &args {
             match k.to_lowercase().as_str() {
@@ -53,6 +56,9 @@ impl InferSchemaArgsParsed {
                 "file_format" => {
                     file_format = Some(string_value(v)?);
                 }
+                "max_records_pre_file" => {
+                    max_records = Some(i64_value(v)? as usize);
+                }
                 _ => {
                     return Err(ErrorCode::BadArguments(format!(
                         "unknown param {} for infer_schema",
@@ -70,6 +76,7 @@ impl InferSchemaArgsParsed {
             connection_name,
             file_format,
             files_info,
+            max_records,
         })
     }
 }
diff --git a/tests/data/csv/max_records.csv b/tests/data/csv/max_records.csv
@@ -0,0 +1,11 @@
+id,value
+1,100
+2,200
+3,300
+4,400
+5,500
+6,foo
+7,bar
+8,baz
+9,qux
+10,quux
diff --git a/tests/data/csv/mixed.csv b/tests/data/csv/mixed.csv
@@ -0,0 +1,4 @@
+id,name,score,active
+1,Alice,88.5,true
+2,Bob,92.0,false
+3,Charlie,,true
diff --git a/tests/data/csv/numbers_with_headers.csv b/tests/data/csv/numbers_with_headers.csv
@@ -0,0 +1,19 @@
+id,value
+0,1
+1,2
+2,3
+3,4
+4,5
+5,6
+6,7
+7,8
+8,9
+9,10
+10,11
+11,12
+12,13
+13,14
+14,15
+15,16
+16,17
+17,18
diff --git a/tests/data/csv/ragged.csv b/tests/data/csv/ragged.csv
@@ -0,0 +1,5 @@
+id,value,comment
+1,10,ok
+2,20
+3,30,missing one field
+4
diff --git a/tests/data/ndjson/max_records.ndjson b/tests/data/ndjson/max_records.ndjson
@@ -0,0 +1,10 @@
+{"id": 1, "value": 100}
+{"id": 2, "value": 200}
+{"id": 3, "value": 300}
+{"id": 4, "value": 400}
+{"id": 5, "value": 500}
+{"id": 6, "value": "foo"}
+{"id": 7, "value": "bar"}
+{"id": 8, "value": "baz"}
+{"id": 9, "value": "qux"}
+{"id": 10, "value": "quux"}
diff --git a/tests/data/ndjson/mixed.ndjson b/tests/data/ndjson/mixed.ndjson
@@ -0,0 +1,3 @@
+{"id": 1, "name": "Alice", "score": 88.5, "active": true}
+{"id": 2, "name": "Bob", "score": 92.0, "active": false}
+{"id": 3, "name": "Charlie", "score": null, "active": true}
diff --git a/tests/data/ndjson/numbers.ndjson b/tests/data/ndjson/numbers.ndjson
@@ -0,0 +1,3 @@
+{"id": 1, "value": 100}
+{"id": 2, "value": 200}
+{"id": 3, "value": 300}
diff --git a/tests/data/ndjson/ragged.ndjson b/tests/data/ndjson/ragged.ndjson
@@ -0,0 +1,4 @@
+{"id": 1, "value": 10, "comment": "ok"}
+{"id": 2, "value": 20}
+{"id": 3, "value": 30, "comment": "missing one field"}
+{"id": 4}
diff --git a/tests/sqllogictests/suites/stage/formats/parquet/infer_schema.test b/tests/sqllogictests/suites/stage/formats/parquet/infer_schema.test
@@ -66,3 +66,60 @@ select * from INFER_SCHEMA(location => 's3://testbucket/data/parquet/tuple.parqu
 ----
 id INT 0 0
 t TUPLE(A INT32, B STRING) 0 1
+
+# CSV
+statement ok
+create or replace file format head_csv_format type = 'CSV' field_delimiter = ',' skip_header = 1;
+
+query TTBI
+select * from infer_schema(location => '@data/csv/numbers_with_headers.csv', file_format => 'CSV');
+----
+column_1 VARCHAR 1 0
+column_2 VARCHAR 1 1
+
+query TTBI
+select * from infer_schema(location => '@data/csv/numbers_with_headers.csv', file_format => 'head_csv_format');
+----
+id BIGINT 1 0
+value BIGINT 1 1
+
+statement error
+select * from infer_schema(location => '@data/csv/ragged.csv', file_format => 'head_csv_format');
+
+query TTBI
+select * from infer_schema(location => '@data/csv/max_records.csv', file_format => 'head_csv_format');
+----
+id BIGINT 1 0
+value VARCHAR 1 1
+
+query TTBI
+select * from infer_schema(location => '@data/csv/max_records.csv', file_format => 'head_csv_format', max_records_pre_file => 5);
+----
+id BIGINT 1 0
+value BIGINT 1 1
+
+# NDJSON
+query TTBI
+select * from infer_schema(location => '@data/ndjson/numbers.ndjson', file_format => 'NDJSON');
+----
+id BIGINT 1 0
+value BIGINT 1 1
+
+query TTBI
+select * from infer_schema(location => '@data/ndjson/ragged.ndjson', file_format => 'NDJSON');
+----
+id BIGINT 1 0
+value BIGINT 1 1
+comment VARCHAR 1 2
+
+query TTBI
+select * from infer_schema(location => '@data/ndjson/max_records.ndjson', file_format => 'NDJSON');
+----
+id BIGINT 1 0
+value VARCHAR 1 1
+
+query TTBI
+select * from infer_schema(location => '@data/ndjson/max_records.ndjson', file_format => 'NDJSON', max_records_pre_file => 5);
+----
+id BIGINT 1 0
+value BIGINT 1 1

Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,12 @@ pub fn bool_value(value: &Scalar) -> Result<bool> {`
`119`	`119`	`}`
`120`	`120`	`}`
`121`	`121`
	`122`	`+pub fn i64_value(value: &Scalar) -> Result<i64> {`
	`123`	`+ value.get_i64().ok_or_else(\|\| {`
	`124`	`+ ErrorCode::BadArguments(format!("invalid value {value} expect to be i64 literal."))`
	`125`	`+ })`
	`126`	`+}`
	`127`	`+`
`122`	`128`	`pub fn string_literal(val: &str) -> Scalar {`
`123`	`129`	`Scalar::String(val.to_string())`
`124`	`130`	`}`
-Original file line number
+Diff line change
@@ @@ -0,0 +1,11 @@ @@
 +id,value
 +1,100
 +2,200
 +3,300
 +4,400
 +5,500
 +6,foo
 +7,bar
 +8,baz
 +9,qux
 +10,quux