diff --git a/src/Microsoft.Data.Analysis/DataFrame.IO.cs b/src/Microsoft.Data.Analysis/DataFrame.IO.cs index 084b66379da..9a6a9154341 100644 --- a/src/Microsoft.Data.Analysis/DataFrame.IO.cs +++ b/src/Microsoft.Data.Analysis/DataFrame.IO.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; +using System.Data; using System.IO; using System.Text; @@ -306,5 +307,74 @@ public static DataFrame LoadCsv(Stream csvStream, return ret; } } + + /// + /// Reads an implementation of IDataReader into a DataFrame. + /// + /// DataReader to be read in + /// column names (can be empty) + /// column types (can be empty) + /// number of rows to read not including the header(if present) + /// add one column with the row index + /// + public static DataFrame FromDataReader(IDataReader reader, string[] columnNames = null, Type[] dataTypes = null, long numberOfRowsToRead = -1, bool addIndexColumn = false) + { + DataTable schemaTable = reader.GetSchemaTable(); + int numberOfColumns = schemaTable.Rows.Count; + + if (columnNames == null) + { + columnNames = new string[numberOfColumns]; + for (int i = 0; i < numberOfColumns; ++i) + { + string columnName = schemaTable.Rows[i]["ColumnName"].ToString(); + columnNames[i] = string.IsNullOrWhiteSpace(columnName) ? $"Column{i}" : columnName; + } + } + + var columns = new List(numberOfColumns); + if (dataTypes == null) + { + for (int i = 0; i < numberOfColumns; ++i) + { + var kind = (Type)schemaTable.Rows[i]["DataType"]; + columns.Add(CreateColumn(kind, columnNames, i)); + } + } + else + { + for (int i = 0; i < numberOfColumns; ++i) + { + columns.Add(CreateColumn(dataTypes[i], columnNames, i)); + } + } + + long rowline = 0; + var ret = new DataFrame(columns); + while (reader.Read() && (numberOfRowsToRead == -1 || rowline < numberOfRowsToRead)) + { + ret.Append(GetRecordValues(reader), inPlace: true); + ++rowline; + } + + if (addIndexColumn) + { + PrimitiveDataFrameColumn indexColumn = new PrimitiveDataFrameColumn("IndexColumn", columns[0].Length); + for (int i = 0; i < columns[0].Length; i++) + { + indexColumn[i] = i; + } + columns.Insert(0, indexColumn); + } + return ret; + + IEnumerable GetRecordValues(IDataRecord record) + { + for (int i = 0; i < record.FieldCount; i++) + { + yield return record[i] == DBNull.Value ? null : record[i]; + } + } + } } }