<commit_before>package com.backtype.hadoop.pail;

import com.backtype.hadoop.*;
import com.backtype.hadoop.formats.RecordInputStream;
import com.backtype.hadoop.formats.RecordOutputStream;
import com.backtype.support.Utils;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.*;

public class Pail<T> extends AbstractPail implements Iterable<T>{
    public static Logger LOG = LoggerFactory.getLogger(Pail.class);

    public static final String META = "pail.meta";

    public class TypedRecordOutputStream implements RecordOutputStream {
        private HashMap<String, RecordOutputStream> _workers = new HashMap<String, RecordOutputStream>();
        private String _userfilename;
        private boolean _overwrite;

        public TypedRecordOutputStream(String userfilename, boolean overwrite) {
            _userfilename = userfilename;
            _overwrite = overwrite;
        }

        public <T> void writeObject(T obj) throws IOException {
            PailStructure<T> structure = ((PailStructure<T>) _structure);
            List<String> rootAttrs = structure.getTarget(obj);
            List<String> attrs = makeRelative(rootAttrs);
            String targetDir = Utils.join(attrs, "/");
            if(!_workers.containsKey(targetDir)) {
                Path p;
                if(targetDir.length()==0) p = new Path(_userfilename);
                else p = new Path(targetDir, _userfilename);
                List<String> totalAttrs = componentsFromRoot(p.toString());
                if(!_structure.isValidTarget(totalAttrs.toArray(new String[totalAttrs.size()]))) {
                    throw new IllegalArgumentException("Cannot write object " + obj.toString() + " to " + p.toString() +
                            ". Conflicts with the structure of the datastore.");
                }
                _workers.put(targetDir, Pail.super.openWrite(p.toString(), _overwrite));
            }
            RecordOutputStream os = _workers.get(targetDir);
            os.writeRaw(structure.serialize(obj));
        }

        public void writeObjects(T... objs) throws IOException {
            for(T obj: objs) {
                writeObject(obj);
            }
        }

        public void close() throws IOException {
            for(RecordOutputStream os: _workers.values()) {
                os.close();
            }
        }

        protected List<String> makeRelative(List<String> attrs) {
            return Utils.stripRoot(getAttrs(), attrs);
        }

        public void writeRaw(byte[] record) throws IOException {
            writeRaw(record, 0, record.length);
        }

        public void writeRaw(byte[] record, int start, int length) throws IOException {
            if(!_workers.containsKey(_userfilename)) {
                checkValidStructure(_userfilename);
                _workers.put(_userfilename, Pail.super.openWrite(_userfilename, _overwrite));
            }
            _workers.get(_userfilename).writeRaw(record, start, length);
        }
    }


    public class TypedRecordInputStream implements RecordInputStream {
        private RecordInputStream is;

        public TypedRecordInputStream(String userFileName) throws IOException {
            is = Pail.super.openRead(userFileName);
        }

        public T readObject() throws IOException {
            byte[] record = readRawRecord();
            if(record==null) return null;
            else return _structure.deserialize(record);
        }

        public void close() throws IOException {
            is.close();
        }

        public byte[] readRawRecord() throws IOException {
            return is.readRawRecord();
        }
    }

    public static Pail create(String path, PailSpec spec) throws IOException {
        return create(Utils.getFS(path), path, spec);
    }

    public static Pail create(FileSystem fs, String path, PailSpec spec) throws IOException {
        return create(fs, path, spec, true);
    }

    public static Pail create(String path) throws IOException {
        return create(Utils.getFS(path), path);
    }

    public static Pail create(FileSystem fs, String path) throws IOException {
        return create(fs, path, (PailSpec) null);
    }

    public static Pail create(String path, PailStructure structure) throws IOException {
        return create(Utils.getFS(path), path, structure);
    }

    public static Pail create(FileSystem fs, String path, PailStructure structure) throws IOException {
        return create(fs, path, new PailSpec(structure));
    }

    public static Pail create(String path, PailStructure structure, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, structure, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, PailStructure structure, boolean failOnExists) throws IOException {
        return create(fs, path, new PailSpec(structure), failOnExists);
    }

    public static Pail create(String path, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, boolean failOnExists) throws IOException {
        return create(fs, path, (PailSpec) null, failOnExists);
    }

