<commit_before>package main

import (
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sync/atomic"

	"github.com/influxdb/influxdb/tsdb/engine/tsm1"
)

var (
	NanFiltered     uint64
	InfFiltered     uint64
	PointsWritten   uint64
	PointsRead      uint64
	TsmFilesCreated uint64
)

type KeyIterator interface {
	Next() bool
	Read() (string, []tsm1.Value, error)
}

// Converter encapsulates the logic for converting b*1 shards to tsm1 shards.
type Converter struct {
	path           string
	maxTSMFileSize uint32
	sequence       int
}

// NewConverter returns a new instance of the Converter.
func NewConverter(path string, sz uint32) *Converter {
	return &Converter{
		path:           path,
		maxTSMFileSize: sz,
	}
}

// Process writes the data provided by iter to a tsm1 shard.
func (c *Converter) Process(iter KeyIterator) error {
	// Ensure the tsm1 directory exists.
	if err := os.MkdirAll(c.path, 0777); err != nil {
		return err
	}

	// Iterate until no more data remains.
	var w tsm1.TSMWriter
	for iter.Next() {
		k, v, err := iter.Read()
		if err != nil {
			return err
		}
		scrubbed := scrubValues(v)

		if w == nil {
			w, err = c.nextTSMWriter()
			if err != nil {
				return err
			}
		}
		if err := w.Write(k, scrubbed); err != nil {
			return err
		}
		atomic.AddUint64(&PointsRead, uint64(len(v)))
		atomic.AddUint64(&PointsWritten, uint64(len(scrubbed)))

		// If we have a max file size configured and we're over it, start a new TSM file.
		if w.Size() > c.maxTSMFileSize {
			if err := w.WriteIndex(); err != nil && err != tsm1.ErrNoValues {
				return err
			}
			if err := w.Close(); err != nil {
				return err
			}
			w = nil
		}
	}

	if w != nil {
		if err := w.WriteIndex(); err != nil && err != tsm1.ErrNoValues {
			return err
		}
		if err := w.Close(); err != nil {
			return err
		}
	}

	return nil
}

// nextTSMWriter returns the next TSMWriter for the Converter.
func (c *Converter) nextTSMWriter() (tsm1.TSMWriter, error) {
	c.sequence++
	fileName := filepath.Join(c.path, fmt.Sprintf("%09d-%09d.%s", 1, c.sequence, tsm1.TSMFileExtension))

	fd, err := os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}

	// Create the writer for the new TSM file.
	w, err := tsm1.NewTSMWriter(fd)
	if err != nil {
		return nil, err
	}<commit_msg>Track bytes written by conversion tool<commit_after>package main

import (
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sync/atomic"

	"github.com/influxdb/influxdb/tsdb/engine/tsm1"
)

var (
	NanFiltered     uint64
	InfFiltered     uint64
	PointsWritten   uint64
	PointsRead      uint64
	TsmFilesCreated uint64
	TsmBytesWritten int64
)

type KeyIterator interface {
	Next() bool
	Read() (string, []tsm1.Value, error)
}

// Converter encapsulates the logic for converting b*1 shards to tsm1 shards.
type Converter struct {
	path           string
	maxTSMFileSize uint32
	sequence       int
}

// NewConverter returns a new instance of the Converter.
func NewConverter(path string, sz uint32) *Converter {
	return &Converter{
		path:           path,
		maxTSMFileSize: sz,
	}
}

// Process writes the data provided by iter to a tsm1 shard.
func (c *Converter) Process(iter KeyIterator) error {
	// Ensure the tsm1 directory exists.
	if err := os.MkdirAll(c.path, 0777); err != nil {
		return err
	}

	// Iterate until no more data remains.
	var w tsm1.TSMWriter
	for iter.Next() {
		k, v, err := iter.Read()
		if err != nil {
			return err
		}
		scrubbed := scrubValues(v)

		if w == nil {
			w, err = c.nextTSMWriter()
			if err != nil {
				return err
			}
		}
		if err := w.Write(k, scrubbed); err != nil {
			return err
		}
		atomic.AddUint64(&PointsRead, uint64(len(v)))
		atomic.AddUint64(&PointsWritten, uint64(len(scrubbed)))

		// If we have a max file size configured and we're over it, start a new TSM file.
		if w.Size() > c.maxTSMFileSize {
			if err := w.WriteIndex(); err != nil && err != tsm1.ErrNoValues {
				return err
			}
			atomic.AddInt64(&TsmBytesWritten, int64(w.Size()))

			if err := w.Close(); err != nil {
				return err
			}
			w = nil
		}
	}

	if w != nil {
		if err := w.WriteIndex(); err != nil && err != tsm1.ErrNoValues {
			return err
		}
		atomic.AddInt64(&TsmBytesWritten, int64(w.Size()))

		if err := w.Close(); err != nil {
			return err
		}
	}

	return nil
}

// nextTSMWriter returns the next TSMWriter for the Converter.
func (c *Converter) nextTSMWriter() (tsm1.TSMWriter, error) {
	c.sequence++
	fileName := filepath.Join(c.path, fmt.Sprintf("%09d-%09d.%s", 1, c.sequence, tsm1.TSMFileExtension))

	fd, err := os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
	if err != nil {
		return nil, err
	}

	// Create the writer for the new TSM file.
	w, err := tsm1.NewTSMWriter(fd)
	if err != nil {
		return nil, err
	}
{'fraction_non_alphanumeric': 0.12555298082999788, 'fraction_numerical': 0.021908573836107014, 'mean_word_length': 3.0169204737732658, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '895119', 'n_tokens_mistral': 1927, 'n_tokens_neox': 1738, 'n_words': 515}