Merge pull request #9327 from youngsofun/fix_csv

fix(csv): fix align_flush with header only.

Author: BohuTANG

src/query/pipeline/sources/src/processors/sources/input_formats/impls/input_format_csv.rs

@@ -299,6 +299,48 @@ impl InputFormatTextBase for InputFormatCSV {
 
     fn align_flush(state: &mut AligningState<Self>) -> Result<Vec<RowBatch>> {
         let mut res = vec![];
+        let num_fields = state.num_fields;
+        let reader = state.csv_reader.as_mut().expect("must success");
+        let field_ends = &mut reader.field_ends[..];
+        let in_tmp = Vec::new();
+        let mut out_tmp = vec![0u8; 1];
+        let mut endlen = reader.n_end;
+
+        if state.rows_to_skip > 0 {
+            let (result, _, _, n_end) =
+                reader
+                    .reader
+                    .read_record(&in_tmp, &mut out_tmp, &mut field_ends[endlen..]);
+            endlen += n_end;
+
+            return match result {
+                ReadRecordResult::InputEmpty => {
+                    reader.n_end = endlen;
+                    Ok(vec![])
+                }
+                ReadRecordResult::OutputFull => Err(output_full_error(&state.path, state.rows)),
+                ReadRecordResult::OutputEndsFull => Err(output_ends_full_error(
+                    num_fields,
+                    field_ends.len(),
+                    &state.path,
+                    state.rows,
+                )),
+                ReadRecordResult::Record => {
+                    Self::check_num_field(num_fields, endlen, field_ends, &state.path, state.rows)?;
+
+                    state.rows_to_skip -= 1;
+                    tracing::debug!(
+                        "csv aligner: skip a header row, remain {}",
+                        state.rows_to_skip
+                    );
+                    Ok(vec![])
+                }
+                ReadRecordResult::End => {
+                    Err(csv_error("unexpect eof in header", &state.path, state.rows))
+                }
+            };
+        }
+
         let num_fields = state.num_fields;
         let reader = state.csv_reader.as_mut().expect("must success");
         let field_ends = &mut reader.field_ends[..];

src/query/pipeline/sources/src/processors/sources/input_formats/source_aligner.rs

@@ -123,8 +123,10 @@ impl<I: InputFormatPipe> Processor for Aligner<I> {
                 let eof = read_batch.is_none();
                 let row_batches = state.align(read_batch)?;
                 for b in row_batches.into_iter() {
-                    process_values.rows += b.rows();
-                    self.row_batches.push_back(b);
+                    if b.size() > 0 {
+                        process_values.rows += b.rows();
+                        self.row_batches.push_back(b);
+                    }
                 }
                 if eof {
                     assert!(self.is_flushing_split);

tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_header_only.result

@@ -0,0 +1 @@
+1

tests/suites/1_stateful/05_formats/05_02_csv/05_02_05_csv_header_only.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../../shell_env.sh
+
+ echo "drop table if exists test_csv_header_only" | $MYSQL_CLIENT_CONNECT
+
+cat << EOF > /tmp/test_csv_header_only.csv
+c1,c2
+EOF
+
+ echo "CREATE TABLE test_csv_header_only
+ (
+     a Int,
+     b Int
+ );" | $MYSQL_CLIENT_CONNECT
+
+ curl -sH "insert_sql:insert into test_csv_header_only file_format = (type = 'CSV' skip_header = 1 record_delimiter = '-')" -F "upload=@/tmp/test_csv_header_only.csv" \
+ -u root: -XPUT "http://localhost:${QUERY_HTTP_HANDLER_PORT}/v1/streaming_load" | grep -c '\"rows\":0'