    public static Pail create(String path, PailSpec spec, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, spec, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, PailSpec spec, boolean failOnExists) throws IOException {
        Path pathp = new Path(path);
        PailFormatFactory.create(spec);
        PailSpec existing = getSpec(fs, pathp);
        if(failOnExists) {
            if(existing!=null) {
                throw new IllegalArgumentException("Pail already exists at path " + path + " with spec " + existing.toString());
            }
            if(fs.exists(pathp))
                throw new IllegalArgumentException("Path " + path + " already exists");
        }
        if(spec!=null && existing!=null) {
            if(spec.getName()!=null) {
                if(!spec.equals(existing))
                    throw new IllegalArgumentException("Specs do not match " + spec.toString() + ", " + existing.toString());
            } else if(spec.getStructure()!=null) {
                if(existing.getStructure()==null || !spec.getStructure().getClass().equals(existing.getStructure().getClass())) {
                    throw new IllegalArgumentException("Specs do not match " + spec.toString() + ", " + existing.toString());
                }
            }
        }
        fs.mkdirs(pathp);
        if(existing==null) {
            if(spec==null) spec = PailFormatFactory.getDefaultCopy();
            if(spec.getName()==null) spec = PailFormatFactory.getDefaultCopy().setStructure(spec.getStructure());
            spec.writeToFileSystem(fs, new Path(pathp, META));
        }


        return new Pail(fs, path);
    }

    private static PailSpec getSpec(FileSystem fs, Path path) throws IOException {
        return (PailSpec) getSpecAndRoot(fs, path)[1];
    }

    private static String getRoot(FileSystem fs, Path path) throws IOException {
        return (String) getSpecAndRoot(fs, path)[0];
    }

    private static Object[] getSpecAndRoot(FileSystem fs, Path path) throws IOException {
        Path curr = path;
        Object[] ret = null;
        while(true) {
            Path meta = new Path(curr, META);
            if(fs.exists(meta)) {
                if(ret!=null) throw new RuntimeException("At least two meta files up directory tree");
                PailSpec spec = PailSpec.readFromFileSystem(fs, meta);
                ret = new Object[] {curr.toString(), spec};
            }
            if(curr.depth()==0) break;
            curr = curr.getParent();
        }
        if(ret==null) ret = new Object[] {null, null};
        return ret;
    }

<commit_msg>Changes getSpecAndRoot(..) as per Oscar<commit_after>package com.backtype.hadoop.pail;

import com.backtype.hadoop.*;
import com.backtype.hadoop.formats.RecordInputStream;
import com.backtype.hadoop.formats.RecordOutputStream;
import com.backtype.support.Utils;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.FileNotFoundException;
import java.io.IOException;
import java.util.*;

public class Pail<T> extends AbstractPail implements Iterable<T>{
    public static Logger LOG = LoggerFactory.getLogger(Pail.class);

    public static final String META = "pail.meta";

    public class TypedRecordOutputStream implements RecordOutputStream {
        private HashMap<String, RecordOutputStream> _workers = new HashMap<String, RecordOutputStream>();
        private String _userfilename;
        private boolean _overwrite;

        public TypedRecordOutputStream(String userfilename, boolean overwrite) {
            _userfilename = userfilename;
            _overwrite = overwrite;
        }

        public <T> void writeObject(T obj) throws IOException {
            PailStructure<T> structure = ((PailStructure<T>) _structure);
            List<String> rootAttrs = structure.getTarget(obj);
            List<String> attrs = makeRelative(rootAttrs);
            String targetDir = Utils.join(attrs, "/");
            if(!_workers.containsKey(targetDir)) {
                Path p;
                if(targetDir.length()==0) p = new Path(_userfilename);
                else p = new Path(targetDir, _userfilename);
                List<String> totalAttrs = componentsFromRoot(p.toString());
                if(!_structure.isValidTarget(totalAttrs.toArray(new String[totalAttrs.size()]))) {
                    throw new IllegalArgumentException("Cannot write object " + obj.toString() + " to " + p.toString() +
                            ". Conflicts with the structure of the datastore.");
                }
                _workers.put(targetDir, Pail.super.openWrite(p.toString(), _overwrite));
            }
            RecordOutputStream os = _workers.get(targetDir);
            os.writeRaw(structure.serialize(obj));
        }

        public void writeObjects(T... objs) throws IOException {
            for(T obj: objs) {
                writeObject(obj);
            }
        }

        public void close() throws IOException {
            for(RecordOutputStream os: _workers.values()) {
                os.close();
            }
        }

        protected List<String> makeRelative(List<String> attrs) {
            return Utils.stripRoot(getAttrs(), attrs);
        }

        public void writeRaw(byte[] record) throws IOException {
            writeRaw(record, 0, record.length);
        }

