From 7dab694480aa0d80f806b6cf4c0794109ba7035d Mon Sep 17 00:00:00 2001 From: "yazhou.zhang" Date: Thu, 13 Aug 2026 18:37:02 +0800 Subject: [PATCH] fix: stream ORC and Parquet files intact Route binary columnar formats through a file-granularity upload path so UTF-8 validation and text delimiters cannot corrupt or reject them. Keep worker concurrency across complete files and preserve redirects, compression, and retry behavior. --- README.md | 8 ++ loader/file_loader_test.go | 241 +++++++++++++++++++++++++++++++++++++ loader/format_test.go | 41 +++++++ loader/stream_loader.go | 112 +++++++++++++++-- main.go | 20 ++- reader/reader.go | 7 ++ 6 files changed, 414 insertions(+), 15 deletions(-) create mode 100644 loader/file_loader_test.go create mode 100644 loader/format_test.go diff --git a/README.md b/README.md index 1905eb9..b8f10fe 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,14 @@ e.g.: doris-streamloader --source_file="data.csv" --url="http://localhost:8330" --header="column_separator:|?columns:col1,col2" --db="testdb" --table="testtbl" ``` +ORC and Parquet files are streamed without text batching. Each source file is +sent intact in a separate Stream Load request, and `workers` controls how many +files can be uploaded concurrently: + +```shell +doris-streamloader --source_file="data/*.orc" --url="http://localhost:8330" --header="format:orc" --db="testdb" --table="testtbl" --workers=4 +``` + For additional details and options, refer to our comprehensive docs below. diff --git a/loader/file_loader_test.go b/loader/file_loader_test.go new file mode 100644 index 0000000..f3a5015 --- /dev/null +++ b/loader/file_loader_test.go @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +package loader + +import ( + "bytes" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "sync" + "testing" + + "doris-streamloader/report" + + "github.com/pierrec/lz4/v4" +) + +func TestLoadFilesPreservesBinaryContents(t *testing.T) { + want := []byte("ORC\x00\xff\n\x80binary-footer") + file := createTempSourceFile(t, want) + + var got []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var err error + got, err = io.ReadAll(r.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + fmt.Fprint(w, `{"Status":"Success","NumberTotalRows":3,"NumberLoadedRows":3,"LoadBytes":20}`) + })) + t.Cleanup(server.Close) + + streamLoad, loadResp := newFileTestStreamLoad(server.URL) + retryInfo := map[int]int{} + streamLoad.LoadFiles([]*os.File{file}, &retryInfo) + streamLoad.wg.Wait() + + if !bytes.Equal(got, want) { + t.Fatalf("request body was modified:\n got: %v\nwant: %v", got, want) + } + if loadResp.TotalRows != 3 { + t.Fatalf("TotalRows = %d, want 3", loadResp.TotalRows) + } +} + +func TestLoadFilesSendsEachFileAsASeparateRequest(t *testing.T) { + wants := [][]byte{ + []byte("ORC-first-file"), + []byte("ORC-second-file"), + } + files := []*os.File{ + createTempSourceFile(t, wants[0]), + createTempSourceFile(t, wants[1]), + } + + var ( + mu sync.Mutex + gots [][]byte + ) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + mu.Lock() + gots = append(gots, body) + mu.Unlock() + fmt.Fprint(w, `{"Status":"Success","NumberTotalRows":1,"NumberLoadedRows":1}`) + })) + t.Cleanup(server.Close) + + streamLoad, _ := newFileTestStreamLoad(server.URL) + retryInfo := map[int]int{} + streamLoad.LoadFiles(files, &retryInfo) + streamLoad.wg.Wait() + + if len(gots) != len(wants) { + t.Fatalf("request count = %d, want %d", len(gots), len(wants)) + } + remaining := make(map[string]int, len(wants)) + for _, want := range wants { + remaining[string(want)]++ + } + for _, got := range gots { + remaining[string(got)]-- + } + for content, count := range remaining { + if count != 0 { + t.Fatalf("request body %q count delta = %d, want 0", content, count) + } + } +} + +func TestLoadFilesRetriesOnlySelectedFiles(t *testing.T) { + files := []*os.File{ + createTempSourceFile(t, []byte("ORC-first-file")), + createTempSourceFile(t, []byte("ORC-second-file")), + createTempSourceFile(t, []byte("ORC-third-file")), + } + + var got []byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var err error + got, err = io.ReadAll(r.Body) + if err != nil { + t.Errorf("read request body: %v", err) + } + fmt.Fprint(w, `{"Status":"Success","NumberTotalRows":1,"NumberLoadedRows":1}`) + })) + t.Cleanup(server.Close) + + streamLoad, _ := newFileTestStreamLoad(server.URL) + retryInfo := map[int]int{1: 1} + streamLoad.LoadFiles(files, &retryInfo) + streamLoad.wg.Wait() + + if want := []byte("ORC-second-file"); !bytes.Equal(got, want) { + t.Fatalf("retried request body = %q, want %q", got, want) + } + if gotWorkers := streamLoad.report.TotalWorkers; gotWorkers != 1 { + t.Fatalf("TotalWorkers = %d, want 1", gotWorkers) + } +} + +func TestLoadFilesCompressesBinaryContentsWhenEnabled(t *testing.T) { + want := []byte("ORC\x00\xff\n\x80binary-footer") + file := createTempSourceFile(t, want) + + var ( + got []byte + readErr error + encoding string + ) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + encoding = r.Header.Get("Content-Encoding") + got, readErr = io.ReadAll(lz4.NewReader(r.Body)) + fmt.Fprint(w, `{"Status":"Success","NumberTotalRows":1,"NumberLoadedRows":1}`) + })) + t.Cleanup(server.Close) + + streamLoad, _ := newFileTestStreamLoad(server.URL) + streamLoad.Compress = true + retryInfo := map[int]int{} + streamLoad.LoadFiles([]*os.File{file}, &retryInfo) + streamLoad.wg.Wait() + + if readErr != nil { + t.Fatalf("decompress request body: %v", readErr) + } + if !bytes.Equal(got, want) { + t.Fatalf("decompressed request body = %v, want %v", got, want) + } + if encoding != "lz4" { + t.Fatalf("Content-Encoding = %q, want lz4", encoding) + } +} + +func TestLoadFilesPreservesBinaryContentsAfterRedirect(t *testing.T) { + want := []byte("ORC\x00\xff\n\x80binary-footer") + file := createTempSourceFile(t, want) + + var got []byte + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var err error + got, err = io.ReadAll(r.Body) + if err != nil { + t.Errorf("read redirected request body: %v", err) + } + fmt.Fprint(w, `{"Status":"Success","NumberTotalRows":1,"NumberLoadedRows":1}`) + })) + t.Cleanup(backend.Close) + + frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Location", backend.URL+r.URL.Path) + w.WriteHeader(http.StatusTemporaryRedirect) + })) + t.Cleanup(frontend.Close) + + streamLoad, _ := newFileTestStreamLoad(frontend.URL) + retryInfo := map[int]int{} + streamLoad.LoadFiles([]*os.File{file}, &retryInfo) + streamLoad.wg.Wait() + + if !bytes.Equal(got, want) { + t.Fatalf("redirected request body = %v, want %v", got, want) + } +} + +func createTempSourceFile(t *testing.T, content []byte) *os.File { + t.Helper() + + file, err := os.CreateTemp(t.TempDir(), "source-*") + if err != nil { + t.Fatal(err) + } + if _, err := file.Write(content); err != nil { + t.Fatal(err) + } + if _, err := file.Seek(0, io.SeekStart); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = file.Close() }) + return file +} + +func newFileTestStreamLoad(serverURL string) (*StreamLoad, *Resp) { + loadResp := &Resp{} + reporter := report.NewReporter(1, 1, 1) + pool := &sync.Pool{New: func() interface{} { return make([]byte, 0, 1024) }} + streamLoad := NewStreamLoad( + serverURL, + "db", + "table", + "user", + "password", + map[string]string{"format": "orc"}, + nil, + pool, + StreamLoadOption{CheckUTF8: true, Timeout: 5}, + reporter, + loadResp, + ) + return streamLoad, loadResp +} diff --git a/loader/format_test.go b/loader/format_test.go new file mode 100644 index 0000000..b7b165b --- /dev/null +++ b/loader/format_test.go @@ -0,0 +1,41 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +package loader + +import "testing" + +func TestIsFileStreamFormat(t *testing.T) { + tests := []struct { + format string + want bool + }{ + {format: "orc", want: true}, + {format: "ORC", want: true}, + {format: " parquet ", want: true}, + {format: "csv", want: false}, + {format: "json", want: false}, + } + + for _, tt := range tests { + t.Run(tt.format, func(t *testing.T) { + if got := IsFileStreamFormat(tt.format); got != tt.want { + t.Fatalf("IsFileStreamFormat(%q) = %t, want %t", tt.format, got, tt.want) + } + }) + } +} diff --git a/loader/stream_loader.go b/loader/stream_loader.go index 80840c4..264f314 100644 --- a/loader/stream_loader.go +++ b/loader/stream_loader.go @@ -24,6 +24,7 @@ import ( "io" "io/ioutil" "net/http" + "os" "strconv" "strings" "sync" @@ -39,6 +40,16 @@ type StreamLoadOption struct { Compress bool CheckUTF8 bool Timeout int + Workers int +} + +func IsFileStreamFormat(format string) bool { + switch strings.ToLower(strings.TrimSpace(format)) { + case "orc", "parquet": + return true + default: + return false + } } type StreamLoad struct { @@ -242,15 +253,15 @@ func (s *StreamLoad) readData(isEOS *atomic.Bool, rawWriter *io.PipeWriter, read } } -func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskIndex int) (*http.Response, error) { +func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskIndex int) (*http.Response, *StreamLoadResp, error) { realUrl := url for { req, err := s.createRequest(realUrl, reader, workerIndex, taskIndex) if err != nil { if req == nil { - return nil, err + return nil, nil, err } else { - return req.Response, err + return req.Response, nil, err } } @@ -266,7 +277,7 @@ func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskInd // send request resp, err := client.Do(req) if err != nil { - return resp, err + return resp, nil, err } defer resp.Body.Close() @@ -274,14 +285,14 @@ func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskInd body, err := ioutil.ReadAll(resp.Body) if err != nil { log.Errorf("Read stream load response failed, response body: %s, error message: %v, it is hard to judge whether load success or not, we suggest:\n1.Do select count(*) to check whether data is partially loaded.\n2.If the data is partially loaded and duplication is unacceptable, consider dropping the table (with caution that all data in the table will be lost) and retry.", body, err) - return resp, nil + return resp, nil, nil } if resp.StatusCode == 307 { realUrl = resp.Header.Get("Location") if realUrl == "" { log.Errorf("Send error, redirectUrl is empty, error message: %v", err) - return resp, nil + return resp, nil, nil } log.Infof("redirect to %s", realUrl) continue @@ -291,7 +302,7 @@ func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskInd if resp.StatusCode != 200 { // print response headers log.Debugf("response headers: %v", resp.Header) - return resp, fmt.Errorf("response code is not 200, response code: %d, response body: %s", resp.StatusCode, body) + return resp, nil, fmt.Errorf("response code is not 200, response code: %d, response body: %s", resp.StatusCode, body) } // parse response @@ -299,10 +310,10 @@ func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskInd err = json.Unmarshal(body, &respMsg) if err != nil { log.Errorf("Parse stream load response failed, response body: %s, error message: %v, it is hard to judge whether load success or not, we suggest:\n1.Do select count(*) to check whether data is partially loaded\n2.If the data is partially loaded and duplication is unacceptable, consider dropping the table (with caution that all data in the table will be lost) and retry.\n", body, err) - return resp, nil + return resp, nil, nil } if respMsg.Status != "Success" && respMsg.Status != "Publish Timeout" { - return resp, fmt.Errorf("stream load failed, response %s", body) + return resp, &respMsg, fmt.Errorf("stream load failed, response %s", body) } log.Infof("resp: %s", body) @@ -313,7 +324,7 @@ func (s *StreamLoad) send(url string, reader io.Reader, workerIndex int, taskInd atomic.AddUint64(&s.loadResp.UnselectedRows, uint64(respMsg.NumberUnselectedRows)) atomic.AddUint64(&s.loadResp.LoadBytes, uint64(respMsg.LoadBytes)) - return resp, nil + return resp, &respMsg, nil } } @@ -361,7 +372,7 @@ func (s *StreamLoad) executeGetAndSend(maxRowsPerTask int, maxBytesPerTask int, workerIndex: workerIndex, taskIndex: taskIndex, }) - if resp, err := s.send(url, NopCloser(pr), workerIndex, taskIndex); err != nil { + if resp, _, err := s.send(url, NopCloser(pr), workerIndex, taskIndex); err != nil { s.handleSendError(workerIndex, taskIndex) log.Errorf("Send error, resp: %v error message: %v", resp, err) return @@ -377,6 +388,85 @@ func (s *StreamLoad) executeGetAndSend(maxRowsPerTask int, maxBytesPerTask int, atomic.AddUint64(&s.report.TotalWorkers, 1) } +// LoadFiles streams every file as a separate request without interpreting or +// modifying its contents. This is required for file formats such as ORC and +// Parquet, whose files cannot be split into text batches. +func (s *StreamLoad) LoadFiles(files []*os.File, retryInfo *map[int]int) { + workers := s.Workers + if workers <= 0 { + workers = 1 + } + + selected := make([]int, 0, len(files)) + for fileIndex, file := range files { + if len(*retryInfo) > 0 { + if _, retry := (*retryInfo)[fileIndex]; !retry { + continue + } + } + selected = append(selected, fileIndex) + s.loadResp.LoadFiles = append(s.loadResp.LoadFiles, file.Name()) + } + atomic.StoreUint64(&s.report.TotalWorkers, uint64(len(selected))) + + semaphore := make(chan struct{}, workers) + for _, fileIndex := range selected { + fileIndex := fileIndex + file := files[fileIndex] + s.wg.Add(1) + go func() { + defer s.wg.Done() + semaphore <- struct{}{} + defer func() { <-semaphore }() + + if _, err := file.Seek(0, io.SeekStart); err != nil { + s.handleSendError(fileIndex, 1) + log.Errorf("Seek file %s failed, error message: %v", file.Name(), err) + return + } + + reader := io.Reader(&reportingReader{reader: file, reporter: s.report}) + if s.Compress { + pipeReader, pipeWriter := io.Pipe() + go compressLZ4(pipeWriter, reader) + reader = pipeReader + } + resp, loadResult, err := s.send(s.createUrl(), reader, fileIndex, 1) + if err != nil { + s.handleSendError(fileIndex, 1) + log.Errorf("Send file %s failed, resp: %v error message: %v", file.Name(), resp, err) + return + } + if loadResult != nil { + atomic.AddUint64(&s.loadResp.TotalRows, uint64(loadResult.NumberTotalRows)) + } + atomic.AddUint64(&s.report.FinishedWorkers, 1) + }() + } +} + +func compressLZ4(writer *io.PipeWriter, reader io.Reader) { + compressor := lz4.NewWriter(writer) + _, copyErr := io.Copy(compressor, reader) + closeErr := compressor.Close() + if copyErr != nil { + _ = writer.CloseWithError(copyErr) + return + } + _ = writer.CloseWithError(closeErr) +} + +type reportingReader struct { + reader io.Reader + reporter *report.Reporter +} + +func (r *reportingReader) Read(p []byte) (int, error) { + n, err := r.reader.Read(p) + r.reporter.IncrSendBytes(int64(n)) + return n, err +} + func (s *StreamLoad) ExecuteGetAndSend(maxRowsPerTask int, maxBytesPerTask int, workerIndex int) { s.wg.Add(1) go s.executeGetAndSend(maxRowsPerTask, maxBytesPerTask, workerIndex) diff --git a/main.go b/main.go index 38d8799..4c88687 100644 --- a/main.go +++ b/main.go @@ -83,6 +83,7 @@ var ( showVersion bool queueSize int lineDelimiter byte = '\n' + fileStreamFormat bool bufferPool = sync.Pool{ New: func() interface{} { return make([]byte, 0, bufferSize) @@ -272,6 +273,9 @@ func paramCheck() { if strings.ToLower(kv[0]) == "format" && strings.ToLower(kv[1]) != "csv" { enableConcurrency = false } + if strings.ToLower(kv[0]) == "format" { + fileStreamFormat = loader.IsFileStreamFormat(kv[1]) + } if strings.ToLower(kv[0]) == "line_delimiter" { @@ -370,6 +374,7 @@ func createStreamLoad(report *report.Reporter, queues []chan []byte) *loader.Str Compress: compress, CheckUTF8: checkUTF8, Timeout: timeout, + Workers: workers, } return loader.NewStreamLoad(url, dbName, tableName, userName, password, headers, queues, &bufferPool, streamLoadOption, report, loadResp) } @@ -397,14 +402,21 @@ func main() { startTime = time.Now() } - maxRowsPerTask = math.MinInt32 - streamLoad.Load(workers, maxRowsPerTask, maxBytesPerTask, &retryInfo) reporter.Report() defer reporter.CloseWait() - reader.Read(reporter, workers, maxBytesPerTask, &retryInfo, loadResp, retryCount, lineDelimiter) - reader.Close() + if fileStreamFormat { + streamLoad.LoadFiles(reader.Files(), &retryInfo) + } else { + maxRowsPerTask = math.MinInt32 + streamLoad.Load(workers, maxRowsPerTask, maxBytesPerTask, &retryInfo) + reader.Read(reporter, workers, maxBytesPerTask, &retryInfo, loadResp, retryCount, lineDelimiter) + reader.Close() + } streamLoad.Wait(loadInfo, retryCount, &retryInfo, startTime) + if fileStreamFormat { + reader.Close() + } if !loadInfo.NeedRetry || retryCount >= loadInfo.RetryTimes-1 { break } diff --git a/reader/reader.go b/reader/reader.go index fa6563f..a749722 100644 --- a/reader/reader.go +++ b/reader/reader.go @@ -217,4 +217,11 @@ func (f *FileReader) Close() { for _, queue := range *f.queues { close(queue) } + for _, file := range f.files { + _ = file.Close() + } +} + +func (f *FileReader) Files() []*os.File { + return f.files }