diff --git a/states/inspect_parquet.go b/states/inspect_parquet.go index a9647ce0..dde0a98f 100644 --- a/states/inspect_parquet.go +++ b/states/inspect_parquet.go @@ -2,6 +2,7 @@ package states import ( "context" + "encoding/binary" "encoding/json" "fmt" "net/url" @@ -31,6 +32,7 @@ type InspectParquetParam struct { FilePath string `name:"file" default:"" desc:"local parquet file path to inspect"` SegmentID int64 `name:"segment" default:"0" desc:"segment ID to inspect binlogs from remote storage"` FieldID int64 `name:"field" default:"0" desc:"only inspect binlogs of this field ID (0 means all fields)"` + StorageVersion string `name:"storageVersion" default:"meta" desc:"segment storage version strategy: meta, auto, or a non-negative version override"` MetadataOnly bool `name:"metadataOnly" default:"true" desc:"print metadata only; set to false to also sample rows"` SampleRows int64 `name:"sampleRows" default:"10" desc:"number of rows to sample when metadataOnly=false"` ShowRowGroups bool `name:"showRowGroups" default:"false" desc:"print per-row-group statistics"` @@ -42,6 +44,38 @@ type InspectParquetParam struct { ExternalFile string `name:"externalFile" default:"" desc:"single parquet object path to inspect in external mode"` } +const inspectAutoStorageVersion int64 = -1 << 63 + +type inspectStorageVersionMode int + +const ( + inspectStorageVersionFromMeta inspectStorageVersionMode = iota + inspectStorageVersionAuto + inspectStorageVersionOverride +) + +type inspectStorageVersionOption struct { + mode inspectStorageVersionMode + version int64 +} + +func parseInspectStorageVersionOption(raw string) (inspectStorageVersionOption, error) { + value := strings.ToLower(strings.TrimSpace(raw)) + switch value { + case "", "meta": + return inspectStorageVersionOption{mode: inspectStorageVersionFromMeta}, nil + case "auto": + return inspectStorageVersionOption{mode: inspectStorageVersionAuto}, nil + } + + version, err := strconv.ParseInt(value, 10, 64) + if err != nil || version < 0 { + return inspectStorageVersionOption{}, errors.Newf( + "invalid --storageVersion %q: expected meta, auto, or a non-negative integer", raw) + } + return inspectStorageVersionOption{mode: inspectStorageVersionOverride, version: version}, nil +} + type externalSourceSpec struct { Format string CloudProvider string @@ -74,6 +108,14 @@ func (s *InstanceState) InspectParquetCommand(ctx context.Context, p *InspectPar } func validateInspectParquetParam(p *InspectParquetParam) error { + storageVersion, err := parseInspectStorageVersionOption(p.StorageVersion) + if err != nil { + return err + } + if storageVersion.mode != inspectStorageVersionFromMeta && (p.External || p.FilePath != "") { + return errors.New("--storageVersion can only be used with --segment") + } + if p.External { if p.FilePath != "" || p.SegmentID != 0 { return errors.New("--external cannot be used with --file or --segment") @@ -129,6 +171,28 @@ func (s *InstanceState) inspectSegmentParquet(ctx context.Context, p *InspectPar fmt.Printf("Segment %d: collection=%d partition=%d storageVersion=%d\n", segment.ID, segment.CollectionID, segment.PartitionID, segment.StorageVersion) + storageVersionOption, err := parseInspectStorageVersionOption(p.StorageVersion) + if err != nil { + return err + } + effectiveStorageVersion := segment.GetStorageVersion() + switch storageVersionOption.mode { + case inspectStorageVersionFromMeta: + fmt.Printf("Storage Version Strategy: meta (effective=%d)\n", effectiveStorageVersion) + case inspectStorageVersionOverride: + effectiveStorageVersion = storageVersionOption.version + fmt.Printf("Storage Version Strategy: override (meta=%d effective=%d)\n", + segment.GetStorageVersion(), effectiveStorageVersion) + case inspectStorageVersionAuto: + if segment.GetManifestPath() != "" { + effectiveStorageVersion = 3 + fmt.Printf("Storage Version Strategy: auto (manifest layout detected)\n") + } else { + effectiveStorageVersion = inspectAutoStorageVersion + fmt.Printf("Storage Version Strategy: auto (detect each binlog object)\n") + } + } + params := []oss.MinioConnectParam{oss.WithSkipCheckBucket(p.SkipBucketCheck)} if p.MinioAddress != "" { params = append(params, oss.WithMinioAddr(p.MinioAddress)) @@ -139,9 +203,9 @@ func (s *InstanceState) inspectSegmentParquet(ctx context.Context, p *InspectPar } rootPath := resolvedStore.RootPath - if segment.GetStorageVersion() >= 3 { + if effectiveStorageVersion >= 3 { if segment.GetManifestPath() == "" { - return errors.Newf("segment %d storage version is %d but got empty manifest", segment.GetID(), segment.GetStorageVersion()) + return errors.Newf("segment %d effective storage version is %d but got empty manifest", segment.GetID(), effectiveStorageVersion) } return inspectV3SegmentParquet(ctx, resolvedStore.Store, rootPath, segment, p) } @@ -153,7 +217,7 @@ func (s *InstanceState) inspectSegmentParquet(ctx context.Context, p *InspectPar for _, binlog := range fieldBinlog.Binlogs { logPath := oss.ResolveObjectKey(rootPath, binlog.LogPath) fmt.Printf("\n===== Field %d | %s =====\n", fieldBinlog.FieldID, logPath) - if err := inspectRemoteBinlog(ctx, resolvedStore.Store, logPath, segment.StorageVersion, p.MetadataOnly, p.SampleRows, p.ShowRowGroups); err != nil { + if err := inspectRemoteBinlog(ctx, resolvedStore.Store, logPath, effectiveStorageVersion, p.MetadataOnly, p.SampleRows, p.ShowRowGroups); err != nil { fmt.Printf("failed to inspect %s: %s\n", logPath, err.Error()) } } @@ -377,29 +441,59 @@ func inspectRemoteBinlog(ctx context.Context, store oss.ObjectStore, logPath str if err != nil { return err } + if closer, ok := obj.(interface{ Close() error }); ok { + defer closer.Close() + } - pqReader, err := openBinlogParquet(obj, storageVersion) + pqReader, detectedStorageVersion, err := openBinlogParquet(obj, storageVersion) if err != nil { return err } defer pqReader.Close() + if storageVersion == inspectAutoStorageVersion { + fmt.Printf("Detected Storage Version: %d (%s)\n", detectedStorageVersion, path.Base(logPath)) + } return printParquetFile(ctx, pqReader, path.Base(logPath), metadataOnly, sampleRows, showRowGroups) } -func openBinlogParquet(r storagecommon.ReadSeeker, storageVersion int64) (*file.Reader, error) { +func openBinlogParquet(r storagecommon.ReadSeeker, storageVersion int64) (*file.Reader, int64, error) { switch storageVersion { + case inspectAutoStorageVersion: + detectedStorageVersion, err := detectBinlogStorageVersion(r) + if err != nil { + return nil, 0, err + } + return openBinlogParquet(r, detectedStorageVersion) case 2: - return file.NewParquetReader(r) + pqReader, err := file.NewParquetReader(r) + return pqReader, storageVersion, err case 0, 1: br, err := binlogv1.NewBinlogReader(r) if err != nil { - return nil, err + return nil, 0, err } - return br.NextParquetReader() + pqReader, err := br.NextParquetReader() + return pqReader, storageVersion, err default: - return nil, errors.Newf("unsupported storage version: %d", storageVersion) + return nil, 0, errors.Newf("unsupported storage version: %d", storageVersion) + } +} + +func detectBinlogStorageVersion(r storagecommon.ReadSeeker) (int64, error) { + var header [4]byte + n, err := r.ReadAt(header[:], 0) + if err != nil { + return 0, errors.Wrapf(err, "read binlog header: got %d of %d bytes", n, len(header)) + } + + if string(header[:]) == "PAR1" { + return 2, nil + } + if binary.LittleEndian.Uint32(header[:]) == uint32(binlogv1.MagicNumberV1) { + return 0, nil } + return 0, errors.Newf("unrecognized binlog header: %x", header) } func printParquetFile(ctx context.Context, pqReader *file.Reader, name string, metadataOnly bool, sampleRows int64, showRowGroups bool) error { diff --git a/states/inspect_parquet_test.go b/states/inspect_parquet_test.go new file mode 100644 index 00000000..d56f8779 --- /dev/null +++ b/states/inspect_parquet_test.go @@ -0,0 +1,270 @@ +package states + +import ( + "bytes" + "encoding/binary" + "testing" + + binlogv1 "github.com/milvus-io/birdwatcher/storage/binlog/v1" +) + +func TestValidateInspectParquetParam(t *testing.T) { + tests := []struct { + name string + param InspectParquetParam + wantErr bool + }{ + { + name: "local file mode", + param: InspectParquetParam{FilePath: "/tmp/a.parquet"}, + wantErr: false, + }, + { + name: "segment mode", + param: InspectParquetParam{SegmentID: 100}, + wantErr: false, + }, + { + name: "segment auto storage version", + param: InspectParquetParam{SegmentID: 100, StorageVersion: "auto"}, + wantErr: false, + }, + { + name: "segment storage version override", + param: InspectParquetParam{SegmentID: 100, StorageVersion: "2"}, + wantErr: false, + }, + { + name: "external manifest mode", + param: InspectParquetParam{External: true, CollectionID: 1, ManifestSegmentID: 10}, + wantErr: false, + }, + { + name: "external file mode", + param: InspectParquetParam{External: true, CollectionID: 1, ExternalFile: "a/b.parquet"}, + wantErr: false, + }, + { + name: "missing selectors", + param: InspectParquetParam{}, + wantErr: true, + }, + { + name: "conflicting local selectors", + param: InspectParquetParam{FilePath: "/tmp/a.parquet", SegmentID: 10}, + wantErr: true, + }, + { + name: "local file rejects auto storage version", + param: InspectParquetParam{FilePath: "/tmp/a.parquet", StorageVersion: "auto"}, + wantErr: true, + }, + { + name: "invalid storage version", + param: InspectParquetParam{SegmentID: 100, StorageVersion: "invalid"}, + wantErr: true, + }, + { + name: "negative storage version", + param: InspectParquetParam{SegmentID: 100, StorageVersion: "-1"}, + wantErr: true, + }, + { + name: "external missing collection", + param: InspectParquetParam{External: true, ManifestSegmentID: 10}, + wantErr: true, + }, + { + name: "external missing target", + param: InspectParquetParam{External: true, CollectionID: 1}, + wantErr: true, + }, + { + name: "external conflicting targets", + param: InspectParquetParam{External: true, CollectionID: 1, ManifestSegmentID: 10, ExternalFile: "a.parquet"}, + wantErr: true, + }, + { + name: "external conflicts with segment", + param: InspectParquetParam{External: true, CollectionID: 1, ManifestSegmentID: 10, SegmentID: 20}, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateInspectParquetParam(&tt.param) + if (err != nil) != tt.wantErr { + t.Fatalf("validateInspectParquetParam() error = %v, wantErr %v", err, tt.wantErr) + } + }) + } +} + +func TestParseInspectStorageVersionOption(t *testing.T) { + tests := []struct { + name string + value string + wantMode inspectStorageVersionMode + wantVersion int64 + wantErr bool + }{ + {name: "empty uses meta", value: "", wantMode: inspectStorageVersionFromMeta}, + {name: "meta", value: "meta", wantMode: inspectStorageVersionFromMeta}, + {name: "auto case insensitive", value: " AUTO ", wantMode: inspectStorageVersionAuto}, + {name: "v1 compatibility value", value: "1", wantMode: inspectStorageVersionOverride, wantVersion: 1}, + {name: "v2", value: "2", wantMode: inspectStorageVersionOverride, wantVersion: 2}, + {name: "future manifest version", value: "4", wantMode: inspectStorageVersionOverride, wantVersion: 4}, + {name: "negative", value: "-1", wantErr: true}, + {name: "invalid", value: "v2", wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseInspectStorageVersionOption(tt.value) + if (err != nil) != tt.wantErr { + t.Fatalf("parseInspectStorageVersionOption() error = %v, wantErr %v", err, tt.wantErr) + } + if tt.wantErr { + return + } + if got.mode != tt.wantMode || got.version != tt.wantVersion { + t.Fatalf("parseInspectStorageVersionOption() = %#v, want mode=%v version=%d", got, tt.wantMode, tt.wantVersion) + } + }) + } +} + +func TestDetectBinlogStorageVersion(t *testing.T) { + v1Header := make([]byte, 4) + binary.LittleEndian.PutUint32(v1Header, uint32(binlogv1.MagicNumberV1)) + + tests := []struct { + name string + data []byte + want int64 + wantErr bool + }{ + {name: "v1 wrapper", data: append(v1Header, 0x01), want: 0}, + {name: "raw parquet", data: []byte("PAR1payload"), want: 2}, + {name: "unknown magic", data: []byte("NOPE"), wantErr: true}, + {name: "short header", data: []byte("PAR"), wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + reader := bytes.NewReader(tt.data) + got, err := detectBinlogStorageVersion(reader) + if (err != nil) != tt.wantErr { + t.Fatalf("detectBinlogStorageVersion() error = %v, wantErr %v", err, tt.wantErr) + } + if !tt.wantErr && got != tt.want { + t.Fatalf("detectBinlogStorageVersion() = %d, want %d", got, tt.want) + } + position, seekErr := reader.Seek(0, 1) + if seekErr != nil { + t.Fatalf("reader.Seek() error = %v", seekErr) + } + if position != 0 { + t.Fatalf("detectBinlogStorageVersion() changed reader position to %d", position) + } + }) + } +} + +func TestParseExternalSpec(t *testing.T) { + raw := `{ + "format": "parquet", + "extfs": { + "cloud_provider": "aliyun", + "external_id": "ext-123", + "region": "cn-hangzhou", + "role_arn": "acs:ram::1:role/demo", + "use_ssl": true + } + }` + + spec, err := parseExternalSpec(raw) + if err != nil { + t.Fatalf("parseExternalSpec() error = %v", err) + } + if spec.Format != "parquet" { + t.Fatalf("unexpected format: %s", spec.Format) + } + if spec.CloudProvider != "aliyun" { + t.Fatalf("unexpected cloud provider: %s", spec.CloudProvider) + } + if spec.Region != "cn-hangzhou" { + t.Fatalf("unexpected region: %s", spec.Region) + } + if spec.RoleARN != "acs:ram::1:role/demo" { + t.Fatalf("unexpected role arn: %s", spec.RoleARN) + } + if spec.ExternalID != "ext-123" { + t.Fatalf("unexpected external id: %s", spec.ExternalID) + } + if spec.UseSSL == nil || !*spec.UseSSL { + t.Fatalf("expected use_ssl=true, got %#v", spec.UseSSL) + } +} + +func TestParseExternalSource(t *testing.T) { + location, err := parseExternalSource("oss://oss-cn-hangzhou.aliyuncs.com/test-oss-0815/testlake/parquet") + if err != nil { + t.Fatalf("parseExternalSource() error = %v", err) + } + if location.Scheme != "oss" { + t.Fatalf("unexpected scheme: %s", location.Scheme) + } + if location.Host != "oss-cn-hangzhou.aliyuncs.com" { + t.Fatalf("unexpected host: %s", location.Host) + } + if location.Bucket != "test-oss-0815" { + t.Fatalf("unexpected bucket: %s", location.Bucket) + } + if location.RootPath != "testlake/parquet" { + t.Fatalf("unexpected root path: %s", location.RootPath) + } +} + +func TestResolveExternalObjectKey(t *testing.T) { + location := externalSourceLocation{ + Host: "oss-cn-hangzhou.aliyuncs.com", + Bucket: "test-oss-0815", + RootPath: "testlake/parquet", + } + + tests := []struct { + name string + file string + want string + }{ + { + name: "relative path", + file: "part-0001.parquet", + want: "testlake/parquet/part-0001.parquet", + }, + { + name: "already rooted", + file: "testlake/parquet/part-0002.parquet", + want: "testlake/parquet/part-0002.parquet", + }, + { + name: "full uri", + file: "oss://oss-cn-hangzhou.aliyuncs.com/test-oss-0815/testlake/parquet/part-0003.parquet", + want: "testlake/parquet/part-0003.parquet", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := resolveExternalObjectKey(location, tt.file) + if err != nil { + t.Fatalf("resolveExternalObjectKey() error = %v", err) + } + if got != tt.want { + t.Fatalf("resolveExternalObjectKey() = %s, want %s", got, tt.want) + } + }) + } +}