        public void writeRaw(byte[] record, int start, int length) throws IOException {
            if(!_workers.containsKey(_userfilename)) {
                checkValidStructure(_userfilename);
                _workers.put(_userfilename, Pail.super.openWrite(_userfilename, _overwrite));
            }
            _workers.get(_userfilename).writeRaw(record, start, length);
        }
    }


    public class TypedRecordInputStream implements RecordInputStream {
        private RecordInputStream is;

        public TypedRecordInputStream(String userFileName) throws IOException {
            is = Pail.super.openRead(userFileName);
        }

        public T readObject() throws IOException {
            byte[] record = readRawRecord();
            if(record==null) return null;
            else return _structure.deserialize(record);
        }

        public void close() throws IOException {
            is.close();
        }

        public byte[] readRawRecord() throws IOException {
            return is.readRawRecord();
        }
    }

    public static Pail create(String path, PailSpec spec) throws IOException {
        return create(Utils.getFS(path), path, spec);
    }

    public static Pail create(FileSystem fs, String path, PailSpec spec) throws IOException {
        return create(fs, path, spec, true);
    }

    public static Pail create(String path) throws IOException {
        return create(Utils.getFS(path), path);
    }

    public static Pail create(FileSystem fs, String path) throws IOException {
        return create(fs, path, (PailSpec) null);
    }

    public static Pail create(String path, PailStructure structure) throws IOException {
        return create(Utils.getFS(path), path, structure);
    }

    public static Pail create(FileSystem fs, String path, PailStructure structure) throws IOException {
        return create(fs, path, new PailSpec(structure));
    }

    public static Pail create(String path, PailStructure structure, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, structure, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, PailStructure structure, boolean failOnExists) throws IOException {
        return create(fs, path, new PailSpec(structure), failOnExists);
    }

    public static Pail create(String path, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, boolean failOnExists) throws IOException {
        return create(fs, path, (PailSpec) null, failOnExists);
    }

    public static Pail create(String path, PailSpec spec, boolean failOnExists) throws IOException {
        return create(Utils.getFS(path), path, spec, failOnExists);
    }

    public static Pail create(FileSystem fs, String path, PailSpec spec, boolean failOnExists) throws IOException {
        Path pathp = new Path(path);
        PailFormatFactory.create(spec);
        PailSpec existing = getSpec(fs, pathp);
        if(failOnExists) {
            if(existing!=null) {
                throw new IllegalArgumentException("Pail already exists at path " + path + " with spec " + existing.toString());
            }
            if(fs.exists(pathp))
                throw new IllegalArgumentException("Path " + path + " already exists");
        }
        if(spec!=null && existing!=null) {
            if(spec.getName()!=null) {
                if(!spec.equals(existing))
                    throw new IllegalArgumentException("Specs do not match " + spec.toString() + ", " + existing.toString());
            } else if(spec.getStructure()!=null) {
                if(existing.getStructure()==null || !spec.getStructure().getClass().equals(existing.getStructure().getClass())) {
                    throw new IllegalArgumentException("Specs do not match " + spec.toString() + ", " + existing.toString());
                }
            }
        }
        fs.mkdirs(pathp);
        if(existing==null) {
            if(spec==null) spec = PailFormatFactory.getDefaultCopy();
            if(spec.getName()==null) spec = PailFormatFactory.getDefaultCopy().setStructure(spec.getStructure());
            spec.writeToFileSystem(fs, new Path(pathp, META));
        }


        return new Pail(fs, path);
    }

    private static PailSpec getSpec(FileSystem fs, Path path) throws IOException {
        return (PailSpec) getSpecAndRoot(fs, path)[1];
    }

    private static String getRoot(FileSystem fs, Path path) throws IOException {
        return (String) getSpecAndRoot(fs, path)[0];
    }

    private static Object[] getSpecAndRoot(FileSystem fs, Path path) throws IOException {
        Path curr = path;
        Object[] ret = null;
        while( curr != null ) { //  changed as per oscar
        //while(true) {
            Path meta = new Path(curr, META);
            if(fs.exists(meta)) {
                if(ret!=null) throw new RuntimeException("At least two meta files up directory tree");
                PailSpec spec = PailSpec.readFromFileSystem(fs, meta);
                ret = new Object[] {curr.toString(), spec};
            }
            if(curr.depth()==0) break;
            curr = curr.getParent();
        }
        if(ret==null) ret = new Object[] {null, null};
        return ret;
    }


{'fraction_non_alphanumeric': 0.10228595942881658, 'fraction_numerical': 0.000858000858000858, 'mean_word_length': 2.53127028781649, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 27, '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': '531528', 'n_tokens_mistral': 4266, 'n_tokens_neox': 4165, 'n_words': 1111}