Skip to main content

rs_dirents2arrow_ipc_stream/
lib.rs

1use std::fs::DirEntry;
2use std::io;
3use std::os::unix::fs::PermissionsExt;
4use std::sync::Arc;
5
6use std::path::Path;
7
8use arrow_array::ArrayRef;
9use arrow_array::RecordBatch;
10use arrow_array::builder::BooleanBuilder;
11use arrow_array::builder::StringBuilder;
12use arrow_array::builder::TimestampMicrosecondBuilder;
13use arrow_array::builder::UInt32Builder;
14use arrow_array::builder::UInt64Builder;
15use arrow_ipc::writer::StreamWriter;
16use arrow_schema::DataType;
17use arrow_schema::Field;
18use arrow_schema::Schema;
19use arrow_schema::TimeUnit;
20use io::BufWriter;
21use io::Write;
22
23pub struct IpcStreamWriter<W: Write>(pub StreamWriter<BufWriter<W>>);
24
25impl<W> IpcStreamWriter<W>
26where
27    W: Write,
28{
29    pub fn write(&mut self, r: &RecordBatch) -> Result<(), io::Error> {
30        self.0.write(r).map_err(io::Error::other)
31    }
32
33    pub fn flush(&mut self) -> Result<(), io::Error> {
34        self.0.flush().map_err(io::Error::other)
35    }
36
37    pub fn finish(&mut self) -> Result<(), io::Error> {
38        self.0.finish().map_err(io::Error::other)
39    }
40}
41
42pub fn dirent_schema_full() -> Schema {
43    Schema::new(vec![
44        Field::new("basename", DataType::Utf8, false),
45        Field::new("filetype", DataType::Utf8, false),
46        Field::new("path", DataType::Utf8, false),
47        Field::new("len", DataType::UInt64, false),
48        Field::new("readonly", DataType::Boolean, false),
49        Field::new("permissions", DataType::UInt32, false),
50        Field::new(
51            "created",
52            DataType::Timestamp(TimeUnit::Microsecond, None),
53            true,
54        ),
55        Field::new(
56            "accessed",
57            DataType::Timestamp(TimeUnit::Microsecond, None),
58            true,
59        ),
60        Field::new(
61            "modified",
62            DataType::Timestamp(TimeUnit::Microsecond, None),
63            true,
64        ),
65    ])
66}
67
68pub fn dirents2batch_full<I>(dirents: I) -> Result<RecordBatch, io::Error>
69where
70    I: Iterator<Item = Result<DirEntry, io::Error>>,
71{
72    let schema = dirent_schema_full();
73    let mut basenames = StringBuilder::new();
74    let mut filetypes = StringBuilder::new();
75    let mut paths = StringBuilder::new();
76    let mut lens = UInt64Builder::new();
77    let mut readonlies = BooleanBuilder::new();
78    let mut permissions = UInt32Builder::new();
79    let mut createds = TimestampMicrosecondBuilder::new();
80    let mut accesseds = TimestampMicrosecondBuilder::new();
81    let mut modifieds = TimestampMicrosecondBuilder::new();
82
83    for dirent in dirents {
84        let dirent = dirent?;
85        let path = dirent.path();
86        let metadata = dirent.metadata()?;
87        let created = metadata.created();
88        let accessed = metadata.accessed();
89        let modified = metadata.modified();
90        basenames.append_value(dirent.file_name().to_string_lossy());
91        let file_type = metadata.file_type();
92        filetypes.append_value(if file_type.is_dir() {
93            "dir"
94        } else if file_type.is_file() {
95            "file"
96        } else if file_type.is_symlink() {
97            "symlink"
98        } else {
99            "unknown"
100        });
101        paths.append_value(path.to_string_lossy());
102        lens.append_value(metadata.len());
103        let perms = metadata.permissions();
104        readonlies.append_value(perms.readonly());
105        permissions.append_value(perms.mode());
106        createds.append_option(
107            created
108                .ok()
109                .and_then(|x| x.duration_since(std::time::UNIX_EPOCH).ok())
110                .map(|d| d.as_micros() as i64),
111        );
112        accesseds.append_option(
113            accessed
114                .ok()
115                .and_then(|x| x.duration_since(std::time::UNIX_EPOCH).ok())
116                .map(|d| d.as_micros() as i64),
117        );
118        modifieds.append_option(
119            modified
120                .ok()
121                .and_then(|x| x.duration_since(std::time::UNIX_EPOCH).ok())
122                .map(|d| d.as_micros() as i64),
123        );
124    }
125
126    let basename_array: ArrayRef = Arc::new(basenames.finish());
127    let filetype_array: ArrayRef = Arc::new(filetypes.finish());
128    let path_array: ArrayRef = Arc::new(paths.finish());
129    let len_array: ArrayRef = Arc::new(lens.finish());
130    let readonly_array: ArrayRef = Arc::new(readonlies.finish());
131    let permission_array: ArrayRef = Arc::new(permissions.finish());
132    let created_array: ArrayRef = Arc::new(createds.finish());
133    let accessed_array: ArrayRef = Arc::new(accesseds.finish());
134    let modified_array: ArrayRef = Arc::new(modifieds.finish());
135
136    RecordBatch::try_new(
137        Arc::new(schema),
138        vec![
139            basename_array,
140            filetype_array,
141            path_array,
142            len_array,
143            readonly_array,
144            permission_array,
145            created_array,
146            accessed_array,
147            modified_array,
148        ],
149    )
150    .map_err(io::Error::other)
151}
152
153pub fn dirname2batch_full<P>(dirname: P) -> Result<RecordBatch, io::Error>
154where
155    P: AsRef<Path>,
156{
157    let dirents = std::fs::read_dir(dirname)?;
158    dirents2batch_full(dirents)
159}